scripts/extract_poses.py

#!/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()