Cap COLMAP threads, single-threaded OpenBLAS

app.py

@@ -31,6 +31,8 @@ sys.path.insert(0, str(APP_DIR))
 os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
 os.environ.setdefault("MAX_JOBS", str(min(32, (os.cpu_count() or 8))))
 os.environ.setdefault("WANDB_MODE", "disabled")
+# pycolmap's bundled OpenBLAS crashes when many matcher threads each spawn BLAS threads.
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
 
 import spaces  # noqa: E402  (must precede any CUDA-touching import)
 import torch  # noqa: E402

colmap_sfm.py

@@ -7,7 +7,7 @@ to a PINHOLE camera model and writes a standard COLMAP scene layout
     <out>/images/*.jpg
     <out>/sparse/{cameras,images,points3D}.bin
 
-which is exactly what ReSplat's reference COLMAP loader expects.
+which is exactly what CodecSplat's reference COLMAP loader expects.
 """
 
 from __future__ import annotations
@@ -18,6 +18,11 @@ from pathlib import Path
 import pycolmap
 from PIL import Image, ImageOps
 
+# ZeroGPU hosts expose ~192 CPUs. With COLMAP's default num_threads=-1 the incremental
+# mapper then fails to register anything beyond the initial pair (reproduced locally by
+# forcing 192 threads), so cap the thread count explicitly.
+NUM_THREADS = 16
+
 IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff")
 
 
@@ -75,6 +80,12 @@ def run_sfm(
             "Upload several overlapping photos of the same scene."
         )
     log(f"[colmap] {len(names)} images prepared at <= {max_image_size}px")
+    try:
+        import os as _os
+        cpu = next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), "?")
+        log(f"[colmap] cpu={cpu} ncpu={_os.cpu_count()} threads={NUM_THREADS}")
+    except Exception:
+        pass
 
     database_path = workdir / "colmap.db"
     if database_path.exists():
@@ -86,6 +97,7 @@ def run_sfm(
     extraction_options.sift.domain_size_pooling = False
     extraction_options.use_gpu = False
     extraction_options.max_image_size = int(max_image_size)
+    extraction_options.num_threads = NUM_THREADS
 
     reader_options = pycolmap.ImageReaderOptions()
     reader_options.camera_model = "SIMPLE_RADIAL"
@@ -100,23 +112,28 @@ def run_sfm(
     )
     log("[colmap] features extracted")
 
+    matching_options = pycolmap.FeatureMatchingOptions()
+    matching_options.num_threads = NUM_THREADS
     if matcher == "sequential":
         pairing = pycolmap.SequentialPairingOptions()
         pairing.overlap = 10
         pairing.quadratic_overlap = True
         pycolmap.match_sequential(
             database_path=database_path,
+            matching_options=matching_options,
             pairing_options=pairing,
             device=pycolmap.Device.cpu,
         )
     else:
-        pycolmap.match_exhaustive(database_path=database_path, device=pycolmap.Device.cpu)
+        pycolmap.match_exhaustive(database_path=database_path, matching_options=matching_options,
+                                  device=pycolmap.Device.cpu)
     log(f"[colmap] {matcher} matching done")
 
     sparse_dir = workdir / "sparse_raw"
     sparse_dir.mkdir(parents=True, exist_ok=True)
     options = pycolmap.IncrementalPipelineOptions()
     options.min_num_matches = 15
+    options.num_threads = NUM_THREADS
     options.multiple_models = False
     options.extract_colors = True
     maps = pycolmap.incremental_mapping(
@@ -134,7 +151,7 @@ def run_sfm(
             "rotation / blurry frames, and keep the scene static."
         )
     # Deterministic world scale (same convention COLMAP itself uses), so that the
-    # fixed near/far planes of the ReSplat checkpoint stay meaningful.
+    # fixed near/far planes of the CodecSplat checkpoint stay meaningful.
     rec.normalize(fixed_scale=False, extent=10.0)
     log(f"[colmap] mapped {rec.num_reg_images()}/{len(names)} images, {rec.num_points3D()} points")
 
@@ -154,7 +171,7 @@ def run_sfm(
     )
     log("[colmap] undistorted to PINHOLE")
 
-    # Keep only what ReSplat reads (the undistorter also writes stereo/ scaffolding).
+    # Keep only what CodecSplat reads (the undistorter also writes stereo/ scaffolding).
     for extra in ("stereo",):
         p = scene_dir / extra
         if p.is_dir():
@@ -166,7 +183,32 @@ def run_sfm(
     except Exception:
         pass
 
+    # Read back the undistorted PINHOLE cameras (normalized K, OpenCV c2w), sorted by name.
+    und = pycolmap.Reconstruction(str(scene_dir / "sparse"))
+    frames = []
+    for _, img in und.images.items():
+        cam = und.cameras[img.camera_id]
+        p = list(cam.params)
+        if cam.model.name == "PINHOLE":
+            fx, fy, cx, cy = p[:4]
+        elif cam.model.name == "SIMPLE_PINHOLE":
+            fx = fy = p[0]
+            cx, cy = p[1:3]
+        else:
+            continue
+        cfw = img.cam_from_world() if callable(img.cam_from_world) else img.cam_from_world
+        w2c = cfw.matrix()  # 3x4
+        path = scene_dir / "images" / img.name
+        if path.exists():
+            frames.append({
+                "name": img.name, "path": str(path),
+                "K": [[fx / cam.width, 0, cx / cam.width], [0, fy / cam.height, cy / cam.height], [0, 0, 1]],
+                "w2c": [list(map(float, row)) for row in w2c],
+            })
+    frames.sort(key=lambda f: f["name"])
+
     return {
+        "frames": frames,
         "scene_dir": str(scene_dir),
         "sparse_sub": "sparse",
         "images_sub": "images",
@@ -175,3 +217,55 @@ def run_sfm(
         "num_points3D": int(rec.num_points3D()),
         "camera_model": cam_model,
     }
+
+
+def run_sfm_subprocess(image_paths, workdir, max_image_size=1280, matcher="exhaustive",
+                       timeout=900, log=print) -> dict:
+    """Run `run_sfm` in a fresh Python process that only imports pycolmap.
+
+    Inside the Gradio process torch / CUDA / OpenMP / BLAS libraries are already loaded,
+    which makes pycolmap's bundled OpenBLAS / OpenMP misbehave (crashes in matching and
+    failed image registration). A clean child process avoids all of that.
+    """
+    import json
+    import os
+    import subprocess
+    import sys
+
+    workdir = Path(workdir)
+    workdir.mkdir(parents=True, exist_ok=True)
+    spec = workdir / "sfm_job.json"
+    result = workdir / "sfm_result.json"
+    spec.write_text(json.dumps({
+        "image_paths": [str(p) for p in image_paths], "workdir": str(workdir),
+        "max_image_size": int(max_image_size), "matcher": matcher, "result": str(result),
+    }))
+    env = {k: v for k, v in os.environ.items() if k not in ("LD_PRELOAD",)}
+    env["LD_LIBRARY_PATH"] = os.environ.get("_ORIG_LD_LIBRARY_PATH", "")
+    env["OPENBLAS_NUM_THREADS"] = "1"
+    env.pop("CUDA_VISIBLE_DEVICES", None)
+    proc = subprocess.run([sys.executable, os.path.abspath(__file__), str(spec)], env=env,
+                          capture_output=True, text=True, timeout=timeout)
+    for line in (proc.stdout or "").splitlines():
+        if line.startswith("[colmap]"):
+            log(line)
+    if not result.exists():
+        tail = "\n".join(((proc.stderr or "") + (proc.stdout or "")).splitlines()[-15:])
+        raise RuntimeError(f"COLMAP subprocess failed (exit {proc.returncode}):\n{tail}")
+    data = json.loads(result.read_text())
+    if "error" in data:
+        raise ValueError(data["error"])
+    return data
+
+
+if __name__ == "__main__":
+    import json
+    import sys
+
+    job = json.loads(Path(sys.argv[1]).read_text())
+    try:
+        out = run_sfm(job["image_paths"], Path(job["workdir"]), job["max_image_size"],
+                      job["matcher"], log=lambda m: print(m, flush=True))
+    except ValueError as exc:
+        out = {"error": str(exc)}
+    Path(job["result"]).write_text(json.dumps(out))