Switching the gaussian-splatting dataparser from ColmapDataParser to NerfStudioDataparser and updating ns-process-data

nerfstudio/configs/method_configs.py

@@ -31,7 +31,6 @@
 from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig
 from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
-from nerfstudio.data.dataparsers.colmap_dataparser import ColmapDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
 from nerfstudio.data.dataparsers.instant_ngp_dataparser import InstantNGPDataParserConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
@@ -600,7 +599,7 @@
     gradient_accumulation_steps={"camera_opt": 100},
     pipeline=VanillaPipelineConfig(
         datamanager=FullImageDatamanagerConfig(
-            dataparser=ColmapDataParserConfig(load_3D_points=True),
+            dataparser=NerfstudioDataParserConfig(load_3D_points=True),
         ),
         model=GaussianSplattingModelConfig(),
     ),

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -73,6 +73,8 @@ class NerfstudioDataParserConfig(DataParserConfig):
     """The interval between frames to use for eval. Only used when eval_mode is eval-interval."""
     depth_unit_scale_factor: float = 1e-3
     """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
+    load_3D_points: bool = False
+    """Whether to load the 3D points from the colmap reconstruction."""
 
 
 @dataclass
@@ -314,11 +316,60 @@ def _generate_dataparser_outputs(self, split="train"):
             applied_scale = float(meta["applied_scale"])
             scale_factor *= applied_scale
 
-        # Load 3D points
+        # reinitialize metadata for dataparser_outputs
         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))
+
+        # _generate_dataparser_outputs might be called more than once so we check if we already loaded the point cloud
+        try:
+            self.prompted_user
+        except AttributeError:
+            self.prompted_user = False
+
+        # Load 3D points
+        if self.config.load_3D_points:
+            colmap_path = self.config.data / "colmap/sparse/0"
+
+            if "ply_file_path" in meta:
+                ply_file_path = data_dir / meta["ply_file_path"]
+
+            elif colmap_path.exists():
+                from rich.prompt import Confirm
+
+                # check if user wants to make a point cloud from colmap points
+                if not self.prompted_user:
+                    self.create_pc = Confirm.ask(
+                        "load_3D_points is true, but the dataset was processed with an outdated ns-process-data that didn't convert colmap points to .ply! Update the colmap dataset automatically?"
+                    )
+
+                if self.create_pc:
+                    import json
+
+                    from nerfstudio.process_data.colmap_utils import create_ply_from_colmap
+
+                    with open(self.config.data / "transforms.json") as f:
+                        transforms = json.load(f)
+
+                    ply_filename = "sparse_pc.ply"
+                    create_ply_from_colmap(filename=ply_filename, recon_dir=colmap_path, output_dir=self.config.data)
+                    ply_file_path = data_dir / ply_filename
+                    transforms["ply_file_path"] = ply_filename
+
+                    with open(self.config.data / "transforms.json", "w", encoding="utf-8") as f:
+                        json.dump(transforms, f, indent=4)
+                else:
+                    ply_file_path = None
+            else:
+                if not self.prompted_user:
+                    CONSOLE.print(
+                        "[bold yellow]Warning: load_3D_points set to true but no point cloud found. gaussian-splatting models will use random point cloud initialization."
+                    )
+                ply_file_path = None
+
+            if ply_file_path:
+                sparse_points = self._load_3D_points(ply_file_path, transform_matrix, scale_factor)
+                if sparse_points is not None:
+                    metadata.update(sparse_points)
+            self.prompted_user = True
 
         dataparser_outputs = DataparserOutputs(
             image_filenames=image_filenames,
@@ -336,10 +387,24 @@ def _generate_dataparser_outputs(self, split="train"):
         return dataparser_outputs
 
     def _load_3D_points(self, ply_file_path: Path, transform_matrix: torch.Tensor, scale_factor: float):
+        """Loads point clouds positions and colors from .ply
+
+        Args:
+            ply_file_path: Path to .ply file
+            transform_matrix: Matrix to transform world coordinates
+            scale_factor: How much to scale the camera origins by.
+
+        Returns:
+            A dictionary of points: points3D_xyz and colors: points3D_rgb
+        """
         import open3d as o3d  # Importing open3d is slow, so we only do it if we need it.
 
         pcd = o3d.io.read_point_cloud(str(ply_file_path))
 
+        # if no points found don't read in an initial point cloud
+        if len(pcd.points) == 0:
+            return None
+
         points3D = torch.from_numpy(np.asarray(pcd.points, dtype=np.float32))
         points3D = (
             torch.cat(

nerfstudio/process_data/colmap_utils.py

@@ -34,6 +34,7 @@
     read_cameras_binary,
     read_images_binary,
     read_points3D_binary,
+    read_points3D_text,
 )
 from nerfstudio.process_data.process_data_utils import CameraModel
 from nerfstudio.utils import colormaps
@@ -391,6 +392,7 @@ def colmap_to_json(
     camera_mask_path: Optional[Path] = None,
     image_id_to_depth_path: Optional[Dict[int, Path]] = None,
     image_rename_map: Optional[Dict[str, str]] = None,
+    ply_filename="sparse_pc.ply",
 ) -> int:
     """Converts COLMAP's cameras.bin and images.bin to a JSON file.
 
@@ -460,6 +462,11 @@ def colmap_to_json(
     applied_transform[2, :] *= -1
     out["applied_transform"] = applied_transform.tolist()
 
+    # create ply from colmap
+    assert ply_filename.endswith(".ply"), f"ply_filename: {ply_filename} does not end with '.ply'"
+    create_ply_from_colmap(ply_filename, recon_dir, output_dir)
+    out["ply_file_path"] = ply_filename
+
     with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
         json.dump(out, f, indent=4)
 
@@ -638,3 +645,43 @@ def get_matching_summary(num_initial_frames: int, num_matched_frames: int) -> st
         result += " or large exposure changes."
         return result
     return f"[bold green]COLMAP found poses for {num_matched_frames / num_initial_frames * 100:.2f}% of the images."
+
+
+def create_ply_from_colmap(filename: str, recon_dir: Path, output_dir: Path):
+    """Writes a ply file from colmap.
+
+    Args:
+        filename: file name for .ply
+        recon_dir: Directory to grab colmap points
+        output_dir: Directory to output .ply
+    """
+    if (recon_dir / "points3D.bin").exists():
+        colmap_points = read_points3D_binary(recon_dir / "points3D.bin")
+    elif (recon_dir / "points3D.txt").exists():
+        colmap_points = read_points3D_text(recon_dir / "points3D.txt")
+    else:
+        raise ValueError(f"Could not find points3D.txt or points3D.bin in {recon_dir}")
+
+    # Load point Positions
+    points3D = torch.from_numpy(np.array([p.xyz for p in colmap_points.values()], dtype=np.float32))
+    # Load point colours
+    points3D_rgb = torch.from_numpy(np.array([p.rgb for p in colmap_points.values()], dtype=np.uint8))
+
+    # write ply
+    with open(output_dir / filename, "w") as f:
+        # Header
+        f.write("ply\n")
+        f.write("format ascii 1.0\n")
+        f.write(f"element vertex {len(points3D)}\n")
+        f.write("property float x\n")
+        f.write("property float y\n")
+        f.write("property float z\n")
+        f.write("property uint8 red\n")
+        f.write("property uint8 green\n")
+        f.write("property uint8 blue\n")
+        f.write("end_header\n")
+
+        for coord, color in zip(points3D, points3D_rgb):
+            x, y, z = coord
+            r, g, b = color
+            f.write(f"{x:8f} {y:8f} {z:8f} {r} {g} {b}\n")

tests/process_data/test_process_images.py

@@ -12,9 +12,11 @@
 from nerfstudio.data.utils.colmap_parsing_utils import (
     Camera,
     Image as ColmapImage,
+    Point3D,
     qvec2rotmat,
     write_cameras_binary,
     write_images_binary,
+    write_points3D_binary,
 )
 from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset
 
@@ -50,6 +52,19 @@ def test_process_images_skip_colmap(tmp_path: Path):
         {1: Camera(1, "OPENCV", width, height, [110, 110, 50, 75, 0, 0, 0, 0, 0, 0])},
         sparse_path / "cameras.bin",
     )
+    write_points3D_binary(
+        {
+            1: Point3D(
+                id=1,
+                xyz=np.array([0, 0, 0]),
+                rgb=np.array([0, 0, 0]),
+                error=np.array([0]),
+                image_ids=np.array([1]),
+                point2D_idxs=np.array([0]),
+            ),
+        },
+        sparse_path / "points3D.bin",
+    )
     frames = {}
     num_frames = 10
     qvecs = random_quaternion(num_frames)