Add a point cloud option for Record3D (#3556)

* Add a point cloud option for Record3D

* Change parameter names and add voxel size parameter

* fix record3d point cloud docs formatting

nerfstudio/process_data/record3d_utils.py

@@ -16,23 +16,32 @@
 
 import json
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
 import numpy as np
+import open3d as o3d
 from scipy.spatial.transform import Rotation
 
 from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
 from nerfstudio.utils import io
 
 
-def record3d_to_json(images_paths: List[Path], metadata_path: Path, output_dir: Path, indices: np.ndarray) -> int:
+def record3d_to_json(
+    images_paths: List[Path],
+    metadata_path: Path,
+    output_dir: Path,
+    indices: np.ndarray,
+    ply_dirname: Optional[Path],
+    voxel_size: Optional[float],
+) -> int:
     """Converts Record3D's metadata and image paths to a JSON file.
 
     Args:
-        images_paths: list if image paths.
+        images_paths: list of image paths.
         metadata_path: Path to the Record3D metadata JSON file.
         output_dir: Path to the output directory.
         indices: Indices to sample the metadata_path. Should be the same length as images_paths.
+        ply_dirname: Path to the directory of exported ply files.
 
     Returns:
         The number of registered images.
@@ -87,6 +96,23 @@ def record3d_to_json(images_paths: List[Path], metadata_path: Path, output_dir:
 
     out["frames"] = frames
 
+    # If .ply directory exists add the sparse point cloud for gsplat point initialization
+    if ply_dirname is not None:
+        assert ply_dirname.exists(), f"Directory not found: {ply_dirname}"
+        assert ply_dirname.is_dir(), f"Path given is not a directory: {ply_dirname}"
+
+        # Create sparce point cloud
+        pcd = o3d.geometry.PointCloud()
+        for ply_filename in ply_dirname.iterdir():
+            temp_pcd = o3d.io.read_point_cloud(str(ply_filename))
+            pcd += temp_pcd.voxel_down_sample(voxel_size=voxel_size)
+
+        # Save point cloud
+        points3D = np.asarray(pcd.points)
+        pcd.points = o3d.utility.Vector3dVector(points3D)
+        o3d.io.write_point_cloud(str(output_dir / "sparse_pc.ply"), pcd, write_ascii=True)
+        out["ply_file_path"] = "sparse_pc.ply"
+
     with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
         json.dump(out, f, indent=4)
 

nerfstudio/scripts/process_data.py

@@ -49,6 +49,11 @@ class ProcessRecord3D(BaseConverterToNerfstudioDataset):
     2. Converts Record3D poses into the nerfstudio format.
     """
 
+    ply_dir: Optional[Path] = None
+    """Path to the Record3D directory of point export ply files."""
+    voxel_size: Optional[float] = 0.8
+    """Voxel size for down sampling dense point cloud"""
+
     num_downscales: int = 3
     """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3
         will downscale the images by 2x, 4x, and 8x."""
@@ -101,7 +106,14 @@ class ProcessRecord3D(BaseConverterToNerfstudioDataset):
             )
 
         metadata_path = self.data / "metadata.json"
-        record3d_utils.record3d_to_json(copied_image_paths, metadata_path, self.output_dir, indices=idx)
+        record3d_utils.record3d_to_json(
+            copied_image_paths,
+            metadata_path,
+            self.output_dir,
+            indices=idx,
+            ply_dirname=self.ply_dir,
+            voxel_size=self.voxel_size,
+        )
         CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:")
 
         for summary in summary_log: