diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py index 55f657a3a4..3a69ceb111 100644 --- a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py +++ b/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, diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 1d9405c81a..e30d176056 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/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) diff --git a/nerfstudio/process_data/images_to_nerfstudio_dataset.py b/nerfstudio/process_data/images_to_nerfstudio_dataset.py index b4d22357c4..1828748383 100644 --- a/nerfstudio/process_data/images_to_nerfstudio_dataset.py +++ b/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 diff --git a/nerfstudio/process_data/video_to_nerfstudio_dataset.py b/nerfstudio/process_data/video_to_nerfstudio_dataset.py index 51a8a0b761..7955e30b39 100644 --- a/nerfstudio/process_data/video_to_nerfstudio_dataset.py +++ b/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 diff --git a/tests/process_data/test_colmap_model_selection.py b/tests/process_data/test_colmap_model_selection.py new file mode 100644 index 0000000000..e62a2693e3 --- /dev/null +++ b/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)