Keep pycolmap out of the main process; CUDA diagnostics

app.py

@@ -69,6 +69,18 @@ if CUDA_HOME:
         except OSError:
             continue
 
+def _cuda_diag(stage: str) -> None:
+    """Log whether the CUDA driver got loaded into the main (GPU-less) process."""
+    try:
+        maps = open("/proc/self/maps").read()
+        libcuda = "libcuda.so" in maps
+    except Exception:
+        libcuda = "?"
+    print(f"[diag] {stage}: libcuda mapped={libcuda} "
+          f"torch.cuda.is_initialized={torch.cuda.is_initialized()}", flush=True)
+
+
+_cuda_diag("after torch import")
 import torch.utils.cpp_extension as _cppext  # noqa: E402
 
 _cppext._check_cuda_version = lambda *a, **k: None
@@ -80,6 +92,7 @@ from gsplat.cuda._backend import _C as _GSPLAT_C  # noqa: E402
 if _GSPLAT_C is None:
     raise RuntimeError(f"gsplat CUDA backend unavailable (CUDA_HOME={CUDA_HOME}).")
 print(f"[setup] gsplat backend ready in {time.time() - _t0:.1f}s", flush=True)
+_cuda_diag("after gsplat build")
 
 
 # ---------------------------------------------------------------------------
@@ -88,11 +101,12 @@ print(f"[setup] gsplat backend ready in {time.time() - _t0:.1f}s", flush=True)
 import imageio.v2 as imageio  # noqa: E402
 import numpy as np  # noqa: E402
 import gradio as gr  # noqa: E402
-import pycolmap  # noqa: E402
 from huggingface_hub import hf_hub_download  # noqa: E402
 from PIL import Image  # noqa: E402
 
 import codecsplat_pipeline as cp  # noqa: E402
+
+_cuda_diag("after pipeline import")
 from colmap_sfm import IMAGE_EXT, run_sfm_subprocess  # noqa: E402
 
 MODEL_REPO = "pengpeng-yu/CodecSplat"
@@ -107,6 +121,7 @@ print("[setup] building CodecSplat...", flush=True)
 ENCODER = cp.build_encoder(ckpt_path).to("cuda")
 DECODER = cp.build_decoder().to("cuda")
 print("[setup] model ready", flush=True)
+_cuda_diag("after model build")
 
 
 def _u8(t: torch.Tensor) -> np.ndarray:
@@ -173,31 +188,17 @@ def codecsplat_gpu(image_paths: list[str], c2w: np.ndarray, K: np.ndarray, lmb:
 # ---------------------------------------------------------------------------
 # 4. Full pipeline (CPU COLMAP -> GPU CodecSplat -> CPU export)
 # ---------------------------------------------------------------------------
-def _read_colmap_scene(scene_dir: str):
-    rec = pycolmap.Reconstruction(os.path.join(scene_dir, "sparse"))
-    items = []
-    for _, img in rec.images.items():
-        cam = rec.cameras[img.camera_id]
-        p = 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
-        K = np.eye(3, dtype=np.float32)
-        K[0, 0], K[1, 1] = fx / cam.width, fy / cam.height
-        K[0, 2], K[1, 2] = cx / cam.width, cy / cam.height
-        w2c = np.eye(4, dtype=np.float64)
-        w2c[:3, :4] = img.cam_from_world().matrix() if callable(img.cam_from_world) else img.cam_from_world.matrix()
-        path = os.path.join(scene_dir, "images", img.name)
-        if os.path.exists(path):
-            items.append((img.name, path, np.linalg.inv(w2c).astype(np.float32), K))
-    items.sort(key=lambda x: x[0])
-    if len(items) < 2:
+def _scene_arrays(frames: list[dict]):
+    if len(frames) < 2:
         raise gr.Error("COLMAP registered fewer than 2 usable images.")
-    return [i[1] for i in items], np.stack([i[2] for i in items]), np.stack([i[3] for i in items])
+    paths, c2ws, Ks = [], [], []
+    for f in frames:
+        w2c = np.eye(4, dtype=np.float64)
+        w2c[:3, :4] = np.asarray(f["w2c"], dtype=np.float64)
+        paths.append(f["path"])
+        c2ws.append(np.linalg.inv(w2c).astype(np.float32))
+        Ks.append(np.asarray(f["K"], dtype=np.float32))
+    return paths, np.stack(c2ws), np.stack(Ks)
 
 
 def reconstruct(images, lmb: float = 1024.0, matcher: str = "exhaustive",
@@ -235,7 +236,8 @@ def reconstruct(images, lmb: float = 1024.0, matcher: str = "exhaustive",
     except (ValueError, RuntimeError) as exc:
         raise gr.Error(str(exc)) from exc
     t_sfm = time.time() - t_start
-    img_paths, c2w, K = _read_colmap_scene(sfm["scene_dir"])
+    img_paths, c2w, K = _scene_arrays(sfm["frames"])
+    _cuda_diag("before GPU call")
 
     progress(0.55, desc="CodecSplat: encode -> bitstream -> decode -> render")
     r = codecsplat_gpu(img_paths, c2w, K, float(lmb))

colmap_sfm.py

@@ -183,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",