Improve ns-process-data images when using existing COLMAP model

nerfstudio/process_data/colmap_utils.py

@@ -604,6 +604,14 @@ def colmap_to_json(
     cameras = read_cameras_binary(cameras_path)
     images = read_images_binary(images_path)
 
+    # Images were renamed to frame_{i:05d}.{ext} and
+    # the filenames needs to be replaced in the transforms.json as well
+    original_filenames = [x.name for x in images.values()]
+    # Sort was used in nerfstudio.process_data.process_data_utils:get_image_filenames
+    original_filenames.sort()
+    # Build the map to the new filenames
+    filename_map = {name: f"frame_{i+1:05d}{os.path.splitext(name)[-1]}" for i, name in enumerate(original_filenames)}
+
     # Only supports one camera
     camera_params = cameras[1].params
 
@@ -619,7 +627,8 @@ def colmap_to_json(
         c2w = c2w[np.array([1, 0, 2, 3]), :]
         c2w[2, :] *= -1
 
-        name = Path(f"./images/{im_data.name}")
+        name = filename_map[im_data.name]
+        name = Path(f"./images/{name}")
 
         frame = {
             "file_path": name.as_posix(),

nerfstudio/process_data/process_data_utils.py

@@ -46,6 +46,19 @@ class CameraModel(Enum):
 }
 
 
+def list_images(data: Path) -> List[Path]:
+    """Lists all supported images in a directory
+
+    Args:
+        data: Path to the directory of images.
+    Returns:
+        Paths to images contained in the directory
+    """
+    allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
+    image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts])
+    return image_paths
+
+
 def get_image_filenames(directory: Path, max_num_images: int = -1) -> Tuple[List[Path], int]:
     """Returns a list of image filenames in a directory.
 
@@ -55,8 +68,7 @@ def get_image_filenames(directory: Path, max_num_images: int = -1) -> Tuple[List
     Returns:
         A tuple of A list of image filenames, number of original image paths.
     """
-    allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
-    image_paths = sorted([p for p in directory.glob("[!.]*") if p.suffix.lower() in allowed_exts])
+    image_paths = list_images(directory)
     num_orig_images = len(image_paths)
 
     if max_num_images != -1 and num_orig_images > max_num_images:
@@ -240,8 +252,7 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int:
         The number of images copied.
     """
     with status(msg="[bold yellow]Copying images...", spinner="bouncingBall", verbose=verbose):
-        allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
-        image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts])
+        image_paths = list_images(data)
 
         if len(image_paths) == 0:
             CONSOLE.log("[bold red]:skull: No usable images in the data folder.")

scripts/process_data.py

@@ -74,6 +74,10 @@ class ProcessImages:
         will downscale the images by 2x, 4x, and 8x."""
     skip_colmap: bool = False
     """If True, skips COLMAP and generates transforms.json if possible."""
+    skip_images: bool = False
+    """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled"""
+    colmap_model_path: str = "colmap/sparse/0"
+    """Optionally sets the path of the colmap model. Used only when --skip-colmap is set to True."""
     colmap_cmd: str = "colmap"
     """How to call the COLMAP executable."""
     gpu: bool = True
@@ -83,6 +87,9 @@ class ProcessImages:
 
     def main(self) -> None:
         """Process images into a nerfstudio dataset."""
+        if not self.skip_colmap and self.colmap_model_path != "colmap/sparse/0":
+            CONSOLE.log("[bold red]The --colmap-model-path can only be used when --skip-colmap is set.")
+            sys.exit(1)
         install_checks.check_ffmpeg_installed()
         install_checks.check_colmap_installed()
 
@@ -92,12 +99,22 @@ def main(self) -> None:
 
         summary_log = []
 
-        # Copy images to output directory
-        num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose)
-        summary_log.append(f"Starting with {num_frames} images")
+        # Copy and downscale images
+        if not self.skip_images:
+            # Copy images to output directory
+            num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose)
+            summary_log.append(f"Starting with {num_frames} images")
 
-        # Downscale images
-        summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose))
+            # Downscale images
+            summary_log.append(
+                process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)
+            )
+        else:
+            num_frames = len(process_data_utils.list_images(self.data))
+            if num_frames == 0:
+                CONSOLE.log("[bold red]:skull: No usable images in the data folder.")
+                sys.exit(1)
+            summary_log.append(f"Starting with {num_frames} images")
 
         # Run COLMAP
         colmap_dir = self.output_dir / "colmap"
@@ -133,11 +150,16 @@ def main(self) -> None:
                 sys.exit(1)
 
         # Save transforms.json
-        if (colmap_dir / "sparse" / "0" / "cameras.bin").exists():
+        if self.skip_colmap:
+            colmap_model_path = Path(self.colmap_model_path)
+        else:
+            colmap_model_path = colmap_dir / "sparse" / "0"
+
+        if (colmap_model_path / "cameras.bin").exists():
             with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):
                 num_matched_frames = colmap_utils.colmap_to_json(
-                    cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin",
-                    images_path=colmap_dir / "sparse" / "0" / "images.bin",
+                    cameras_path=colmap_model_path / "cameras.bin",
+                    images_path=colmap_model_path / "images.bin",
                     output_dir=self.output_dir,
                     camera_model=CAMERA_MODELS[self.camera_type],
                 )