1#![allow(clippy::missing_panics_doc)] use gloss_img::DynImage;
8use image::{GenericImageView, Pixel};
9extern crate nalgebra as na;
10use core::f32;
11use itertools::izip;
12#[allow(unused_imports)]
13use log::{error, info, warn};
14use na::DMatrix;
15use obj_exporter::{Geometry, ObjSet, Object, Primitive, Shape, TVertex, Vertex};
16use std::ops::AddAssign;
17
18use ply_rs::ply::{self, Addable, Ply};
19
20use burn::tensor::{backend::Backend, Float, Int, Tensor};
21use gloss_utils::tensor;
22
23#[derive(PartialEq)]
24pub enum PerVertexNormalsWeightingType {
25 Uniform,
27 Area,
29}
30
31#[derive(PartialEq)]
32pub enum SplatType {
33 Avg,
35 Sum,
37}
38
39#[derive(PartialEq)]
40pub enum IndirRemovalPolicy {
41 RemoveInvalidRows,
42 RemoveInvalidCols,
43}
44
45pub fn save_obj(verts: &DMatrix<f32>, faces: Option<&DMatrix<u32>>, uv: Option<&DMatrix<f32>>, normals: Option<&DMatrix<f32>>, path: &str) {
46 let verts_obj: Vec<Vertex> = verts
47 .row_iter()
48 .map(|row| Vertex {
49 x: f64::from(row[0]),
50 y: f64::from(row[1]),
51 z: f64::from(row[2]),
52 })
53 .collect();
54 let faces_obj: Vec<Shape> = if let Some(faces) = faces {
55 faces
56 .row_iter()
57 .map(|row| Shape {
58 primitive: Primitive::Triangle(
59 (row[0] as usize, uv.map(|_| row[0] as usize), normals.map(|_| row[0] as usize)),
63 (row[1] as usize, uv.map(|_| row[1] as usize), normals.map(|_| row[1] as usize)),
64 (row[2] as usize, uv.map(|_| row[2] as usize), normals.map(|_| row[2] as usize)),
65 ),
66 groups: vec![],
67 smoothing_groups: vec![],
68 })
69 .collect()
70 } else {
71 Vec::new()
72 };
73 let uv_obj: Vec<TVertex> = if let Some(uv) = uv {
74 uv.row_iter()
75 .map(|row| TVertex {
76 u: f64::from(row[0]),
77 v: f64::from(row[1]),
78 w: 0.0,
79 })
80 .collect()
81 } else {
82 Vec::<TVertex>::new()
83 };
84
85 let normals_obj: Vec<Vertex> = if let Some(normals) = normals {
86 normals
87 .row_iter()
88 .map(|row| Vertex {
89 x: f64::from(row[0]),
90 y: f64::from(row[1]),
91 z: f64::from(row[2]),
92 })
93 .collect()
94 } else {
95 Vec::<Vertex>::new()
96 };
97
98 let set = ObjSet {
99 material_library: None,
100 objects: vec![Object {
101 name: "exported_mesh".to_owned(),
102 vertices: verts_obj,
103 tex_vertices: uv_obj,
104 normals: normals_obj,
105 geometry: vec![Geometry {
106 material_name: None,
107 shapes: faces_obj,
108 }],
109 }],
110 };
111
112 obj_exporter::export_to_file(&set, path).unwrap();
113}
114
115#[allow(clippy::too_many_lines)]
116pub fn save_ply(
117 verts: &DMatrix<f32>,
118 faces: Option<&DMatrix<u32>>,
119 uvs: Option<&DMatrix<f32>>,
120 normals: Option<&DMatrix<f32>>,
121 colors: Option<&DMatrix<f32>>,
122 path: &str,
123) {
124 let ply_float_prop = ply::PropertyType::Scalar(ply::ScalarType::Float);
125 let ply_uchar_prop = ply::PropertyType::Scalar(ply::ScalarType::UChar);
126
127 #[allow(clippy::approx_constant)]
128 let mut ply = {
129 let mut ply = Ply::<ply::DefaultElement>::new();
130 ply.header.encoding = ply::Encoding::Ascii; ply.header.comments.push("Gloss Ply file".to_string());
133
134 let mut point_element = ply::ElementDef::new("vertex".to_string());
138 let p = ply::PropertyDef::new("x".to_string(), ply_float_prop.clone());
139 point_element.properties.add(p);
140 let p = ply::PropertyDef::new("y".to_string(), ply_float_prop.clone());
141 point_element.properties.add(p);
142 let p = ply::PropertyDef::new("z".to_string(), ply_float_prop.clone());
143 point_element.properties.add(p);
144
145 if normals.is_some() {
146 let p = ply::PropertyDef::new("nx".to_string(), ply_float_prop.clone());
147 point_element.properties.add(p);
148 let p = ply::PropertyDef::new("ny".to_string(), ply_float_prop.clone());
149 point_element.properties.add(p);
150 let p = ply::PropertyDef::new("nz".to_string(), ply_float_prop.clone());
151 point_element.properties.add(p);
152 }
153
154 if uvs.is_some() {
156 let p = ply::PropertyDef::new("s".to_string(), ply_float_prop.clone());
157 point_element.properties.add(p);
158 let p = ply::PropertyDef::new("t".to_string(), ply_float_prop.clone());
159 point_element.properties.add(p);
160 }
161
162 if colors.is_some() {
164 let p = ply::PropertyDef::new("red".to_string(), ply_uchar_prop.clone());
165 point_element.properties.add(p);
166 let p = ply::PropertyDef::new("green".to_string(), ply_uchar_prop.clone());
167 point_element.properties.add(p);
168 let p = ply::PropertyDef::new("blue".to_string(), ply_uchar_prop.clone());
169 point_element.properties.add(p);
170 }
171
172 ply.header.elements.add(point_element);
173
174 let mut face_element = ply::ElementDef::new("face".to_string());
176 let f = ply::PropertyDef::new(
178 "vertex_indices".to_string(),
179 ply::PropertyType::List(ply::ScalarType::UChar, ply::ScalarType::UInt), );
181 face_element.properties.add(f);
182 ply.header.elements.add(face_element);
183
184 let mut points_list = Vec::new();
186 for (idx, vert) in verts.row_iter().enumerate() {
187 let mut point_elem = ply::DefaultElement::new();
188 point_elem.insert("x".to_string(), ply::Property::Float(vert[0]));
189 point_elem.insert("y".to_string(), ply::Property::Float(vert[1]));
190 point_elem.insert("z".to_string(), ply::Property::Float(vert[2]));
191
192 if let Some(normals) = normals {
193 let normal = normals.row(idx);
194 point_elem.insert("nx".to_string(), ply::Property::Float(normal[0]));
195 point_elem.insert("ny".to_string(), ply::Property::Float(normal[1]));
196 point_elem.insert("nz".to_string(), ply::Property::Float(normal[2]));
197 }
198
199 if let Some(uvs) = uvs {
200 let uv = uvs.row(idx);
201 point_elem.insert("s".to_string(), ply::Property::Float(uv[0]));
202 point_elem.insert("t".to_string(), ply::Property::Float(uv[1]));
203 }
204
205 #[allow(clippy::cast_sign_loss)]
206 if let Some(colors) = colors {
207 let color = colors.row(idx);
208 #[allow(clippy::cast_possible_truncation)]
209 point_elem.insert("red".to_string(), ply::Property::UChar((color[0] * 255.0) as u8));
210 #[allow(clippy::cast_possible_truncation)]
211 point_elem.insert("green".to_string(), ply::Property::UChar((color[1] * 255.0) as u8));
212 #[allow(clippy::cast_possible_truncation)]
213 point_elem.insert("blue".to_string(), ply::Property::UChar((color[2] * 255.0) as u8));
214 }
215
216 points_list.push(point_elem);
217 }
218 ply.payload.insert("vertex".to_string(), points_list);
219
220 if let Some(faces) = faces {
222 let mut faces_list = Vec::new();
223 for face in faces.row_iter() {
224 let mut face_elem = ply::DefaultElement::new();
225 face_elem.insert("vertex_indices".to_string(), ply::Property::ListUInt(face.iter().copied().collect()));
226 faces_list.push(face_elem);
227 }
228 ply.payload.insert("face".to_string(), faces_list);
229 }
230
231 ply.make_consistent().unwrap();
235 ply
236 };
237
238 let mut file = std::fs::File::create(path).unwrap();
239 let w = ply_rs::writer::Writer::new();
240 let written = w.write_ply(&mut file, &mut ply).unwrap();
241 println!("{written} bytes written");
242}
243
244pub fn compute_per_vertex_normals(
246 verts: &na::DMatrix<f32>,
247 faces: &na::DMatrix<u32>,
248 weighting_type: &PerVertexNormalsWeightingType,
249) -> na::DMatrix<f32> {
250 match weighting_type {
251 PerVertexNormalsWeightingType::Uniform => {
252 let faces_row_major = faces.transpose(); let verts_row_major = verts.transpose(); let mut per_vertex_normal_row_major = DMatrix::<f32>::zeros(3, verts.nrows());
255 for face in faces_row_major.column_iter() {
256 let v0 = verts_row_major.fixed_view::<3, 1>(0, face[0] as usize);
261 let v1 = verts_row_major.fixed_view::<3, 1>(0, face[1] as usize);
262 let v2 = verts_row_major.fixed_view::<3, 1>(0, face[2] as usize);
263 let d1 = v1 - v0;
264 let d2 = v2 - v0;
265 let face_normal = d1.cross(&d2).normalize();
266 per_vertex_normal_row_major.column_mut(face[0] as usize).add_assign(&face_normal);
268 per_vertex_normal_row_major.column_mut(face[1] as usize).add_assign(&face_normal);
269 per_vertex_normal_row_major.column_mut(face[2] as usize).add_assign(&face_normal);
270 }
271 for mut row in per_vertex_normal_row_major.column_iter_mut() {
273 row /= row.norm();
274 }
275 per_vertex_normal_row_major.transpose()
276 }
277 PerVertexNormalsWeightingType::Area => {
278 let face_normals = compute_per_face_normals(verts, faces);
280 let w = compute_double_face_areas(verts, faces);
282 let mut per_vertex_normal = DMatrix::<f32>::zeros(verts.nrows(), 3);
284 for ((face, face_normal), &face_weight) in faces.row_iter().zip(face_normals.row_iter()).zip(w.iter()) {
285 for &v_idx in face.iter() {
286 per_vertex_normal.row_mut(v_idx as usize).add_assign(face_weight * face_normal);
287 }
288 }
289 for mut row in per_vertex_normal.row_iter_mut() {
291 row /= row.norm();
292 }
293 per_vertex_normal
294 }
295 }
296}
297
298pub fn compute_per_vertex_normals_burn<B: Backend>(
300 verts: &Tensor<B, 2, Float>, faces: &Tensor<B, 2, Int>, weighting_type: &PerVertexNormalsWeightingType,
303) -> Tensor<B, 2, Float> {
304 let num_verts = verts.shape().dims[0];
305 let num_faces = faces.shape().dims[0];
306 let mut per_vertex_normals = Tensor::<B, 2, Float>::zeros([num_verts, 3], &verts.device());
307 match weighting_type {
309 PerVertexNormalsWeightingType::Uniform => {
310 let all_idxs: Tensor<B, 1, Int> = faces.clone().flatten(0, 1);
311 let all_tris_as_verts = verts.clone().select(0, all_idxs).reshape([num_faces, 3, 3]);
312 let all_tris_as_verts_chunks = all_tris_as_verts.chunk(3, 1); let v0 = all_tris_as_verts_chunks[0].clone().squeeze(1); let v1 = all_tris_as_verts_chunks[1].clone().squeeze(1); let v2 = all_tris_as_verts_chunks[2].clone().squeeze(1); let d1 = v1.sub(v0.clone()); let d2 = v2.sub(v0.clone()); let cross_product = tensor::cross_product(&d1, &d2); let face_normals = tensor::normalize_tensor(cross_product);
329
330 let face_indices_expanded: Tensor<B, 1, Int> = faces.clone().flatten(0, 1); let mut face_normals_repeated: Tensor<B, 3> = face_normals.unsqueeze_dim(1);
333
334 face_normals_repeated = face_normals_repeated.repeat(&[1, 3, 1]);
335
336 let face_normals_to_scatter = face_normals_repeated.reshape([num_faces * 3, 3]); per_vertex_normals = per_vertex_normals.select_assign(0, face_indices_expanded, face_normals_to_scatter);
340 per_vertex_normals = tensor::normalize_tensor(per_vertex_normals);
342 per_vertex_normals
343 }
344 PerVertexNormalsWeightingType::Area => {
345 let all_idxs: Tensor<B, 1, Int> = faces.clone().flatten(0, 1);
347 let all_tris_as_verts = verts.clone().select(0, all_idxs).reshape([num_faces, 3, 3]);
348 let all_tris_as_verts_chunks = all_tris_as_verts.chunk(3, 1); let v0 = all_tris_as_verts_chunks[0].clone().squeeze(1); let v1 = all_tris_as_verts_chunks[1].clone().squeeze(1); let v2 = all_tris_as_verts_chunks[2].clone().squeeze(1); let d1 = v1.sub(v0.clone()); let d2 = v2.sub(v0.clone()); let face_normals = tensor::cross_product(&d1, &d2); let face_areas = face_normals.clone().powf_scalar(2.0).sum_dim(1).sqrt(); let weighted_face_normals = face_normals.div(face_areas); let face_indices_expanded: Tensor<B, 1, Int> = faces.clone().flatten(0, 1); let mut weighted_face_normals_repeated: Tensor<B, 3> = weighted_face_normals.unsqueeze_dim(1);
366 weighted_face_normals_repeated = weighted_face_normals_repeated.repeat(&[1, 3, 1]);
367
368 let weighted_normals_to_scatter = weighted_face_normals_repeated.reshape([num_faces * 3, 3]);
369
370 per_vertex_normals = per_vertex_normals.select_assign(0, face_indices_expanded, weighted_normals_to_scatter);
371
372 per_vertex_normals = tensor::normalize_tensor(per_vertex_normals);
373 per_vertex_normals
374 }
375 }
376}
377
378pub fn compute_per_face_normals(verts: &na::DMatrix<f32>, faces: &na::DMatrix<u32>) -> na::DMatrix<f32> {
380 let verts_row_major = verts.transpose(); let mut face_normals = unsafe { DMatrix::<core::mem::MaybeUninit<f32>>::uninit(na::Dyn(faces.nrows()), na::Dyn(3)).assume_init() };
383
384 for (face, mut face_normal) in faces.row_iter().zip(face_normals.row_iter_mut()) {
385 let v0 = verts_row_major.fixed_view::<3, 1>(0, face[0] as usize);
389 let v1 = verts_row_major.fixed_view::<3, 1>(0, face[1] as usize);
390 let v2 = verts_row_major.fixed_view::<3, 1>(0, face[2] as usize);
391 let d1 = v1 - v0;
392 let d2 = v2 - v0;
393 let normal = d1.cross(&d2).normalize();
394 face_normal.copy_from(&normal.transpose());
397 }
398
399 face_normals
400}
401
402pub fn compute_double_face_areas(verts: &na::DMatrix<f32>, faces: &na::DMatrix<u32>) -> na::DMatrix<f32> {
404 let proj_doublearea =
406 |x, y, f| -> f32 {
408 let fx = faces[(f,0)] as usize;
409 let fy = faces[(f,1)] as usize;
410 let fz = faces[(f,2)] as usize;
411 let rx = verts[(fx,x)]-verts[(fz,x)];
412 let sx = verts[(fy,x)]-verts[(fz,x)];
413 let ry = verts[(fx,y)]-verts[(fz,y)];
414 let sy = verts[(fy,y)]-verts[(fz,y)];
415 rx*sy - ry*sx
416 };
417
418 let mut double_face_areas = DMatrix::<f32>::zeros(faces.nrows(), 1);
419
420 for f in 0..faces.nrows() {
421 for d in 0..3 {
422 let double_area = proj_doublearea(d, (d + 1) % 3, f);
423 double_face_areas[(f, 0)] += double_area * double_area;
424 }
425 }
426 double_face_areas = double_face_areas.map(f32::sqrt);
428
429 double_face_areas
430}
431
432#[allow(clippy::similar_names)]
433pub fn compute_tangents(verts: &na::DMatrix<f32>, faces: &na::DMatrix<u32>, normals: &na::DMatrix<f32>, uvs: &na::DMatrix<f32>) -> na::DMatrix<f32> {
435 let faces_row_major = faces.transpose(); let verts_row_major = verts.transpose(); let normals_row_major = normals.transpose(); let uvs_row_major = uvs.transpose(); let mut tangents_rm = DMatrix::<f32>::zeros(3, verts.nrows());
451 let mut handness_rm = DMatrix::<f32>::zeros(1, verts.nrows());
452 let mut bitangents_rm = DMatrix::<f32>::zeros(3, verts.nrows());
453
454 let mut v_uv_rm = DMatrix::<f32>::zeros(3 + 2, verts.nrows());
456 v_uv_rm.view_mut((0, 0), (3, verts.nrows())).copy_from(&verts_row_major);
457 v_uv_rm.view_mut((3, 0), (2, verts.nrows())).copy_from(&uvs_row_major);
458
459 for face in faces_row_major.column_iter() {
460 let id0 = face[0] as usize;
465 let id1 = face[1] as usize;
466 let id2 = face[2] as usize;
467
468 let v_uv0 = v_uv_rm.fixed_view::<5, 1>(0, id0);
481 let v_uv1 = v_uv_rm.fixed_view::<5, 1>(0, id1);
482 let v_uv2 = v_uv_rm.fixed_view::<5, 1>(0, id2);
483
484 let v0 = v_uv0.fixed_rows::<3>(0);
486 let v1 = v_uv1.fixed_rows::<3>(0);
487 let v2 = v_uv2.fixed_rows::<3>(0);
488 let uv0 = v_uv0.fixed_rows::<2>(3);
490 let uv1 = v_uv1.fixed_rows::<2>(3);
491 let uv2 = v_uv2.fixed_rows::<2>(3);
492
493 let delta_pos1 = v1 - v0;
495 let delta_pos2 = v2 - v0;
496
497 let delta_uv1 = uv1 - uv0;
499 let delta_uv2 = uv2 - uv0;
500
501 let r = 1.0 / (delta_uv1[0] * delta_uv2[1] - delta_uv1[1] * delta_uv2[0]);
502 let tangent = (delta_pos1 * delta_uv2[1] - delta_pos2 * delta_uv1[1]) * r;
503 let bitangent = (delta_pos2 * delta_uv1[0] - delta_pos1 * delta_uv2[0]) * r;
504
505 if r.is_finite() {
510 tangents_rm.column_mut(id0).add_assign(&tangent);
511 tangents_rm.column_mut(id1).add_assign(&tangent);
512 tangents_rm.column_mut(id2).add_assign(&tangent);
513 bitangents_rm.column_mut(id0).add_assign(&bitangent);
514 bitangents_rm.column_mut(id1).add_assign(&bitangent);
515 bitangents_rm.column_mut(id2).add_assign(&bitangent);
516 }
517 }
518
519 for (idx, ((mut tangent, normal), bitangent)) in tangents_rm
523 .column_iter_mut()
524 .zip(normals_row_major.column_iter())
525 .zip(bitangents_rm.column_iter())
526 .enumerate()
527 {
528 let dot = tangent.dot(&normal);
530 let new_tangent = (&tangent - normal * dot).normalize();
532 let cross_vec = tangent.cross(&bitangent);
534 let dot: f32 = cross_vec.dot(&normal);
535 handness_rm.column_mut(idx).fill(dot.signum()); tangent.copy_from_slice(new_tangent.data.as_slice());
538 }
542
543 let mut tangents_and_handness_rm = DMatrix::<f32>::zeros(4, verts.nrows());
558 tangents_and_handness_rm.view_mut((0, 0), (3, verts.nrows())).copy_from(&tangents_rm);
559 tangents_and_handness_rm.view_mut((3, 0), (1, verts.nrows())).copy_from(&handness_rm);
560
561 tangents_and_handness_rm.transpose()
562}
563
564#[allow(clippy::similar_names)]
565#[allow(clippy::too_many_lines)]
566pub fn compute_tangents_burn<B: Backend>(
568 verts: &Tensor<B, 2, Float>, faces: &Tensor<B, 2, Int>, normals: &Tensor<B, 2, Float>, uvs: &Tensor<B, 2, Float>, ) -> Tensor<B, 2, Float> {
573 let num_faces = faces.shape().dims[0];
578 let num_verts = verts.shape().dims[0];
579 let all_idxs: Tensor<B, 1, Int> = faces.clone().flatten(0, 1);
583 let all_tris_as_verts = verts.clone().select(0, all_idxs.clone()).reshape([num_faces, 3, 3]);
587 let all_tris_as_uvs = uvs.clone().select(0, all_idxs.clone()).reshape([num_faces, 3, 2]);
588 let v0: Tensor<B, 2> = all_tris_as_verts.clone().slice([0..num_faces, 0..1, 0..3]).squeeze(1);
591 let v1: Tensor<B, 2> = all_tris_as_verts.clone().slice([0..num_faces, 1..2, 0..3]).squeeze(1);
592 let v2: Tensor<B, 2> = all_tris_as_verts.clone().slice([0..num_faces, 2..3, 0..3]).squeeze(1);
593 let uv0 = all_tris_as_uvs.clone().slice([0..num_faces, 0..1, 0..2]).squeeze(1);
596 let uv1 = all_tris_as_uvs.clone().slice([0..num_faces, 1..2, 0..2]).squeeze(1);
597 let uv2 = all_tris_as_uvs.clone().slice([0..num_faces, 2..3, 0..2]).squeeze(1);
598
599 let delta_pos1 = v1.clone().sub(v0.clone());
600 let delta_pos2 = v2.clone().sub(v0.clone());
601
602 let delta_uv1 = uv1.clone().sub(uv0.clone());
603 let delta_uv2 = uv2.clone().sub(uv0.clone());
604 let delta_uv1_0 = delta_uv1.clone().select(1, Tensor::<B, 1, Int>::from_ints([0], &verts.device())); let delta_uv1_1 = delta_uv1.clone().select(1, Tensor::<B, 1, Int>::from_ints([1], &verts.device())); let delta_uv2_0 = delta_uv2.clone().select(1, Tensor::<B, 1, Int>::from_ints([0], &verts.device())); let delta_uv2_1 = delta_uv2.clone().select(1, Tensor::<B, 1, Int>::from_ints([1], &verts.device())); let denominator = delta_uv1_0
614 .clone()
615 .mul(delta_uv2_1.clone())
616 .sub(delta_uv1_1.clone().mul(delta_uv2_0.clone()));
617
618 let r_mask = denominator.clone().abs().greater_elem(f32::EPSILON).float(); let r = denominator.recip().mul(r_mask.clone()); let tangent = delta_pos1
623 .clone()
624 .mul(delta_uv2_1.clone().repeat(&[1, 3]))
625 .sub(delta_pos2.clone().mul(delta_uv1_1.clone().repeat(&[1, 3])))
626 .mul(r.clone());
627 let bitangent = delta_pos2
630 .clone()
631 .mul(delta_uv1_0.clone())
632 .sub(delta_pos1.clone().mul(delta_uv2_0.clone()))
633 .mul(r.clone()); let mut tangents_rm = Tensor::<B, 2, Float>::zeros([verts.dims()[0], 3], &verts.device());
637 let mut handness_rm = Tensor::<B, 2, Float>::zeros([verts.dims()[0], 1], &verts.device());
638 let mut bitangents_rm = Tensor::<B, 2, Float>::zeros([verts.dims()[0], 3], &verts.device());
639 let face_indices_expanded: Tensor<B, 1, Int> = faces.clone().flatten(0, 1); let mut face_tangents_repeated: Tensor<B, 3> = tangent.unsqueeze_dim(1);
643 face_tangents_repeated = face_tangents_repeated.repeat(&[1, 3, 1]);
644
645 let face_tangents_to_scatter = face_tangents_repeated.reshape([num_faces * 3, 3]); tangents_rm = tangents_rm.select_assign(0, face_indices_expanded.clone(), face_tangents_to_scatter.clone());
649 let mut face_bitangents_repeated: Tensor<B, 3> = bitangent.unsqueeze_dim(1);
652 face_bitangents_repeated = face_bitangents_repeated.repeat(&[1, 3, 1]);
653
654 let face_bitangents_to_scatter = face_bitangents_repeated.reshape([num_faces * 3, 3]); bitangents_rm = bitangents_rm.select_assign(0, face_indices_expanded, face_bitangents_to_scatter.clone());
658 let dot_product = tangents_rm.clone().mul(normals.clone()).sum_dim(1); let new_tangents = tensor::normalize_tensor(tangents_rm.clone().sub(normals.clone().mul(dot_product))); let cross_vec = tensor::cross_product(&tangents_rm, &bitangents_rm);
664 let handedness_sign = cross_vec.clone().mul(normals.clone()).sum_dim(1).sign(); let handness_indices: Tensor<B, 1, Int> = Tensor::<B, 1, Int>::arange(
667 0..(i64::try_from(handness_rm.shape().dims[0]).expect("Dimension size exceeds i64 range")),
668 &verts.device(),
669 );
670 handness_rm = handness_rm.select_assign(0, handness_indices.clone(), handedness_sign);
671 tangents_rm = tangents_rm.slice_assign([0..num_verts, 0..3], new_tangents);
672 let tangent_x = tangents_rm.clone().slice([0..num_verts, 0..1]);
678 let tangent_y = tangents_rm.clone().slice([0..num_verts, 1..2]);
679 let tangent_z = tangents_rm.clone().slice([0..num_verts, 2..3]);
680
681 let handness = handness_rm.clone().slice([0..num_verts, 0..1]);
682 let tangents_and_handness_rm: Tensor<B, 2, Float> = Tensor::<B, 1>::stack(
685 vec![tangent_x.squeeze(1), tangent_y.squeeze(1), tangent_z.squeeze(1), handness.squeeze(1)],
686 1,
687 );
688 tangents_and_handness_rm
689}
690
691pub fn transform_verts(verts: &na::DMatrix<f32>, model_matrix: &na::SimilarityMatrix3<f32>) -> na::DMatrix<f32> {
692 let mut verts_transformed = verts.clone();
693 for (vert, mut vert_transformed) in verts.row_iter().zip(verts_transformed.row_iter_mut()) {
694 let v_modif = model_matrix * na::Point3::from(vert.fixed_columns::<3>(0).transpose());
697 vert_transformed.copy_from_slice(v_modif.coords.as_slice());
698 }
699 verts_transformed
700}
701
702pub fn transform_vectors(verts: &na::DMatrix<f32>, model_matrix: &na::SimilarityMatrix3<f32>) -> na::DMatrix<f32> {
703 let mut verts_transformed = verts.clone();
704 for (vert, mut vert_transformed) in verts.row_iter().zip(verts_transformed.row_iter_mut()) {
705 let v_modif = model_matrix * vert.fixed_columns::<3>(0).transpose();
707 vert_transformed.copy_from_slice(v_modif.data.as_slice());
708 }
709 verts_transformed
710}
711
712pub fn get_bounding_points(verts: &na::DMatrix<f32>, model_matrix: Option<na::SimilarityMatrix3<f32>>) -> (na::Point3<f32>, na::Point3<f32>) {
713 let mut min_point_global = na::Point3::<f32>::new(f32::MAX, f32::MAX, f32::MAX);
714 let mut max_point_global = na::Point3::<f32>::new(f32::MIN, f32::MIN, f32::MIN);
715
716 let min_coord_vec: Vec<f32> = verts.column_iter().map(|c| c.min()).collect();
718 let max_coord_vec: Vec<f32> = verts.column_iter().map(|c| c.max()).collect();
719 let min_point = na::Point3::<f32>::from_slice(&min_coord_vec);
720 let max_point = na::Point3::<f32>::from_slice(&max_coord_vec);
721
722 let (min_point_w, max_point_w) = if let Some(model_matrix) = model_matrix {
724 (model_matrix * min_point, model_matrix * max_point)
725 } else {
726 (min_point, max_point)
727 };
728
729 min_point_global = min_point_global.inf(&min_point_w);
731 max_point_global = max_point_global.sup(&max_point_w);
732
733 (min_point_global, max_point_global)
734}
735
736pub fn get_centroid(verts: &na::DMatrix<f32>, model_matrix: Option<na::SimilarityMatrix3<f32>>) -> na::Point3<f32> {
737 let (min_point_global, max_point_global) = get_bounding_points(verts, model_matrix);
738
739 min_point_global.lerp(&max_point_global, 0.5)
741}
742
743#[allow(clippy::cast_possible_truncation)]
744#[allow(clippy::cast_precision_loss)]
745#[allow(clippy::cast_sign_loss)]
746pub fn sample_img_with_uvs(img: &DynImage, uvs: &na::DMatrix<f32>, is_srgb: bool) -> DMatrix<f32> {
747 let mut sampled_pixels =
748 unsafe { DMatrix::<core::mem::MaybeUninit<f32>>::uninit(na::Dyn(uvs.nrows()), na::Dyn(img.channels() as usize)).assume_init() };
749
750 for (i, uv) in uvs.row_iter().enumerate() {
751 let x = (uv[0] * img.width() as f32) - 0.5;
752 let y = ((1.0 - uv[1]) * img.height() as f32) - 0.5;
753
754 let x = x.floor() as u32;
756 let y = y.floor() as u32;
757 let mut sample = img.get_pixel(x, y);
758
759 if is_srgb {
760 sample.apply(|v| (v / 255.0).powf(2.2) * 255.0);
761 }
762
763 sampled_pixels.row_mut(i).copy_from_slice(&sample.0[0..img.channels() as usize]);
764 }
765
766 sampled_pixels
771}
772
773pub fn filter_rows<T: na::Scalar>(mat: &na::DMatrix<T>, mask: &[bool], keep: bool) -> (DMatrix<T>, Vec<i32>, Vec<i32>) {
778 assert_eq!(
779 mat.nrows(),
780 mask.len(),
781 "Mat and mask need to have the same nr of rows. Mat has nr rows{} while mask have length {}",
782 mat.nrows(),
783 mask.len()
784 );
785
786 let nr_filtered_rows: u32 = mask.iter().map(|v| u32::from(*v == keep)).sum();
788
789 let mut selected =
790 unsafe { DMatrix::<core::mem::MaybeUninit<T>>::uninit(na::Dyn(nr_filtered_rows as usize), na::Dyn(mat.ncols())).assume_init() };
791
792 let mut orig2filtered: Vec<i32> = vec![-1; mat.nrows()];
794 let mut filtered2orig: Vec<i32> = vec![-1; nr_filtered_rows as usize];
795
796 let mut idx_filtered = 0;
797 for (idx_orig, (val, mask_val)) in izip!(mat.row_iter(), mask.iter()).enumerate() {
798 if *mask_val == keep {
799 selected.row_mut(idx_filtered).copy_from(&val);
800
801 orig2filtered[idx_orig] = i32::try_from(idx_filtered).unwrap();
803 filtered2orig[idx_filtered] = i32::try_from(idx_orig).unwrap();
804
805 idx_filtered += 1;
806 }
807 }
808
809 (selected, orig2filtered, filtered2orig)
810}
811
812pub fn filter_cols<T: na::Scalar>(mat: &na::DMatrix<T>, mask: &[bool], keep: bool) -> (DMatrix<T>, Vec<i32>, Vec<i32>) {
817 assert_eq!(
818 mat.ncols(),
819 mask.len(),
820 "Mat and mask need to have the same nr of cols. Mat has nr cols{} while mask have length {}",
821 mat.ncols(),
822 mask.len()
823 );
824
825 let nr_filtered_cols: u32 = mask.iter().map(|v| u32::from(*v == keep)).sum();
827
828 let mut selected =
829 unsafe { DMatrix::<core::mem::MaybeUninit<T>>::uninit(na::Dyn(mat.nrows()), na::Dyn(nr_filtered_cols as usize)).assume_init() };
830
831 let mut orig2filtered: Vec<i32> = vec![-1; mat.ncols()];
833 let mut filtered2orig: Vec<i32> = vec![-1; nr_filtered_cols as usize];
834
835 let mut idx_filtered = 0;
836 for (idx_orig, (val, mask_val)) in izip!(mat.column_iter(), mask.iter()).enumerate() {
837 if *mask_val == keep {
838 selected.column_mut(idx_filtered).copy_from(&val);
839
840 orig2filtered[idx_orig] = i32::try_from(idx_filtered).unwrap();
842 filtered2orig[idx_filtered] = i32::try_from(idx_orig).unwrap();
843
844 idx_filtered += 1;
845 }
846 }
847
848 (selected, orig2filtered, filtered2orig)
849}
850
851pub fn splat_rows(mat: &na::DMatrix<f32>, indices_orig2splatted: &[u32], splat_type: &SplatType) -> DMatrix<f32> {
859 assert_eq!(
860 mat.nrows(),
861 indices_orig2splatted.len(),
862 "Mat and indices need to have the same nr of rows. Mat has nr rows{} while indices have length {}",
863 mat.nrows(),
864 indices_orig2splatted.len()
865 );
866
867 let max_idx_splatted = *indices_orig2splatted.iter().max().unwrap() as usize;
872
873 let mut splatted = na::DMatrix::zeros(max_idx_splatted + 1, mat.ncols());
874 let mut normalization: Vec<f32> = vec![0.0; max_idx_splatted + 1];
875
876 for (val, idx_destin) in izip!(mat.row_iter(), indices_orig2splatted.iter()) {
877 let mut splatted_row = splatted.row_mut(*idx_destin as usize);
878 splatted_row.add_assign(val);
879
880 #[allow(clippy::single_match)]
881 match splat_type {
882 SplatType::Avg => {
883 normalization[*idx_destin as usize] += 1.0;
884 }
885 SplatType::Sum => {}
886 }
887 }
888
889 match splat_type {
891 SplatType::Avg => {
892 for (mut splatted_row, normalization_row) in izip!(splatted.row_iter_mut(), normalization.iter()) {
893 for e in splatted_row.iter_mut() {
894 *e /= normalization_row;
895 }
896 }
897 }
898 SplatType::Sum => {}
899 }
900
901 splatted
902}
903
904pub fn apply_indirection(mat: &na::DMatrix<u32>, indices_orig2destin: &[i32], removal_policy: &IndirRemovalPolicy) -> (DMatrix<u32>, Vec<bool>) {
905 let reindexed: DMatrix<i32> = DMatrix::from_iterator(
908 mat.nrows(),
909 mat.ncols(),
910 mat.iter().map(|v| indices_orig2destin.get(*v as usize).map_or(-1, |v| *v)),
911 );
912
913 let (reindexed_only_valid, mask) = match removal_policy {
915 IndirRemovalPolicy::RemoveInvalidRows => {
916 let mut valid_rows = vec![true; mat.nrows()];
917 reindexed
918 .row_iter()
919 .enumerate()
920 .filter(|(_, r)| r.iter().any(|x| *x == -1))
921 .for_each(|(idx, _)| valid_rows[idx] = false);
922 let (filtered, _, _) = filter_rows(&reindexed, &valid_rows, true);
923 (filtered, valid_rows)
924 }
925 IndirRemovalPolicy::RemoveInvalidCols => {
926 let mut valid_cols = vec![true; mat.ncols()];
927 reindexed
928 .column_iter()
929 .enumerate()
930 .filter(|(_, c)| c.iter().any(|x| *x == -1))
931 .for_each(|(idx, _)| valid_cols[idx] = false);
932 let (filtered, _, _) = filter_cols(&reindexed, &valid_cols, true);
933 (filtered, valid_cols)
934 }
935 };
936
937 let reindexed_filtered = reindexed_only_valid.try_cast::<u32>().unwrap();
938
939 (reindexed_filtered, mask)
940}
941
942pub fn compute_dummy_uvs<B: Backend>(nr_verts: usize, device: &B::Device) -> Tensor<B, 2, Float> {
952 Tensor::<B, 2, Float>::zeros([nr_verts, 2], device)
953}
954pub fn compute_dummy_colors<B: Backend>(nr_verts: usize, device: &B::Device) -> Tensor<B, 2, Float> {
955 Tensor::<B, 2, Float>::zeros([nr_verts, 3], device)
956}
957pub fn compute_dummy_tangents<B: Backend>(nr_verts: usize, device: &B::Device) -> Tensor<B, 2, Float> {
958 Tensor::<B, 2, Float>::zeros([nr_verts, 4], device)
959}
960
961pub fn create_frustum_verts_and_edges(
962 extrinsics: &na::Matrix4<f32>,
963 fovy: f32,
964 aspect_ratio: f32,
965 near: f32, far: f32,
967) -> (DMatrix<f32>, DMatrix<u32>) {
968 let rot = extrinsics.fixed_view::<3, 3>(0, 0);
970 let trans = extrinsics.fixed_view::<3, 1>(0, 3);
971
972 let lookfrom = -rot.transpose() * trans;
974
975 let forward = -rot.column(2);
976 let right = rot.column(0);
977 let up = rot.column(1);
978
979 let near_height = 2.0 * (fovy / 2.0).tan() * near;
981 let near_width = near_height * aspect_ratio;
982
983 let far_height = 2.0 * (fovy / 2.0).tan() * far;
984 let far_width = far_height * aspect_ratio;
985
986 let near_center = lookfrom + forward * near;
988 let near_top_left = near_center + (up * (near_height / 2.0)) - (right * (near_width / 2.0));
989 let near_top_right = near_center + (up * (near_height / 2.0)) + (right * (near_width / 2.0));
990 let near_bottom_left = near_center - (up * (near_height / 2.0)) - (right * (near_width / 2.0));
991 let near_bottom_right = near_center - (up * (near_height / 2.0)) + (right * (near_width / 2.0));
992
993 let far_center = lookfrom + forward * far;
994 let far_top_left = far_center + (up * (far_height / 2.0)) - (right * (far_width / 2.0));
995 let far_top_right = far_center + (up * (far_height / 2.0)) + (right * (far_width / 2.0));
996 let far_bottom_left = far_center - (up * (far_height / 2.0)) - (right * (far_width / 2.0));
997 let far_bottom_right = far_center - (up * (far_height / 2.0)) + (right * (far_width / 2.0));
998
999 let verts = DMatrix::from_row_slice(
1001 8,
1002 3,
1003 &[
1004 near_top_left.x,
1005 near_top_left.y,
1006 near_top_left.z,
1007 near_top_right.x,
1008 near_top_right.y,
1009 near_top_right.z,
1010 near_bottom_left.x,
1011 near_bottom_left.y,
1012 near_bottom_left.z,
1013 near_bottom_right.x,
1014 near_bottom_right.y,
1015 near_bottom_right.z,
1016 far_top_left.x,
1017 far_top_left.y,
1018 far_top_left.z,
1019 far_top_right.x,
1020 far_top_right.y,
1021 far_top_right.z,
1022 far_bottom_left.x,
1023 far_bottom_left.y,
1024 far_bottom_left.z,
1025 far_bottom_right.x,
1026 far_bottom_right.y,
1027 far_bottom_right.z,
1028 ],
1029 );
1030
1031 let edges = DMatrix::from_row_slice(12, 2, &[0, 1, 1, 3, 3, 2, 2, 0, 4, 5, 5, 7, 7, 6, 6, 4, 0, 4, 1, 5, 2, 6, 3, 7]);
1032
1033 (verts, edges)
1034}