Support multiple camera modes

nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py

@@ -31,6 +31,8 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset):
 
     camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective"
     """Camera model to use."""
+    camera_mode: Literal["auto", "single", "per_image"] = "single"
+    """What camera mode to use for the reconstruction. Only works with hloc sfm_tool"""
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed
     and accuracy. Exhaustive is slower but more accurate. Sequential is faster but
@@ -219,6 +221,7 @@ def _run_colmap(self, mask_path: Optional[Path] = None):
             assert feature_type is not None
             assert matcher_type is not None
             assert matcher_type != "NN"  # Only used for colmap.
+
             hloc_utils.run_hloc(
                 image_dir=image_dir,
                 colmap_dir=self.absolute_colmap_path,
@@ -228,6 +231,7 @@ def _run_colmap(self, mask_path: Optional[Path] = None):
                 feature_type=feature_type,
                 matcher_type=matcher_type,
                 refine_pixsfm=self.refine_pixsfm,
+                camera_mode=camera_mode,
             )
         else:
             raise RuntimeError("Invalid combination of sfm_tool, feature_type, and matcher_type, " "exiting")

nerfstudio/process_data/colmap_utils.py

@@ -451,9 +451,17 @@ def colmap_to_json(
         frames.append(frame)
 
     if set(cam_id_to_camera.keys()) != {1}:
-        raise RuntimeError("Only single camera shared for all images is supported.")
-    out = parse_colmap_camera_params(cam_id_to_camera[1])
-    out["frames"] = frames
+        cameras = {k: parse_colmap_camera_params(v) for k, v in cam_id_to_camera.items()}
+        camera_models = list(set(x["camera_model"] for x in cameras.values()))
+        assert len(camera_models) == 1, "All cameras should be of the same model"
+
+        for i, frame in enumerate(frames):
+            frames[i] = {**frame, **cameras[frame["colmap_im_id"]]}
+
+        out = {"camera_model": camera_models[0], "frames": frames}
+    else:
+        out = parse_colmap_camera_params(cam_id_to_camera[1])
+        out["frames"] = frames
 
     applied_transform = np.eye(4)[:3, :]
     applied_transform = applied_transform[np.array([1, 0, 2]), :]

nerfstudio/process_data/hloc_utils.py

@@ -71,6 +71,7 @@ def run_hloc(
     ] = "superglue",
     num_matched: int = 50,
     refine_pixsfm: bool = False,
+    camera_mode: Literal["auto", "single", "per_image"] = "single"
 ) -> None:
     """Runs hloc on the images.
 
@@ -85,6 +86,7 @@ def run_hloc(
         matcher_type: Type of feature matcher to use.
         num_matched: Number of image pairs for loc.
         refine_pixsfm: If True, refine the reconstruction using pixel-perfect-sfm.
+        camera_mode: What camera mode to use for the reconstruction
     """
     if not _HAS_HLOC:
         CONSOLE.print(
@@ -119,6 +121,8 @@ def run_hloc(
     match_features.main(matcher_conf, sfm_pairs, features=features, matches=matches)  # type: ignore
 
     image_options = pycolmap.ImageReaderOptions(camera_model=camera_model.value)  # type: ignore
+
+    camera_mode = pycolmap.CameraMode(camera_mode.upper())
     if refine_pixsfm:
         sfm = PixSfM(  # type: ignore
             conf={
@@ -134,7 +138,7 @@ def run_hloc(
             features,
             matches,
             image_list=references,
-            camera_mode=pycolmap.CameraMode.SINGLE,  # type: ignore
+            camera_mode=camera_mode,
             image_options=image_options,
             verbose=verbose,
         )
@@ -147,7 +151,7 @@ def run_hloc(
             sfm_pairs,
             features,
             matches,
-            camera_mode=pycolmap.CameraMode.SINGLE,  # type: ignore
+            camera_mode=camera_mode,
             image_options=image_options,
             verbose=verbose,
         )