Implement mesh clamping for MeshWithData, implement keep_vertices

w1th0utnam3 · w1th0utnam3 · commit ad93f37d3d9d · 2023-09-13T16:42:08.000+02:00
diff --git a/splashsurf/src/reconstruction.rs b/splashsurf/src/reconstruction.rs
@@ -210,7 +210,7 @@ pub struct ReconstructSubcommandArgs {
         ignore_case = true,
         require_equals = true
     )]
-    pub keep_vertices: Switch,
+    pub keep_verts: Switch,
     /// Whether to compute surface normals at the mesh vertices and write them to the output file
     #[arg(
         help_heading = ARGS_INTERP,
@@ -283,6 +283,7 @@ pub struct ReconstructSubcommandArgs {
         require_equals = true
     )]
     pub output_raw_mesh: Switch,
+
     /// Whether to try to convert triangles to quads if they meet quality criteria
     #[arg(
         help_heading = ARGS_INTERP,
@@ -303,7 +304,7 @@ pub struct ReconstructSubcommandArgs {
     #[arg(help_heading = ARGS_INTERP, long, default_value = "135")]
     pub quad_max_interior_angle: f64,
 
-    /// Lower corner of the bounding-box for the surface mesh, mesh outside gets cut away (requires mesh-max to be specified)
+    /// Lower corner of the bounding-box for the surface mesh, triangles completely outside are removed (requires mesh-aabb-max to be specified)
     #[arg(
         help_heading = ARGS_POSTPROC,
         long,
@@ -313,7 +314,7 @@ pub struct ReconstructSubcommandArgs {
         requires = "mesh_aabb_max",
     )]
     pub mesh_aabb_min: Option<Vec<f64>>,
-    /// Upper corner of the bounding-box for the surface mesh, mesh outside gets cut away (requires mesh-min to be specified)
+    /// Upper corner of the bounding-box for the surface mesh, triangles completely outside are removed (requires mesh-aabb-min to be specified)
     #[arg(
         help_heading = ARGS_POSTPROC,
         long,
@@ -323,6 +324,16 @@ pub struct ReconstructSubcommandArgs {
         requires = "mesh_aabb_min",
     )]
     pub mesh_aabb_max: Option<Vec<f64>>,
+    /// Whether to clamp vertices outside of the specified mesh AABB to the AABB (only has an effect if mesh-aabb-min/max are specified)
+    #[arg(
+        help_heading = ARGS_POSTPROC,
+        long,
+        default_value = "off",
+        value_name = "off|on",
+        ignore_case = true,
+        require_equals = true
+    )]
+    pub mesh_aabb_clamp_verts: Switch,
 
     /// Optional filename for writing the point cloud representation of the intermediate density map to disk
     #[arg(help_heading = ARGS_DEBUG, long, value_parser = value_parser!(PathBuf))]
@@ -456,6 +467,7 @@ mod arguments {
         pub output_raw_normals: bool,
         pub output_raw_mesh: bool,
         pub mesh_aabb: Option<Aabb3d<f64>>,
+        pub mesh_aabb_clamp_vertices: bool,
     }
 
     /// All arguments that can be supplied to the surface reconstruction tool converted to useful types
@@ -590,7 +602,7 @@ mod arguments {
                 check_mesh: args.check_mesh.into_bool(),
                 mesh_cleanup: args.mesh_cleanup.into_bool(),
                 decimate_barnacles: args.decimate_barnacles.into_bool(),
-                keep_vertices: args.keep_vertices.into_bool(),
+                keep_vertices: args.keep_verts.into_bool(),
                 compute_normals: args.normals.into_bool(),
                 sph_normals: args.sph_normals.into_bool(),
                 normals_smoothing_iters: args.normals_smoothing_iters,
@@ -606,6 +618,7 @@ mod arguments {
                 output_raw_normals: args.output_raw_normals.into_bool(),
                 output_raw_mesh: args.output_raw_mesh.into_bool(),
                 mesh_aabb,
+                mesh_aabb_clamp_vertices: args.mesh_aabb_clamp_verts.into_bool(),
             };
 
             Ok(ReconstructionRunnerArgs {
@@ -1308,6 +1321,22 @@ pub(crate) fn reconstruction_pipeline_generic<I: Index, R: Real>(
         }
     }
 
+    // Remove and clamp cells outside of AABB
+    let mesh_with_data = if let Some(mesh_aabb) = &postprocessing.mesh_aabb {
+        profile!("clamp mesh to aabb");
+        info!("Post-processing: Clamping mesh to AABB...");
+
+        mesh_with_data.par_clamp_with_aabb(
+            &mesh_aabb
+                .try_convert()
+                .ok_or_else(|| anyhow!("Failed to convert mesh AABB"))?,
+            postprocessing.mesh_aabb_clamp_vertices,
+            postprocessing.keep_vertices,
+        )
+    } else {
+        mesh_with_data
+    };
+
     // Convert triangles to quads
     let (tri_mesh, tri_quad_mesh) = if postprocessing.generate_quads {
         info!("Post-processing: Convert triangles to quads...");
diff --git a/splashsurf_lib/src/mesh.rs b/splashsurf_lib/src/mesh.rs
@@ -310,20 +310,91 @@ where
     }
 
     /// Returns a new mesh containing only the specified cells and removes all unreferenced vertices
-    fn keep_cells(&self, cell_indices: &[usize]) -> Self {
+    fn keep_cells(&self, cell_indices: &[usize], keep_vertices: bool) -> Self {
+        if keep_vertices {
+            keep_cells_impl(self, cell_indices, &[])
+        } else {
+            let vertex_keep_table = vertex_keep_table(self, cell_indices);
+            keep_cells_impl(self, cell_indices, &vertex_keep_table)
+        }
+    }
+
+    /// Removes all cells from the mesh that are completely outside of the given AABB and clamps the remaining cells to the boundary
+    fn par_clamp_with_aabb(
+        &self,
+        aabb: &Aabb3d<R>,
+        clamp_vertices: bool,
+        keep_vertices: bool,
+    ) -> Self
+    where
+        Self::Cell: Sync,
+    {
+        // Find all triangles with at least one vertex inside of AABB
         let vertices = self.vertices();
-        let cells = self.cells();
-
-        // Each entry is true if this vertex should be kept, false otherwise
-        let vertex_keep_table = {
-            let mut table = vec![false; vertices.len()];
-            for cell in cell_indices.iter().copied().map(|c_i| &cells[c_i]) {
-                for &vertex_index in cell.vertices() {
-                    table[vertex_index] = true;
-                }
+        let cells_to_keep = self
+            .cells()
+            .par_iter()
+            .enumerate()
+            .filter(|(_, cell)| {
+                cell.vertices()
+                    .iter()
+                    .copied()
+                    .any(|v| aabb.contains_point(&vertices[v]))
+            })
+            .map(|(i, _)| i)
+            .collect::<Vec<_>>();
+        // Remove all other cells from mesh
+        let mut new_mesh = self.keep_cells(&cells_to_keep, keep_vertices);
+        // Clamp remaining vertices to AABB
+        if clamp_vertices {
+            new_mesh.vertices_mut().par_iter_mut().for_each(|v| {
+                let min = aabb.min();
+                let max = aabb.max();
+                v.x = v.x.clamp(min.x, max.x);
+                v.y = v.y.clamp(min.y, max.y);
+                v.z = v.z.clamp(min.z, max.z);
+            });
+        }
+
+        new_mesh
+    }
+}
+
+/// Returns the list of vertices that should remain in the given mesh after keeping only the given cells
+fn vertex_keep_table<R: Real, MeshT: Mesh3d<R>>(mesh: &MeshT, cell_indices: &[usize]) -> Vec<bool> {
+    let vertices = mesh.vertices();
+    let cells = mesh.cells();
+
+    // Each entry is true if this vertex should be kept, false otherwise
+    let vertex_keep_table = {
+        let mut table = vec![false; vertices.len()];
+        for cell in cell_indices.iter().copied().map(|c_i| &cells[c_i]) {
+            for &vertex_index in cell.vertices() {
+                table[vertex_index] = true;
             }
-            table
-        };
+        }
+        table
+    };
+
+    vertex_keep_table
+}
+
+/// Returns a new mesh keeping only the given cells and vertices in the mesh
+fn keep_cells_impl<R: Real, MeshT: Mesh3d<R>>(
+    mesh: &MeshT,
+    cell_indices: &[usize],
+    vertex_keep_table: &[bool],
+) -> MeshT {
+    let vertices = mesh.vertices();
+    let cells = mesh.cells();
+
+    if vertex_keep_table.is_empty() {
+        MeshT::from_vertices_and_connectivity(
+            mesh.vertices().to_vec(),
+            cell_indices.iter().map(|&i| &cells[i]).cloned().collect(),
+        )
+    } else {
+        assert_eq!(mesh.vertices().len(), vertex_keep_table.len());
 
         let old_to_new_label_map = {
             let mut label_map = MapType::default();
@@ -353,45 +424,13 @@ where
 
         let relabeled_vertices: Vec<_> = vertex_keep_table
             .iter()
+            .copied()
             .enumerate()
-            .filter_map(|(i, should_keep)| if *should_keep { Some(i) } else { None })
+            .filter_map(|(i, should_keep)| if should_keep { Some(i) } else { None })
             .map(|index| vertices[index].clone())
             .collect();
 
-        Self::from_vertices_and_connectivity(relabeled_vertices, relabeled_cells)
-    }
-
-    /// Removes all cells from the mesh that are completely outside of the given AABB and clamps the remaining cells to the boundary
-    fn par_clamp_with_aabb(&self, aabb: &Aabb3d<R>) -> Self
-    where
-        Self::Cell: Sync,
-    {
-        // Find all triangles with at least one vertex inside of AABB
-        let vertices = self.vertices();
-        let cells_to_keep = self
-            .cells()
-            .par_iter()
-            .enumerate()
-            .filter(|(_, cell)| {
-                cell.vertices()
-                    .iter()
-                    .copied()
-                    .any(|v| aabb.contains_point(&vertices[v]))
-            })
-            .map(|(i, _)| i)
-            .collect::<Vec<_>>();
-        // Remove all other cells from mesh
-        let mut new_mesh = self.keep_cells(&cells_to_keep);
-        // Clamp remaining vertices to AABB
-        new_mesh.vertices_mut().par_iter_mut().for_each(|v| {
-            let min = aabb.min();
-            let max = aabb.max();
-            v.x = v.x.clamp(min.x, max.x);
-            v.y = v.y.clamp(min.y, max.y);
-            v.z = v.z.clamp(min.z, max.z);
-        });
-
-        new_mesh
+        MeshT::from_vertices_and_connectivity(relabeled_vertices, relabeled_cells)
     }
 }
 
@@ -927,6 +966,45 @@ impl<R: Real, MeshT: Mesh3d<R>> Mesh3d<R> for MeshWithData<R, MeshT> {
             connectivity,
         ))
     }
+
+    /// Returns a new mesh containing only the specified cells and removes all unreferenced vertices and attributes
+    fn keep_cells(&self, cell_indices: &[usize], keep_all_vertices: bool) -> Self {
+        // Filter internal mesh
+        let mut new_mesh = if keep_all_vertices {
+            let mut new_mesh = keep_cells_impl(self, cell_indices, &[]);
+            new_mesh.point_attributes = self.point_attributes.clone();
+            new_mesh
+        } else {
+            let vertex_keep_table = vertex_keep_table(self, cell_indices);
+            let mut new_mesh = keep_cells_impl(self, cell_indices, &vertex_keep_table);
+
+            let vertex_indices = vertex_keep_table
+                .iter()
+                .copied()
+                .enumerate()
+                .filter_map(|(i, should_keep)| if should_keep { Some(i) } else { None })
+                .collect::<Vec<_>>();
+
+            // Filter the point attributes
+            new_mesh.point_attributes = self
+                .point_attributes
+                .iter()
+                .map(|attr| attr.keep_indices(&vertex_indices))
+                .collect();
+
+            new_mesh
+        };
+
+        // Filter the cell attributes
+        let cell_attributes = self
+            .cell_attributes
+            .iter()
+            .map(|attr| attr.keep_indices(&cell_indices))
+            .collect();
+        new_mesh.cell_attributes = cell_attributes;
+
+        new_mesh
+    }
 }
 
 /// Returns an mesh data wrapper with a default mesh and without attached attributes
@@ -1026,6 +1104,26 @@ impl<R: Real> MeshAttribute<R> {
                 .with_data(vec3r_vec.iter().flatten().copied().collect::<Vec<R>>()),
         }
     }
+
+    /// Returns a new attribute keeping only the entries with the given index
+    fn keep_indices(&self, indices: &[usize]) -> Self {
+        let data = match &self.data {
+            AttributeData::ScalarU64(d) => {
+                AttributeData::ScalarU64(indices.iter().copied().map(|i| d[i].clone()).collect())
+            }
+            AttributeData::ScalarReal(d) => {
+                AttributeData::ScalarReal(indices.iter().copied().map(|i| d[i].clone()).collect())
+            }
+            AttributeData::Vector3Real(d) => {
+                AttributeData::Vector3Real(indices.iter().copied().map(|i| d[i].clone()).collect())
+            }
+        };
+
+        Self {
+            name: self.name.clone(),
+            data,
+        }
+    }
 }
 
 impl<R: Real> AttributeData<R> {
