Merge branch 'gaussian-splatting' of https://github.com/nerfstudio-project/nerfstudio into gaussian-splatting

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -20,6 +20,7 @@ from pathlib import Path
 from typing import Literal, Optional, Type
 
 import numpy as np
+import open3d as o3d
 import torch
 from PIL import Image
 
@@ -309,6 +310,12 @@ class Nerfstudio(DataParser):
             applied_scale = float(meta["applied_scale"])
             scale_factor *= applied_scale
 
+        # Load 3D points
+        metadata = {}
+        if "ply_file_path" in meta:
+            ply_file_path = data_dir / meta["ply_file_path"]
+            metadata.update(self._load_3D_points(ply_file_path, transform_matrix, scale_factor))
+
         dataparser_outputs = DataparserOutputs(
             image_filenames=image_filenames,
             cameras=cameras,
@@ -319,10 +326,34 @@ class Nerfstudio(DataParser):
             metadata={
                 "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
                 "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
+                **metadata,
             },
         )
         return dataparser_outputs
 
+    def _load_3D_points(self, ply_file_path: Path, transform_matrix: torch.Tensor, scale_factor: float):
+        pcd = o3d.io.read_point_cloud(str(ply_file_path))
+
+        points3D = torch.from_numpy(np.asarray(pcd.points, dtype=np.float32))
+        points3D = (
+            torch.cat(
+                (
+                    points3D,
+                    torch.ones_like(points3D[..., :1]),
+                ),
+                -1,
+            )
+            @ transform_matrix.T
+        )
+        points3D *= scale_factor
+        points3D_rgb = torch.from_numpy((np.asarray(pcd.colors) * 255).astype(np.uint8))
+
+        out = {
+            "points3D_xyz": points3D,
+            "points3D_rgb": points3D_rgb,
+        }
+        return out
+
     def _get_fname(self, filepath: Path, data_dir: Path, downsample_folder_prefix="images_") -> Path:
         """Get the filename of the image file.
         downsample_folder_prefix can be used to point to auxiliary image data, e.g. masks

nerfstudio/scripts/exporter.py

@@ -121,6 +121,9 @@ class ExportPointCloud(Exporter):
     """Number of rays to evaluate per batch. Decrease if you run out of memory."""
     std_ratio: float = 10.0
     """Threshold based on STD of the average distances across the point cloud to remove outliers."""
+    save_world_frame: bool = True
+    """If true, saves in the frame of the transform.json file, if false saves in the frame of the scaled 
+        dataparser transform"""
 
     def main(self) -> None:
         """Export point cloud."""
@@ -157,6 +160,15 @@ class ExportPointCloud(Exporter):
             crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
+        if self.save_world_frame:
+            # apply the inverse dataparser transform to the point cloud
+            points = np.asarray(pcd.points)
+            poses = np.eye(4,dtype=np.float32)[None,...].repeat(points.shape[0], axis=0)[:,:3,:]
+            poses[:,:3,3] = points
+            poses = pipeline.datamanager.train_dataparser_outputs.transform_poses_to_original_space(torch.from_numpy(poses))
+            points = poses[:,:3,3].numpy()
+            pcd.points = o3d.utility.Vector3dVector(points)
+
         torch.cuda.empty_cache()
 
         CONSOLE.print(f"[bold green]:white_check_mark: Generated {pcd}")

nerfstudio/viewer_beta/export_panel.py

@@ -81,6 +81,7 @@ def populate_point_cloud_tab(
 ) -> None:
     server.add_gui_markdown("<small>Render depth, project to an oriented point cloud, and filter.</small> ")
     num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1)
+    world_frame = server.add_gui_checkbox("Save in world frame", True, hint="Save the point cloud in the transforms.json frame, rather than scaled scene frame")
     remove_outliers = server.add_gui_checkbox("Remove outliers", True)
     normals = server.add_gui_dropdown(
         "Normals",
@@ -104,6 +105,7 @@ def populate_point_cloud_tab(
                 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),
             ]
         )