Spaces:
Running on Zero
Running on Zero
| """COLMAP structure-from-motion preprocessing (runs inside the Space, on CPU). | |
| Takes a set of raw multi-view images of a scene, estimates camera intrinsics / | |
| poses / a sparse point cloud with COLMAP (via `pycolmap`), undistorts the images | |
| 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 CodecSplat's reference COLMAP loader expects. | |
| """ | |
| from __future__ import annotations | |
| import shutil | |
| 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") | |
| def prepare_images(image_paths, dest: Path, max_image_size: int = 1280) -> list[str]: | |
| """Copy the uploaded images into `dest`, EXIF-rotated, RGB, size-capped.""" | |
| dest.mkdir(parents=True, exist_ok=True) | |
| names: list[str] = [] | |
| for i, src in enumerate(str(p) for p in image_paths): # caller decides the order | |
| if not src.lower().endswith(IMAGE_EXT): | |
| continue | |
| try: | |
| img = Image.open(src) | |
| img = ImageOps.exif_transpose(img).convert("RGB") | |
| except Exception as exc: | |
| print(f"[colmap] skipping unreadable image {src}: {exc}", flush=True) | |
| continue | |
| scale = max_image_size / max(img.size) | |
| if scale < 1.0: | |
| img = img.resize( | |
| (max(1, int(round(img.width * scale))), max(1, int(round(img.height * scale)))), | |
| Image.LANCZOS, | |
| ) | |
| name = f"{i:04d}.jpg" | |
| img.save(dest / name, quality=95) | |
| names.append(name) | |
| return names | |
| def _largest_reconstruction(maps) -> pycolmap.Reconstruction | None: | |
| recs = list(maps.values()) if isinstance(maps, dict) else list(maps) | |
| if not recs: | |
| return None | |
| return max(recs, key=lambda r: r.num_reg_images()) | |
| def run_sfm( | |
| image_paths, | |
| workdir: Path, | |
| max_image_size: int = 1280, | |
| matcher: str = "exhaustive", | |
| max_num_features: int = 8192, | |
| log=print, | |
| ) -> dict: | |
| """Run COLMAP SfM on `image_paths`; return info about the written scene. | |
| Returns a dict with `scene_dir`, `sparse_sub`, `images_sub`, `num_input`, | |
| `num_registered`, `num_points3D`, `camera_model`. | |
| """ | |
| workdir = Path(workdir) | |
| image_dir = workdir / "sfm_images" | |
| names = prepare_images(image_paths, image_dir, max_image_size) | |
| if len(names) < 3: | |
| raise ValueError( | |
| f"COLMAP needs at least 3 usable images, got {len(names)}. " | |
| "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(): | |
| database_path.unlink() | |
| extraction_options = pycolmap.FeatureExtractionOptions() | |
| extraction_options.sift.max_num_features = int(max_num_features) | |
| extraction_options.sift.estimate_affine_shape = False | |
| 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" | |
| pycolmap.extract_features( | |
| database_path=database_path, | |
| image_path=image_dir, | |
| camera_mode=pycolmap.CameraMode.SINGLE, | |
| reader_options=reader_options, | |
| extraction_options=extraction_options, | |
| device=pycolmap.Device.cpu, | |
| ) | |
| 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, 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( | |
| database_path=database_path, | |
| image_path=image_dir, | |
| output_path=sparse_dir, | |
| options=options, | |
| ) | |
| rec = _largest_reconstruction(maps) | |
| if rec is None or rec.num_reg_images() < 3: | |
| raise ValueError( | |
| "COLMAP could not reconstruct this image set (only " | |
| f"{0 if rec is None else rec.num_reg_images()} of {len(names)} images registered). " | |
| "Use photos with plenty of visual overlap between neighbours, avoid pure " | |
| "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 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") | |
| model_dir = workdir / "sparse_model" | |
| if model_dir.exists(): | |
| shutil.rmtree(model_dir) | |
| model_dir.mkdir(parents=True) | |
| rec.write(str(model_dir)) | |
| scene_dir = workdir / "colmap_scene" | |
| if scene_dir.exists(): | |
| shutil.rmtree(scene_dir) | |
| pycolmap.undistort_images( | |
| output_path=scene_dir, | |
| input_path=model_dir, | |
| image_path=image_dir, | |
| ) | |
| log("[colmap] undistorted to PINHOLE") | |
| # Keep only what CodecSplat reads (the undistorter also writes stereo/ scaffolding). | |
| for extra in ("stereo",): | |
| p = scene_dir / extra | |
| if p.is_dir(): | |
| shutil.rmtree(p, ignore_errors=True) | |
| cam_model = "PINHOLE" | |
| try: | |
| cam_model = list(rec.cameras.values())[0].model.name | |
| 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", | |
| "num_input": len(names), | |
| "num_registered": int(rec.num_reg_images()), | |
| "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)) | |