From 0397ee8e54965864ad3333375580c27a545ac131 Mon Sep 17 00:00:00 2001
From: jamesoncrate <jamesoncrate@gmail.com>
Date: Sun, 22 Dec 2024 12:05:53 -0800
Subject: [PATCH] Change parameter names and add voxel size parameter

---
 docs/quickstart/custom_dataset.md         |  4 ++++
 nerfstudio/process_data/record3d_utils.py |  5 +++--
 nerfstudio/scripts/process_data.py        | 12 +++++++++---
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md
index dbda0bf3e2..fbe65f93e9 100644
--- a/docs/quickstart/custom_dataset.md
+++ b/docs/quickstart/custom_dataset.md
@@ -293,6 +293,10 @@ follow these steps:
 ns-process-data record3d --data {data directory} --ply {ply directory} --output-dir {output directory}
 ```
 
+Additionally you can specify `--voxel-size {float}` which determines the level of sparsity when downsampling from the dense
+point clouds generated by Record3D to the sparse point cloud used in Nerfstudio. The default value is 0.8, lower is less
+sparse, higher is more sparse.
+
 (spectacularai)=
 
 ## Spectacular AI
diff --git a/nerfstudio/process_data/record3d_utils.py b/nerfstudio/process_data/record3d_utils.py
index 407826390a..50212c2b20 100644
--- a/nerfstudio/process_data/record3d_utils.py
+++ b/nerfstudio/process_data/record3d_utils.py
@@ -31,7 +31,8 @@ def record3d_to_json(
     metadata_path: Path,
     output_dir: Path,
     indices: np.ndarray,
-    ply_dirname: Optional[Path] = None,
+    ply_dirname: Optional[Path],
+    voxel_size: Optional[float],
 ) -> int:
     """Converts Record3D's metadata and image paths to a JSON file.
 
@@ -104,7 +105,7 @@ def record3d_to_json(
         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=0.8)
+            pcd += temp_pcd.voxel_down_sample(voxel_size=voxel_size)
 
         # Save point cloud
         points3D = np.asarray(pcd.points)
diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py
index d734f67927..1fdd36f7f2 100644
--- a/nerfstudio/scripts/process_data.py
+++ b/nerfstudio/scripts/process_data.py
@@ -49,8 +49,10 @@ class ProcessRecord3D(BaseConverterToNerfstudioDataset):
     2. Converts Record3D poses into the nerfstudio format.
     """
 
-    ply: Optional[Path] = None
+    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
@@ -104,9 +106,13 @@ def main(self) -> None:
             )
 
         metadata_path = self.data / "metadata.json"
-        ply_path = self.ply if hasattr(self, "ply") else None
         record3d_utils.record3d_to_json(
-            copied_image_paths, metadata_path, self.output_dir, indices=idx, ply_dirname=ply_path
+            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:")