#!/usr/bin/env python3 """ EndoGaussian-4D Camera Pose Extraction Pipeline Multi-stage SfM pipeline for extracting camera poses from endoscopic video, addressing the "unposed video" challenge on texture-less surgical tissue. Pipeline Stages (tried in order, falls back on failure): Stage 1: COLMAP Sequential Matcher (tuned for endoscopy) Stage 2: COLMAP Exhaustive Matcher (slower, more robust) Stage 3: Depth-Anything + PnP-RANSAC (learning-based fallback) The pipeline also implements Holistic Gaussian Initialization (HGI): P = ∪_t K⁻¹ · T_t · D_t · (I_t ⊙ M_t) Usage: # Auto mode: tries COLMAP first, falls back to Depth+PnP python scripts/extract_poses.py --input ./data/endonerf/cutting --mode auto # Force specific method python scripts/extract_poses.py --input ./data/endonerf/cutting --mode colmap_sequential python scripts/extract_poses.py --input ./data/endonerf/cutting --mode depth_pnp # Run HGI after pose extraction python scripts/extract_poses.py --input ./data/endonerf/cutting --hgi --subsample 0.001 """ import argparse import json import os import shutil import subprocess import tempfile from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np # --------------------------------------------------------------------------- # COLMAP Configuration (tuned for endoscopy) # --------------------------------------------------------------------------- COLMAP_FEATURE_CONFIG = { # Lower peak threshold to detect more features on smooth tissue "SiftExtraction.peak_threshold": "0.004", # More octaves for multi-scale matching "SiftExtraction.num_octaves": "4", # Max features per image (increase for detail-poor endoscopy) "SiftExtraction.max_num_features": "8192", # Enable GPU if available "SiftExtraction.use_gpu": "1", } COLMAP_SEQUENTIAL_CONFIG = { # Overlap window for sequential matching (endoscopy = smooth camera motion) "SiftMatching.guided_matching": "1", "SequentialMatching.overlap": "15", "SequentialMatching.loop_detection": "1", } COLMAP_MAPPER_CONFIG = { # Lower triangulation angle for close-range endoscopy "Mapper.init_min_tri_angle": "2.0", "Mapper.multiple_models": "0", # More permissive registration for texture-poor scenes "Mapper.abs_pose_min_num_inliers": "10", "Mapper.ba_global_max_num_iterations": "50", } # --------------------------------------------------------------------------- # COLMAP Pipeline # --------------------------------------------------------------------------- class COLMAPRunner: """Runs COLMAP SfM pipeline with endoscopy-tuned parameters.""" def __init__(self, image_dir: str, work_dir: str, use_gpu: bool = True): self.image_dir = Path(image_dir) self.work_dir = Path(work_dir) self.work_dir.mkdir(parents=True, exist_ok=True) self.db_path = self.work_dir / "database.db" self.sparse_dir = self.work_dir / "sparse" self.use_gpu = use_gpu def _check_colmap(self) -> bool: """Check if COLMAP is installed.""" try: result = subprocess.run(["colmap", "--help"], capture_output=True, timeout=10) return result.returncode == 0 except (FileNotFoundError, subprocess.TimeoutExpired): return False def _run_cmd(self, args: List[str], desc: str = "") -> bool: """Run a COLMAP command.""" print(f" [COLMAP] {desc}...") try: result = subprocess.run( args, capture_output=True, text=True, timeout=600 ) if result.returncode != 0: print(f" [COLMAP] {desc} FAILED: {result.stderr[:500]}") return False return True except subprocess.TimeoutExpired: print(f" [COLMAP] {desc} TIMEOUT") return False def extract_features(self) -> bool: """Extract SIFT features tuned for endoscopy.""" args = [ "colmap", "feature_extractor", "--database_path", str(self.db_path), "--image_path", str(self.image_dir), ] for k, v in COLMAP_FEATURE_CONFIG.items(): if k == "SiftExtraction.use_gpu" and not self.use_gpu: args.extend([f"--{k}", "0"]) else: args.extend([f"--{k}", v]) return self._run_cmd(args, "Feature extraction") def match_sequential(self) -> bool: """Sequential matching (exploits temporal continuity).""" args = [ "colmap", "sequential_matcher", "--database_path", str(self.db_path), ] for k, v in COLMAP_SEQUENTIAL_CONFIG.items(): args.extend([f"--{k}", v]) return self._run_cmd(args, "Sequential matching") def match_exhaustive(self) -> bool: """Exhaustive matching (slower but more robust).""" args = [ "colmap", "exhaustive_matcher", "--database_path", str(self.db_path), ] return self._run_cmd(args, "Exhaustive matching") def reconstruct(self) -> bool: """Run incremental SfM mapper.""" self.sparse_dir.mkdir(parents=True, exist_ok=True) args = [ "colmap", "mapper", "--database_path", str(self.db_path), "--image_path", str(self.image_dir), "--output_path", str(self.sparse_dir), ] for k, v in COLMAP_MAPPER_CONFIG.items(): args.extend([f"--{k}", v]) return self._run_cmd(args, "Incremental SfM") def get_registration_rate(self) -> float: """Check what fraction of images were registered.""" model_dir = self.sparse_dir / "0" if not model_dir.exists(): return 0.0 try: # Read images.txt to count registered images images_txt = model_dir / "images.txt" if images_txt.exists(): with open(images_txt) as f: lines = [l for l in f.readlines() if l.strip() and not l.startswith("#")] # Every other line is an image entry n_registered = len(lines) // 2 else: # Try binary format images_bin = model_dir / "images.bin" if images_bin.exists(): # Approximate: count by file size n_registered = max(1, os.path.getsize(images_bin) // 200) else: return 0.0 n_total = len(list(self.image_dir.glob("*.png"))) + \ len(list(self.image_dir.glob("*.jpg"))) return n_registered / max(n_total, 1) except Exception: return 0.0 def extract_poses(self) -> Optional[Dict]: """Extract poses from COLMAP reconstruction.""" model_dir = self.sparse_dir / "0" if not model_dir.exists(): return None try: # Try using pycolmap for clean extraction import pycolmap reconstruction = pycolmap.Reconstruction(str(model_dir)) poses = {} intrinsics = None for img_id, image in reconstruction.images.items(): cam = reconstruction.cameras[image.camera_id] # Camera-to-world transform R = image.cam_from_world.rotation.matrix() t = image.cam_from_world.translation # World-to-camera w2c = np.eye(4, dtype=np.float64) w2c[:3, :3] = R w2c[:3, 3] = t # Camera-to-world c2w = np.linalg.inv(w2c) poses[image.name] = c2w.astype(np.float32) if intrinsics is None: params = cam.params if cam.model_name in ("SIMPLE_PINHOLE", "SIMPLE_RADIAL"): fx = fy = params[0] cx, cy = params[1], params[2] elif cam.model_name in ("PINHOLE", "RADIAL"): fx, fy = params[0], params[1] cx, cy = params[2], params[3] else: fx = fy = params[0] cx, cy = cam.width / 2, cam.height / 2 intrinsics = np.array([ [fx, 0, cx], [0, fy, cy], [0, 0, 1] ], dtype=np.float32) return {"poses": poses, "intrinsics": intrinsics} except ImportError: print(" [COLMAP] pycolmap not available, reading text format...") return self._parse_colmap_text(model_dir) def _parse_colmap_text(self, model_dir: Path) -> Optional[Dict]: """Parse COLMAP text-format output.""" images_txt = model_dir / "images.txt" cameras_txt = model_dir / "cameras.txt" if not images_txt.exists(): # Convert binary to text self._run_cmd([ "colmap", "model_converter", "--input_path", str(model_dir), "--output_path", str(model_dir), "--output_type", "TXT", ], "Convert to text") if not images_txt.exists(): return None poses = {} with open(images_txt) as f: lines = [l.strip() for l in f.readlines() if l.strip() and not l.startswith("#")] for i in range(0, len(lines), 2): parts = lines[i].split() # IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME if len(parts) < 10: continue qw, qx, qy, qz = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4]) tx, ty, tz = float(parts[5]), float(parts[6]), float(parts[7]) name = parts[9] # Quaternion to rotation matrix R = _quat_to_rotation_matrix(qw, qx, qy, qz) w2c = np.eye(4, dtype=np.float32) w2c[:3, :3] = R w2c[:3, 3] = [tx, ty, tz] c2w = np.linalg.inv(w2c) poses[name] = c2w intrinsics = None if cameras_txt.exists(): with open(cameras_txt) as f: for line in f: if line.startswith("#"): continue parts = line.strip().split() if len(parts) >= 5: model = parts[1] params = [float(p) for p in parts[4:]] if model in ("SIMPLE_PINHOLE", "SIMPLE_RADIAL"): fx = fy = params[0] cx, cy = params[1], params[2] elif model in ("PINHOLE",): fx, fy = params[0], params[1] cx, cy = params[2], params[3] else: fx = fy = params[0] cx = float(parts[2]) / 2 cy = float(parts[3]) / 2 intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32) break return {"poses": poses, "intrinsics": intrinsics} def _quat_to_rotation_matrix(qw, qx, qy, qz): """Convert quaternion to 3x3 rotation matrix.""" R = np.array([ [1 - 2*(qy*qy + qz*qz), 2*(qx*qy - qz*qw), 2*(qx*qz + qy*qw)], [2*(qx*qy + qz*qw), 1 - 2*(qx*qx + qz*qz), 2*(qy*qz - qx*qw)], [2*(qx*qz - qy*qw), 2*(qy*qz + qx*qw), 1 - 2*(qx*qx + qy*qy)], ], dtype=np.float32) return R # --------------------------------------------------------------------------- # Depth-Anything + PnP-RANSAC Fallback # --------------------------------------------------------------------------- class DepthPnPPipeline: """ Learning-based pose estimation for when COLMAP fails on texture-less tissue. Pipeline: 1. Estimate monocular depth for all frames using Depth-Anything-Small 2. Extract ORB features and match between consecutive frames 3. Use PnP-RANSAC with depth to estimate relative poses 4. Chain relative poses into a global trajectory This handles the fundamental challenge of endoscopy: smooth, specular, texture-less tissue surfaces that defeat traditional SfM. """ def __init__(self, device: str = "cuda"): self.device = device self._depth_model = None self._depth_processor = None def _load_depth_model(self): """Lazy-load Depth-Anything-Small from HuggingFace.""" if self._depth_model is not None: return print(" [Depth] Loading Depth-Anything-V2-Small...") try: from transformers import AutoImageProcessor, AutoModelForDepthEstimation import torch model_id = "depth-anything/Depth-Anything-V2-Small-hf" self._depth_processor = AutoImageProcessor.from_pretrained(model_id) self._depth_model = AutoModelForDepthEstimation.from_pretrained(model_id) self._depth_model.to(self.device) self._depth_model.eval() print(" [Depth] Model loaded ✓") except Exception as e: print(f" [Depth] Failed to load model: {e}") raise def estimate_depth(self, image: np.ndarray) -> np.ndarray: """ Estimate monocular depth for a single image. Args: image: [H, W, 3] uint8 RGB image Returns: [H, W] float32 relative depth map (larger = farther) """ import torch from PIL import Image self._load_depth_model() pil_image = Image.fromarray(image) inputs = self._depth_processor(images=pil_image, return_tensors="pt") inputs = {k: v.to(self.device) for k, v in inputs.items()} with torch.no_grad(): outputs = self._depth_model(**inputs) predicted_depth = outputs.predicted_depth # Interpolate to original size depth = torch.nn.functional.interpolate( predicted_depth.unsqueeze(1), size=image.shape[:2], mode="bicubic", align_corners=False, ).squeeze().cpu().numpy() return depth.astype(np.float32) def extract_poses( self, image_dir: str, intrinsics: Optional[np.ndarray] = None, ) -> Dict: """ Extract poses using Depth + ORB + PnP-RANSAC. Args: image_dir: Directory with image files intrinsics: [3, 3] camera matrix (estimated if not provided) Returns: Dict with "poses" (name → [4,4]) and "intrinsics" ([3,3]) """ import cv2 from PIL import Image img_dir = Path(image_dir) image_paths = sorted( list(img_dir.glob("*.png")) + list(img_dir.glob("*.jpg")) ) if not image_paths: raise FileNotFoundError(f"No images in {image_dir}") n_images = len(image_paths) print(f" [DepthPnP] Processing {n_images} images...") # Load first image to get dimensions first_img = np.array(Image.open(image_paths[0]).convert("RGB")) H, W = first_img.shape[:2] # Estimate intrinsics if not provided if intrinsics is None: f = max(H, W) * 1.2 # Rough focal length estimate intrinsics = np.array([ [f, 0, W / 2], [0, f, H / 2], [0, 0, 1] ], dtype=np.float32) # Initialize ORB detector orb = cv2.ORB_create(nfeatures=2000) bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True) # Process frame pairs poses = {} cumulative_pose = np.eye(4, dtype=np.float32) poses[image_paths[0].name] = cumulative_pose.copy() prev_img = cv2.cvtColor(first_img, cv2.COLOR_RGB2GRAY) prev_depth = self.estimate_depth(first_img) prev_kp, prev_desc = orb.detectAndCompute(prev_img, None) n_success = 0 for i in range(1, n_images): curr_rgb = np.array(Image.open(image_paths[i]).convert("RGB")) curr_gray = cv2.cvtColor(curr_rgb, cv2.COLOR_RGB2GRAY) # Features curr_kp, curr_desc = orb.detectAndCompute(curr_gray, None) if prev_desc is None or curr_desc is None or len(prev_kp) < 10 or len(curr_kp) < 10: poses[image_paths[i].name] = cumulative_pose.copy() prev_img = curr_gray prev_kp, prev_desc = curr_kp, curr_desc continue # Match matches = bf.match(prev_desc, curr_desc) matches = sorted(matches, key=lambda m: m.distance)[:500] if len(matches) < 8: poses[image_paths[i].name] = cumulative_pose.copy() prev_img = curr_gray prev_kp, prev_desc = curr_kp, curr_desc continue # Get 3D-2D correspondences using depth obj_points = [] img_points = [] for m in matches: pt_prev = prev_kp[m.queryIdx].pt pt_curr = curr_kp[m.trainIdx].pt u, v = int(round(pt_prev[0])), int(round(pt_prev[1])) if 0 <= v < H and 0 <= u < W: d = prev_depth[v, u] if d > 1e-3: # Backproject to 3D x = (u - intrinsics[0, 2]) * d / intrinsics[0, 0] y = (v - intrinsics[1, 2]) * d / intrinsics[1, 1] z = d obj_points.append([x, y, z]) img_points.append([pt_curr[0], pt_curr[1]]) if len(obj_points) < 6: poses[image_paths[i].name] = cumulative_pose.copy() prev_img = curr_gray prev_depth = self.estimate_depth(curr_rgb) prev_kp, prev_desc = curr_kp, curr_desc continue obj_points = np.array(obj_points, dtype=np.float32) img_points = np.array(img_points, dtype=np.float32) # PnP-RANSAC success, rvec, tvec, inliers = cv2.solvePnPRansac( obj_points, img_points, intrinsics, None, iterationsCount=1000, reprojectionError=5.0, flags=cv2.SOLVEPNP_ITERATIVE, ) if success and inliers is not None and len(inliers) >= 6: R, _ = cv2.Rodrigues(rvec) rel_pose = np.eye(4, dtype=np.float32) rel_pose[:3, :3] = R rel_pose[:3, 3] = tvec.squeeze() cumulative_pose = cumulative_pose @ np.linalg.inv(rel_pose) n_success += 1 poses[image_paths[i].name] = cumulative_pose.copy() # Update previous frame prev_img = curr_gray prev_depth = self.estimate_depth(curr_rgb) prev_kp, prev_desc = curr_kp, curr_desc if (i + 1) % 20 == 0: print(f" Frame {i+1}/{n_images} | PnP success: {n_success}/{i}") rate = n_success / max(n_images - 1, 1) print(f" [DepthPnP] Registration rate: {rate:.1%} ({n_success}/{n_images-1})") return {"poses": poses, "intrinsics": intrinsics} # --------------------------------------------------------------------------- # Multi-Stage SfM Pipeline # --------------------------------------------------------------------------- class EndoSfMPipeline: """ Multi-stage pose extraction pipeline for endoscopic video. Automatically tries methods in order of accuracy: 1. COLMAP sequential (fastest, works on textured regions) 2. COLMAP exhaustive (slower, catches more matches) 3. Depth-Anything + PnP (learning-based, handles texture-less) Falls back to next stage if registration rate < threshold. """ REGISTRATION_THRESHOLD = 0.7 # 70% of images must be registered def __init__(self, input_dir: str, output_dir: Optional[str] = None): self.input_dir = Path(input_dir) self.output_dir = Path(output_dir) if output_dir else self.input_dir # Find image directory self.image_dir = self._find_image_dir() def _find_image_dir(self) -> Path: """Find the image directory within the input.""" for name in ["images", "color", "Frames", "rgb"]: d = self.input_dir / name if d.is_dir(): return d # Check if input dir itself contains images if list(self.input_dir.glob("*.png")) or list(self.input_dir.glob("*.jpg")): return self.input_dir raise FileNotFoundError(f"No image directory found in {self.input_dir}") def run(self, mode: str = "auto") -> Dict: """ Run pose extraction. Args: mode: "auto", "colmap_sequential", "colmap_exhaustive", "depth_pnp" Returns: Dict with "poses", "intrinsics", "method" """ print(f"\n{'='*60}") print(f"EndoGaussian-4D Pose Extraction") print(f"Input: {self.input_dir}") print(f"Images: {self.image_dir}") print(f"Mode: {mode}") print(f"{'='*60}\n") if mode == "auto": return self._run_auto() elif mode == "colmap_sequential": return self._run_colmap("sequential") elif mode == "colmap_exhaustive": return self._run_colmap("exhaustive") elif mode == "depth_pnp": return self._run_depth_pnp() else: raise ValueError(f"Unknown mode: {mode}") def _run_auto(self) -> Dict: """Auto mode: try methods in order.""" # Stage 1: COLMAP sequential print("[Stage 1/3] COLMAP Sequential Matcher") result = self._run_colmap("sequential") if result and result.get("registration_rate", 0) >= self.REGISTRATION_THRESHOLD: result["method"] = "colmap_sequential" self._save_result(result) return result print(f" Registration rate too low, trying next stage...\n") # Stage 2: COLMAP exhaustive print("[Stage 2/3] COLMAP Exhaustive Matcher") result = self._run_colmap("exhaustive") if result and result.get("registration_rate", 0) >= self.REGISTRATION_THRESHOLD: result["method"] = "colmap_exhaustive" self._save_result(result) return result print(f" Registration rate too low, trying next stage...\n") # Stage 3: Depth + PnP print("[Stage 3/3] Depth-Anything + PnP-RANSAC") result = self._run_depth_pnp() result["method"] = "depth_pnp" self._save_result(result) return result def _run_colmap(self, matching: str) -> Optional[Dict]: """Run COLMAP pipeline.""" work_dir = self.output_dir / f"colmap_{matching}" runner = COLMAPRunner(str(self.image_dir), str(work_dir)) if not runner._check_colmap(): print(" [COLMAP] Not installed, skipping") return None if not runner.extract_features(): return None if matching == "sequential": if not runner.match_sequential(): return None else: if not runner.match_exhaustive(): return None if not runner.reconstruct(): return None rate = runner.get_registration_rate() print(f" Registration rate: {rate:.1%}") result = runner.extract_poses() if result: result["registration_rate"] = rate return result def _run_depth_pnp(self) -> Dict: """Run Depth-Anything + PnP pipeline.""" pipeline = DepthPnPPipeline() return pipeline.extract_poses(str(self.image_dir)) def _save_result(self, result: Dict): """Save poses in LLFF format + JSON metadata.""" poses = result.get("poses", {}) intrinsics = result.get("intrinsics") if not poses: print(" [Save] No poses to save") return # Sort by filename sorted_names = sorted(poses.keys()) n = len(sorted_names) # Build LLFF poses_bounds.npy: [N, 17] = [3x5 pose | near, far] if intrinsics is not None: H, W = 480, 640 # Default; should be read from images # Try to get actual dimensions for name in sorted_names: img_path = self.image_dir / name if img_path.exists(): from PIL import Image img = Image.open(img_path) W, H = img.size break f = intrinsics[0, 0] poses_bounds = np.zeros((n, 17), dtype=np.float64) for i, name in enumerate(sorted_names): c2w = poses[name] # LLFF format: [R|t|hwf] hwf = np.array([H, W, f], dtype=np.float64) pose_3x5 = np.concatenate([c2w[:3, :4], hwf.reshape(3, 1)], axis=1) poses_bounds[i, :15] = pose_3x5.reshape(-1) poses_bounds[i, 15] = 0.01 # near poses_bounds[i, 16] = 100.0 # far out_path = self.output_dir / "poses_bounds.npy" np.save(str(out_path), poses_bounds) print(f" [Save] Saved {n} poses to {out_path}") # Also save JSON for easier inspection json_data = { "method": result.get("method", "unknown"), "n_poses": n, "registration_rate": result.get("registration_rate", -1), "intrinsics": intrinsics.tolist() if intrinsics is not None else None, "frames": [ { "file_path": name, "transform_matrix": poses[name].tolist(), } for name in sorted_names ], } json_path = self.output_dir / "transforms.json" with open(json_path, "w") as f: json.dump(json_data, f, indent=2) print(f" [Save] Saved transforms.json to {json_path}") # --------------------------------------------------------------------------- # Holistic Gaussian Initialization (HGI) # --------------------------------------------------------------------------- def holistic_gaussian_init( sequence_dir: str, subsample: float = 0.001, exclude_tools: bool = True, ) -> Tuple[np.ndarray, np.ndarray]: """ Holistic Gaussian Initialization via depth backprojection. P = ∪_t K⁻¹ · T_t · D_t · (I_t ⊙ M_t) Backprojects depth maps from ALL frames into world coordinates, creating a dense union point cloud that covers the full scene. Tool regions are excluded via mask M_t. This avoids the sparse-initialization problem of vanilla 3DGS (which only uses COLMAP sparse points) and provides coverage of regions only visible from certain viewpoints. Args: sequence_dir: Path to organized sequence directory subsample: Fraction of points to keep (0.001 = 0.1%) exclude_tools: Whether to exclude tool regions from initialization Returns: (points [N, 3], colors [N, 3]) ready for Gaussian initialization """ # Use the unified dataset loader from scripts.download_datasets import EndoDataset dataset = EndoDataset(sequence_dir) return dataset.get_point_cloud(subsample=subsample) # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="EndoGaussian-4D Camera Pose Extraction", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--input", type=str, required=True, help="Input sequence directory") parser.add_argument("--output", type=str, default=None, help="Output directory (default: same as input)") parser.add_argument("--mode", type=str, default="auto", choices=["auto", "colmap_sequential", "colmap_exhaustive", "depth_pnp"], help="Pose extraction method") parser.add_argument("--hgi", action="store_true", help="Run Holistic Gaussian Initialization after pose extraction") parser.add_argument("--subsample", type=float, default=0.001, help="Point cloud subsample ratio for HGI (default: 0.001)") parser.add_argument("--no-gpu", action="store_true", help="Disable GPU for COLMAP") args = parser.parse_args() pipeline = EndoSfMPipeline(args.input, args.output) result = pipeline.run(mode=args.mode) print(f"\nResult: {result.get('method', 'unknown')} | " f"{len(result.get('poses', {}))} poses extracted") if args.hgi: print("\n[HGI] Running Holistic Gaussian Initialization...") points, colors = holistic_gaussian_init( args.input, subsample=args.subsample ) out_dir = Path(args.output or args.input) np.save(str(out_dir / "hgi_points.npy"), points) np.save(str(out_dir / "hgi_colors.npy"), colors) print(f"[HGI] Saved {len(points):,} points to {out_dir}") if __name__ == "__main__": main()