Fix COLMAP sparse model selection when multiple reconstructions exist

nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py

@@ -79,7 +79,7 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset):
     """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled"""
     colmap_model_path: Path = Path("colmap/sparse/0")
     """Optionally sets the path of the colmap model. Used only when --skip-colmap is set to True. The path is relative
-       to the output directory.
+       to the output directory. Non-default model paths are preserved; the default selects the best sparse model.
     """
     colmap_cmd: str = "colmap"
     """How to call the COLMAP executable."""
@@ -117,6 +117,19 @@ def absolute_colmap_model_path(self) -> Path:
     def absolute_colmap_path(self) -> Path:
         return self.output_dir / "colmap"
 
+    def _select_colmap_model(self) -> None:
+        """Select the best generated COLMAP model when using the default path."""
+        if self.colmap_model_path != self.default_colmap_path():
+            return
+        sparse_dir = self.absolute_colmap_path / "sparse"
+        selected_model = colmap_utils.find_best_colmap_model(sparse_dir)
+        if selected_model is None:
+            return
+        relative_model_path = selected_model.relative_to(sparse_dir)
+        self.colmap_model_path = Path("colmap") / "sparse" / relative_model_path
+        num_images = colmap_utils.get_colmap_model_image_count(selected_model)
+        CONSOLE.log(f"Using COLMAP model {self.colmap_model_path} ({num_images} registered images)")
+
     def _save_transforms(
         self,
         num_frames: int,
@@ -160,9 +173,7 @@ def _export_depth(self) -> Tuple[Optional[Dict[int, Path]], List[str]]:
             depth_dir = self.output_dir / "depth"
             depth_dir.mkdir(parents=True, exist_ok=True)
             image_id_to_depth_path = colmap_utils.create_sfm_depth(
-                recon_dir=self.absolute_colmap_model_path
-                if self.skip_colmap
-                else self.output_dir / self.default_colmap_path(),
+                recon_dir=self.absolute_colmap_model_path,
                 output_dir=depth_dir,
                 include_depth_debug=self.include_depth_debug,
                 input_images_dir=self.image_dir,

nerfstudio/process_data/colmap_utils.py

@@ -17,8 +17,9 @@
 """
 
 import json
+import struct
 from pathlib import Path
-from typing import Any, Dict, Literal, Optional, Union
+from typing import Any, Dict, Literal, Optional, Tuple, Union
 
 import appdirs
 import cv2
@@ -89,6 +90,59 @@ def get_vocab_tree() -> Path:
     return vocab_tree_filename
 
 
+def _read_colmap_model(recon_dir: Path) -> Tuple[Dict[int, Any], Dict[int, Any]]:
+    """Read the binary camera and registered-image tables from a COLMAP model."""
+    cameras_path = recon_dir / "cameras.bin"
+    images_path = recon_dir / "images.bin"
+    if not cameras_path.exists() or not images_path.exists():
+        raise ValueError(f"Could not find cameras.bin and images.bin in {recon_dir}")
+    return read_cameras_binary(cameras_path), read_images_binary(images_path)
+
+
+def get_colmap_model_image_count(recon_dir: Path) -> int:
+    """Return the number of registered images in a usable COLMAP model."""
+    _, images = _read_colmap_model(recon_dir)
+    if not (recon_dir / "points3D.bin").exists():
+        raise ValueError(f"Could not find points3D.bin in {recon_dir}")
+    if not images:
+        raise ValueError(f"COLMAP model {recon_dir} contains no registered images")
+    return len(images)
+
+
+def _iter_colmap_model_paths(sparse_dir: Path):
+    """Yield direct sparse models and the known one-level HLoc nested models."""
+    for model_dir in sorted((path for path in sparse_dir.iterdir() if path.is_dir()), key=lambda path: path.name):
+        yield model_dir
+        nested_dir = model_dir / "models"
+        if nested_dir.is_dir():
+            yield from sorted((path for path in nested_dir.iterdir() if path.is_dir()), key=lambda path: path.name)
+
+
+def find_best_colmap_model(sparse_dir: Path) -> Optional[Path]:
+    """Select the best readable reconstruction in a COLMAP sparse directory.
+
+    Registered image count is the primary quality signal. Exact ties use the
+    relative model path for deterministic output. Invalid or incomplete models
+    are skipped with a warning.
+    """
+    if not sparse_dir.is_dir():
+        return None
+
+    candidates = []
+    for model_dir in _iter_colmap_model_paths(sparse_dir):
+        try:
+            num_images = get_colmap_model_image_count(model_dir)
+        except (OSError, ValueError, KeyError, IndexError, AssertionError, struct.error) as exc:
+            CONSOLE.print(f"[bold yellow]Skipping invalid COLMAP model {model_dir}: {exc}")
+            continue
+        relative_path = model_dir.relative_to(sparse_dir).as_posix()
+        candidates.append((num_images, relative_path, model_dir))
+
+    if not candidates:
+        return None
+    return sorted(candidates, key=lambda candidate: (-candidate[0], candidate[1]))[0][2]
+
+
 def run_colmap(
     image_dir: Path,
     colmap_dir: Path,
@@ -173,11 +227,14 @@ def run_colmap(
     CONSOLE.log("[bold green]:tada: Done COLMAP bundle adjustment.")
 
     if refine_intrinsics:
+        selected_model = find_best_colmap_model(sparse_dir)
+        if selected_model is None:
+            raise RuntimeError(f"Could not find a valid COLMAP reconstruction in {sparse_dir}.")
         with status(msg="[bold yellow]Refine intrinsics...", spinner="dqpb", verbose=verbose):
             bundle_adjuster_cmd = [
                 f"{colmap_cmd} bundle_adjuster",
-                f"--input_path {sparse_dir}/0",
-                f"--output_path {sparse_dir}/0",
+                f"--input_path {selected_model}",
+                f"--output_path {selected_model}",
                 "--BundleAdjustment.refine_principal_point 1",
             ]
             run_command(" ".join(bundle_adjuster_cmd), verbose=verbose)

nerfstudio/process_data/images_to_nerfstudio_dataset.py

@@ -118,6 +118,11 @@ def main(self) -> None:
             # Colmap uses renamed images
             image_rename_map = None
 
+        # COLMAP may produce multiple disconnected sparse reconstructions.
+        # Non-default --colmap-model-path values are preserved; the default
+        # selects the best readable reconstruction before exporting anything.
+        self._select_colmap_model()
+
         # Export depth maps
         image_id_to_depth_path, log_tmp = self._export_depth()
         summary_log += log_tmp

nerfstudio/process_data/video_to_nerfstudio_dataset.py

@@ -140,6 +140,11 @@ def main(self) -> None:
         if not self.skip_colmap:
             self._run_colmap(mask_path)
 
+        # COLMAP may produce multiple disconnected sparse reconstructions.
+        # Non-default --colmap-model-path values are preserved; the default
+        # selects the best readable reconstruction before exporting anything.
+        self._select_colmap_model()
+
         # Export depth maps
         image_id_to_depth_path, log_tmp = self._export_depth()
         summary_log += log_tmp

tests/process_data/test_colmap_model_selection.py

@@ -0,0 +1,134 @@
+"""Tests for selecting a COLMAP reconstruction when multiple models are present."""
+
+import json
+import os
+from pathlib import Path
+from typing import List
+
+import numpy as np
+from PIL import Image
+
+from nerfstudio.data.utils.colmap_parsing_utils import (
+    Camera,
+    Image as ColmapImage,
+    Point3D,
+    write_cameras_binary,
+    write_images_binary,
+    write_points3D_binary,
+)
+from nerfstudio.process_data.colmap_utils import find_best_colmap_model, get_colmap_model_image_count
+from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset
+
+
+def _write_binary_model(model_dir: Path, image_names: List[str], num_points: int) -> None:
+    model_dir.mkdir(parents=True, exist_ok=True)
+    write_cameras_binary(
+        {1: Camera(1, "PINHOLE", 32, 32, np.array([25.0, 25.0, 16.0, 16.0]))},
+        model_dir / "cameras.bin",
+    )
+    images = {
+        image_id: ColmapImage(
+            image_id,
+            np.array([1.0, 0.0, 0.0, 0.0]),
+            np.array([float(image_id), 0.0, 0.0]),
+            1,
+            image_name,
+            np.empty((0, 2)),
+            np.empty((0,), dtype=np.int64),
+        )
+        for image_id, image_name in enumerate(image_names, start=1)
+    }
+    write_images_binary(images, model_dir / "images.bin")
+    points = {
+        point_id: Point3D(
+            point_id,
+            np.array([float(point_id), 0.0, 1.0]),
+            np.array([255, 0, 0], dtype=np.uint8),
+            0.0,
+            np.empty((0,), dtype=np.int32),
+            np.empty((0,), dtype=np.int32),
+        )
+        for point_id in range(1, num_points + 1)
+    }
+    write_points3D_binary(points, model_dir / "points3D.bin")
+
+
+def _make_source_images(source_dir: Path, num_images: int) -> List[str]:
+    source_dir.mkdir()
+    image_names = [f"image_{i}.png" for i in range(num_images)]
+    for image_name in image_names:
+        Image.new("RGB", (32, 32)).save(source_dir / image_name)
+    return image_names
+
+
+def _mock_external_tools(tmp_path: Path, monkeypatch) -> None:
+    mocked_bin = tmp_path / "bin"
+    mocked_bin.mkdir()
+    for executable in ("colmap", "ffmpeg"):
+        path = mocked_bin / executable
+        path.touch(mode=0o755)
+    monkeypatch.setenv("PATH", f"{mocked_bin}{os.pathsep}{os.environ['PATH']}")
+
+
+def test_find_best_colmap_model_prefers_registered_images(tmp_path: Path):
+    sparse_dir = tmp_path / "sparse"
+    _write_binary_model(sparse_dir / "0", ["a.png", "b.png", "c.png"], 100)
+    _write_binary_model(sparse_dir / "5", ["a.png", "b.png", "c.png"], 101)
+    _write_binary_model(sparse_dir / "2", ["a.png", "b.png", "c.png", "d.png"], 1)
+
+    assert find_best_colmap_model(sparse_dir) == sparse_dir / "2"
+    assert get_colmap_model_image_count(sparse_dir / "2") == 4
+
+
+def test_find_best_colmap_model_breaks_exact_ties_by_model_path(tmp_path: Path):
+    sparse_dir = tmp_path / "sparse"
+    _write_binary_model(sparse_dir / "5", ["a.png", "b.png"], 3)
+    _write_binary_model(sparse_dir / "2", ["a.png", "b.png"], 3)
+
+    assert find_best_colmap_model(sparse_dir) == sparse_dir / "2"
+
+
+def test_find_best_colmap_model_skips_invalid_models(tmp_path: Path):
+    sparse_dir = tmp_path / "sparse"
+    invalid_model = sparse_dir / "0"
+    invalid_model.mkdir(parents=True)
+    (invalid_model / "cameras.bin").write_bytes(b"corrupt")
+    _write_binary_model(sparse_dir / "2", ["a.png", "b.png"], 1)
+
+    assert find_best_colmap_model(sparse_dir) == sparse_dir / "2"
+
+
+def test_images_converter_uses_best_direct_sparse_model(tmp_path: Path, monkeypatch):
+    source_dir = tmp_path / "source"
+    output_dir = tmp_path / "output"
+    image_names = _make_source_images(source_dir, 8)
+
+    sparse_dir = output_dir / "colmap" / "sparse"
+    _write_binary_model(sparse_dir / "0", image_names[:3], 2)
+    _write_binary_model(sparse_dir / "1", image_names, 4)
+    _mock_external_tools(tmp_path, monkeypatch)
+
+    converter = ImagesToNerfstudioDataset(data=source_dir, output_dir=output_dir, skip_colmap=True)
+    converter.main()
+
+    assert converter.colmap_model_path == Path("colmap/sparse/1")
+    frames = json.loads((output_dir / "transforms.json").read_text())["frames"]
+    assert len(frames) == len(image_names)
+
+
+def test_images_converter_uses_best_hloc_nested_model(tmp_path: Path, monkeypatch):
+    source_dir = tmp_path / "source"
+    output_dir = tmp_path / "output"
+    image_names = _make_source_images(source_dir, 8)
+
+    sparse_dir = output_dir / "colmap" / "sparse"
+    _write_binary_model(sparse_dir / "0", image_names[:3], 2)
+    _write_binary_model(sparse_dir / "0" / "models" / "0", image_names, 4)
+    _mock_external_tools(tmp_path, monkeypatch)
+
+    converter = ImagesToNerfstudioDataset(data=source_dir, output_dir=output_dir, skip_colmap=True)
+    converter.main()
+
+    assert converter.colmap_model_path == Path("colmap/sparse/0/models/0")
+    frames = json.loads((output_dir / "transforms.json").read_text())["frames"]
+    assert len(frames) == len(image_names)