Adding cropping for gaussian splat exports (#3105)

* added bounding box to splat export params

Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>

nerfstudio/exporter/exporter_utils.py

@@ -89,9 +89,6 @@ def generate_point_cloud(
     rgb_output_name: str = "rgb",
     depth_output_name: str = "depth",
     normal_output_name: Optional[str] = None,
-    use_bounding_box: bool = True,
-    bounding_box_min: Optional[Tuple[float, float, float]] = None,
-    bounding_box_max: Optional[Tuple[float, float, float]] = None,
     crop_obb: Optional[OrientedBox] = None,
     std_ratio: float = 10.0,
 ) -> o3d.geometry.PointCloud:
@@ -106,9 +103,6 @@ def generate_point_cloud(
         rgb_output_name: Name of the RGB output.
         depth_output_name: Name of the depth output.
         normal_output_name: Name of the normal output.
-        use_bounding_box: Whether to use a bounding box to sample points.
-        bounding_box_min: Minimum of the bounding box.
-        bounding_box_max: Maximum of the bounding box.
         std_ratio: Threshold based on STD of the average distances across the point cloud to remove outliers.
 
     Returns:
@@ -126,8 +120,6 @@ def generate_point_cloud(
     rgbs = []
     normals = []
     view_directions = []
-    if use_bounding_box and (crop_obb is not None and bounding_box_max is not None):
-        CONSOLE.print("Provided aabb and crop_obb at the same time, using only the obb", style="bold yellow")
     with progress as progress_bar:
         task = progress_bar.add_task("Generating Point Cloud", total=num_points)
         while not progress_bar.finished:
@@ -171,21 +163,13 @@ def generate_point_cloud(
             if normal is not None:
                 normal = normal[mask]
 
-            if use_bounding_box:
-                if crop_obb is None:
-                    comp_l = torch.tensor(bounding_box_min, device=point.device)
-                    comp_m = torch.tensor(bounding_box_max, device=point.device)
-                    assert torch.all(
-                        comp_l < comp_m
-                    ), f"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}"
-                    mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
-                else:
-                    mask = crop_obb.within(point)
-                point = point[mask]
-                rgb = rgb[mask]
-                view_direction = view_direction[mask]
-                if normal is not None:
-                    normal = normal[mask]
+            if crop_obb is not None:
+                mask = crop_obb.within(point)
+            point = point[mask]
+            rgb = rgb[mask]
+            view_direction = view_direction[mask]
+            if normal is not None:
+                normal = normal[mask]
 
             points.append(point)
             rgbs.append(rgb)

nerfstudio/scripts/exporter.py

@@ -107,12 +107,6 @@ class ExportPointCloud(Exporter):
     """Name of the depth output."""
     rgb_output_name: str = "rgb"
     """Name of the RGB output."""
-    use_bounding_box: bool = True
-    """Only query points within the bounding box"""
-    bounding_box_min: Optional[Tuple[float, float, float]] = (-1, -1, -1)
-    """Minimum of the bounding box, used if use_bounding_box is True."""
-    bounding_box_max: Optional[Tuple[float, float, float]] = (1, 1, 1)
-    """Maximum of the bounding box, used if use_bounding_box is True."""
 
     obb_center: Optional[Tuple[float, float, float]] = None
     """Center of the oriented bounding box."""
@@ -160,9 +154,6 @@ class ExportPointCloud(Exporter):
             rgb_output_name=self.rgb_output_name,
             depth_output_name=self.depth_output_name,
             normal_output_name=self.normal_output_name if self.normal_method == "model_output" else None,
-            use_bounding_box=self.use_bounding_box,
-            bounding_box_min=self.bounding_box_min,
-            bounding_box_max=self.bounding_box_max,
             crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
@@ -346,9 +337,6 @@ class ExportPoissonMesh(Exporter):
             rgb_output_name=self.rgb_output_name,
             depth_output_name=self.depth_output_name,
             normal_output_name=self.normal_output_name if self.normal_method == "model_output" else None,
-            use_bounding_box=self.use_bounding_box,
-            bounding_box_min=self.bounding_box_min,
-            bounding_box_max=self.bounding_box_max,
             crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
@@ -491,6 +479,13 @@ class ExportGaussianSplat(Exporter):
     Export 3D Gaussian Splatting model to a .ply
     """
 
+    obb_center: Optional[Tuple[float, float, float]] = None
+    """Center of the oriented bounding box."""
+    obb_rotation: Optional[Tuple[float, float, float]] = None
+    """Rotation of the oriented bounding box. Expressed as RPY Euler angles in radians"""
+    obb_scale: Optional[Tuple[float, float, float]] = None
+    """Scale of the oriented bounding box along each axis."""
+
     @staticmethod
     def write_ply(filename: str, count: int, map_to_tensors: OrderedDict[str, np.ndarray]):
         """
@@ -585,6 +580,16 @@ class ExportGaussianSplat(Exporter):
             for i in range(4):
                 map_to_tensors[f"rot_{i}"] = quats[:, i, None]
 
+            if self.obb_center is not None and self.obb_rotation is not None and self.obb_scale is not None:
+                crop_obb = OrientedBox.from_params(self.obb_center, self.obb_rotation, self.obb_scale)
+                assert crop_obb is not None
+                mask = crop_obb.within(torch.from_numpy(positions)).numpy()
+                for k, t in map_to_tensors.items():
+                    map_to_tensors[k] = map_to_tensors[k][mask]
+
+                n = map_to_tensors["x"].shape[0]
+                count = n
+
         # post optimization, it is possible have NaN/Inf values in some attributes
         # to ensure the exported ply file has finite values, we enforce finite filters.
         select = np.ones(n, dtype=bool)

nerfstudio/viewer/export_panel.py

@@ -73,10 +73,12 @@ def show_command_modal(client: viser.ClientHandle, what: Literal["mesh", "point
             modal.close()
 
 
-def get_crop_string(obb: OrientedBox):
+def get_crop_string(obb: OrientedBox, crop_viewport: bool):
     """Takes in an oriented bounding box and returns a string of the form "--obb_{center,rotation,scale}
     and each arg formatted with spaces around it
     """
+    if not crop_viewport:
+        return ""
     rpy = vtf.SO3.from_matrix(obb.R.numpy(force=True)).as_rpy_radians()
     pos = obb.T.squeeze().tolist()
     scale = obb.S.squeeze().tolist()
@@ -125,9 +127,8 @@ def populate_point_cloud_tab(
                     f"--num-points {num_points.value}",
                     f"--remove-outliers {remove_outliers.value}",
                     f"--normal-method {normals.value}",
-                    f"--use_bounding_box {control_panel.crop_viewport}",
                     f"--save-world-frame {world_frame.value}",
-                    get_crop_string(control_panel.crop_obb),
+                    get_crop_string(control_panel.crop_obb, control_panel.crop_viewport),
                 ]
             )
             show_command_modal(event.client, "point cloud", command)
@@ -174,8 +175,7 @@ def populate_mesh_tab(
                     f"--num-points {num_points.value}",
                     f"--remove-outliers {remove_outliers.value}",
                     f"--normal-method {normals.value}",
-                    f"--use_bounding_box {control_panel.crop_viewport}",
-                    get_crop_string(control_panel.crop_obb),
+                    get_crop_string(control_panel.crop_obb, control_panel.crop_viewport),
                 ]
             )
             show_command_modal(event.client, "mesh", command)
@@ -204,6 +204,7 @@ def populate_splat_tab(
                     "ns-export gaussian-splat",
                     f"--load-config {config_path}",
                     f"--output-dir {output_directory.value}",
+                    get_crop_string(control_panel.crop_obb, control_panel.crop_viewport),
                 ]
             )
             show_command_modal(event.client, "splat", command)