source/colmap_worker.py

"""Isolated PyCOLMAP worker for EDGS preprocessing.

PyCOLMAP can terminate the Python process if a native dependency segfaults.
Running COLMAP in this child process protects the Gradio backend process and
converts a native crash into a normal Python RuntimeError in the parent.
"""
import os

_THREAD_DEFAULTS = {
    "OPENBLAS_NUM_THREADS": "1",
    "OMP_NUM_THREADS": "1",
    "MKL_NUM_THREADS": "1",
    "NUMEXPR_NUM_THREADS": "1",
    "VECLIB_MAXIMUM_THREADS": "1",
    "BLIS_NUM_THREADS": "1",
    "OPENCV_OPENCL_RUNTIME": "disabled",
    "OPENCV_OPENCL_DEVICE": "disabled",
}
for _name, _value in _THREAD_DEFAULTS.items():
    os.environ.setdefault(_name, _value)

import argparse
import time
from pathlib import Path

import pycolmap


def _get_int_env(name: str, default: int, minimum: int = 1) -> int:
    try:
        return max(minimum, int(os.getenv(name, str(default))))
    except Exception:
        return default


def _set_if_present(obj, name: str, value):
    if hasattr(obj, name):
        try:
            setattr(obj, name, value)
            return True
        except Exception as exc:
            print(f"[EDGS][COLMAP] Could not set {name}={value!r}: {exc!r}", flush=True)
    return False


def _make_extraction_options(threads: int, max_image_size: int, max_num_features: int):
    options = pycolmap.FeatureExtractionOptions()
    _set_if_present(options, "num_threads", threads)
    _set_if_present(options, "use_gpu", False)
    _set_if_present(options, "max_image_size", max_image_size)
    if hasattr(options, "sift"):
        _set_if_present(options.sift, "max_num_features", max_num_features)
        _set_if_present(options.sift, "first_octave", 0)
    return options


def _make_matching_options(threads: int, max_num_matches: int):
    options = pycolmap.FeatureMatchingOptions()
    _set_if_present(options, "num_threads", threads)
    _set_if_present(options, "use_gpu", False)
    _set_if_present(options, "max_num_matches", max_num_matches)
    if hasattr(options, "sift"):
        # Avoid the FAISS CPU matcher path that triggered OpenBLAS crashes in logs.
        _set_if_present(options.sift, "cpu_brute_force_matcher", True)
    return options


def _make_mapping_options(threads: int):
    options = pycolmap.IncrementalPipelineOptions()
    options.min_num_matches = 15
    options.multiple_models = True
    options.max_num_models = 50
    options.max_model_overlap = 20
    options.min_model_size = 10
    options.extract_colors = True
    options.num_threads = threads
    options.mapper.init_min_num_inliers = 30
    options.mapper.init_max_error = 8.0
    options.mapper.init_min_tri_angle = 5.0
    return options


def run_colmap_on_scene(scene_dir: str) -> None:
    start_time = time.time()
    scene_path = Path(scene_dir)
    database_path = scene_path / "database.db"
    sparse_path = scene_path / "sparse"
    image_dir = scene_path / "images"
    sparse_path.mkdir(parents=True, exist_ok=True)

    threads = _get_int_env("EDGS_COLMAP_THREADS", 2)
    max_image_size = _get_int_env("EDGS_COLMAP_MAX_IMAGE_SIZE", 1024)
    max_num_features = _get_int_env("EDGS_COLMAP_MAX_NUM_FEATURES", 4096)
    max_num_matches = _get_int_env("EDGS_COLMAP_MAX_NUM_MATCHES", 16384)

    print(
        "[EDGS][COLMAP] isolated worker settings: "
        f"threads={threads}, max_image_size={max_image_size}, "
        f"max_num_features={max_num_features}, max_num_matches={max_num_matches}, "
        f"OPENBLAS_NUM_THREADS={os.getenv('OPENBLAS_NUM_THREADS')}, "
        f"OMP_NUM_THREADS={os.getenv('OMP_NUM_THREADS')}",
        flush=True,
    )

    extraction_options = _make_extraction_options(threads, max_image_size, max_num_features)
    matching_options = _make_matching_options(threads, max_num_matches)

    print(f"[EDGS][COLMAP] extracting features in {image_dir}", flush=True)
    pycolmap.extract_features(
        database_path=database_path,
        image_path=image_dir,
        extraction_options=extraction_options,
    )
    print(f"[EDGS][COLMAP] feature extraction done in {(time.time() - start_time):.2f}s", flush=True)

    print("[EDGS][COLMAP] matching features", flush=True)
    pycolmap.match_exhaustive(
        database_path=database_path,
        matching_options=matching_options,
    )
    print(f"[EDGS][COLMAP] feature matching done in {(time.time() - start_time):.2f}s", flush=True)

    print("[EDGS][COLMAP] incremental mapping", flush=True)
    pycolmap.incremental_mapping(
        database_path=database_path,
        image_path=image_dir,
        output_path=sparse_path,
        options=_make_mapping_options(threads),
    )
    print(f"[EDGS][COLMAP] incremental mapping done in {(time.time() - start_time):.2f}s", flush=True)

    recon_path = sparse_path / "0"
    if not recon_path.exists():
        raise RuntimeError(
            "COLMAP did not produce sparse/0. Try fewer reference views, a smaller video crop, "
            "or higher-overlap footage."
        )

    reconstruction = pycolmap.Reconstruction(recon_path)
    for cam in reconstruction.cameras.values():
        cam.model = "SIMPLE_PINHOLE"
        cam.params = cam.params[:3]
    reconstruction.write(recon_path)
    print(f"[EDGS][COLMAP] total worker time: {(time.time() - start_time):.2f}s", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Run EDGS PyCOLMAP preprocessing in an isolated child process.")
    parser.add_argument("scene_dir")
    args = parser.parse_args()
    run_colmap_on_scene(args.scene_dir)


if __name__ == "__main__":
    main()