foundation_stereo.py

"""
FoundationStereo — Clean Python API Wrapper
============================================

A clean, function-based interface to NVlabs/FoundationStereo (CVPR 2025 Best Paper Nomination).
Zero-shot stereo matching with foundation model quality.

Source: https://github.com/NVlabs/FoundationStereo

Setup:
------
1. Clone the repo:
     git clone https://github.com/NVlabs/FoundationStereo.git
     cd FoundationStereo

2. Install dependencies:
     conda env create -f environment.yml
     conda activate foundation_stereo
     pip install flash-attn  # optional, needs GPU compute >= 8.0

3. Download pretrained weights (pick one):
     # Option A: From Google Drive (see README)
     # Option B: From HuggingFace mirror
     #   huggingface-cli download vitaebin/foundation-stereo-model --local-dir pretrained_models

4. Place this file in the repo root or add the repo root to sys.path.

Usage:
------
    from foundation_stereo import FoundationStereoInference

    # Initialize once
    stereo = FoundationStereoInference(
        repo_dir="/path/to/FoundationStereo",
        ckpt_path="pretrained_models/23-51-11/model_best_bp2.pth",
    )

    # Run on a stereo pair
    disp = stereo.predict("left.png", "right.png")

    # Or with numpy arrays directly
    disp = stereo.predict_arrays(left_rgb, right_rgb)

    # Get depth map
    depth = stereo.disparity_to_depth(disp, focal_length=754.668, baseline=0.063)

    # Visualize
    colored = stereo.visualize_disparity(disp)
"""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple

import cv2
import numpy as np


# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION — All options are set here, no argparse, no CLI
# ══════════════════════════════════════════════════════════════════════════════


@dataclass
class FoundationStereoConfig:
    """All configurable parameters for FoundationStereo inference.

    Attributes:
        repo_dir: Path to the cloned FoundationStereo repository root.
        ckpt_path: Path to the .pth checkpoint file (relative to repo_dir or absolute).
                   E.g. "pretrained_models/23-51-11/model_best_bp2.pth"
        vit_size: Vision Transformer backbone size.
                  - "vitl" = ViT-Large (best accuracy, slower, for 23-51-11 model)
                  - "vits" = ViT-Small (faster, lighter, for 11-33-40 model)
                  - "vitb" = ViT-Base (medium)
        valid_iters: Number of GRU refinement iterations at inference.
                     Higher = more accurate but slower. Recommended: 32 (best), 16 (fast).
        scale: Input image downscale factor (0 < scale <= 1.0).
               Use < 1.0 to reduce memory for very large images.
        mixed_precision: Use AMP (automatic mixed precision) for inference.
                        Faster and uses less VRAM, negligible quality impact.
        low_memory: Trade speed for lower VRAM usage. Enable for GPUs with <16GB.
        use_hierarchical: Use hierarchical (coarse-to-fine) inference.
                         Recommended for images > 1K resolution.
        hierarchical_ratio: Resolution ratio for the first (coarse) pass in hierarchical mode.
                           Typically 0.5 (half resolution first pass).
        max_disp: Maximum disparity search range (pixels). Default 416 from config.
                  Increase for close-range scenes with large disparities.
        device: CUDA device string. E.g. "cuda:0", "cuda:1".
        seed: Random seed for reproducibility.
    """

    repo_dir: str = "./FoundationStereo"
    ckpt_path: str = "pretrained_models/23-51-11/model_best_bp2.pth"
    vit_size: str = "vitl"
    valid_iters: int = 32
    scale: float = 1.0
    mixed_precision: bool = True
    low_memory: bool = False
    use_hierarchical: bool = False
    hierarchical_ratio: float = 0.5
    max_disp: int = 416
    device: str = "cuda:0"
    seed: int = 0


# ══════════════════════════════════════════════════════════════════════════════
# MAIN INFERENCE CLASS
# ══════════════════════════════════════════════════════════════════════════════


class FoundationStereoInference:
    """Clean Python interface for FoundationStereo inference.

    Handles model loading, preprocessing, inference, and postprocessing.
    Call predict() or predict_arrays() for stereo disparity estimation.

    Example:
        stereo = FoundationStereoInference(
            repo_dir="/path/to/FoundationStereo",
            ckpt_path="pretrained_models/23-51-11/model_best_bp2.pth",
        )
        disp = stereo.predict("left.png", "right.png")
    """

    def __init__(
        self,
        repo_dir: str = "./FoundationStereo",
        ckpt_path: str = "pretrained_models/23-51-11/model_best_bp2.pth",
        vit_size: str = "vitl",
        valid_iters: int = 32,
        scale: float = 1.0,
        mixed_precision: bool = True,
        low_memory: bool = False,
        use_hierarchical: bool = False,
        hierarchical_ratio: float = 0.5,
        max_disp: int = 416,
        device: str = "cuda:0",
        seed: int = 0,
    ):
        """Initialize FoundationStereo. Alternatively use from_config()."""
        self.config = FoundationStereoConfig(
            repo_dir=repo_dir,
            ckpt_path=ckpt_path,
            vit_size=vit_size,
            valid_iters=valid_iters,
            scale=scale,
            mixed_precision=mixed_precision,
            low_memory=low_memory,
            use_hierarchical=use_hierarchical,
            hierarchical_ratio=hierarchical_ratio,
            max_disp=max_disp,
            device=device,
            seed=seed,
        )
        self._model = None
        self._cfg = None
        self._setup_imports()
        self._load_model()

    @classmethod
    def from_config(cls, config: FoundationStereoConfig) -> "FoundationStereoInference":
        """Create instance from a config dataclass."""
        return cls(**config.__dict__)

    # ──────────────────────────────────────────────────────────────────────
    # PUBLIC API
    # ──────────────────────────────────────────────────────────────────────

    def predict(
        self,
        left_path: str,
        right_path: str,
        scale: Optional[float] = None,
        valid_iters: Optional[int] = None,
        use_hierarchical: Optional[bool] = None,
    ) -> np.ndarray:
        """Run stereo matching on a pair of image files.

        Args:
            left_path: Path to the left (reference) image. Any format imageio supports.
            right_path: Path to the right image.
            scale: Override config scale for this call (0 < scale <= 1.0).
            valid_iters: Override GRU iterations for this call.
            use_hierarchical: Override hierarchical mode for this call.

        Returns:
            Disparity map as float32 numpy array, shape (H, W).
            Values are in pixels (at the possibly-scaled resolution).
            Higher values = closer to camera.
        """
        import imageio

        left_img = imageio.imread(left_path)   # Returns RGB uint8 (H, W, 3)
        right_img = imageio.imread(right_path)

        return self.predict_arrays(
            left_img, right_img,
            scale=scale,
            valid_iters=valid_iters,
            use_hierarchical=use_hierarchical,
        )

    def predict_arrays(
        self,
        left_img: np.ndarray,
        right_img: np.ndarray,
        scale: Optional[float] = None,
        valid_iters: Optional[int] = None,
        use_hierarchical: Optional[bool] = None,
    ) -> np.ndarray:
        """Run stereo matching on numpy arrays.

        Args:
            left_img: Left image as numpy array, shape (H, W, 3), RGB, uint8.
            right_img: Right image as numpy array, shape (H, W, 3), RGB, uint8.
            scale: Override config scale (0 < scale <= 1.0).
            valid_iters: Override GRU iterations.
            use_hierarchical: Override hierarchical mode.

        Returns:
            Disparity map as float32 numpy array, shape (H, W).
        """
        import torch

        # Resolve parameters (per-call overrides > config defaults)
        _scale = scale if scale is not None else self.config.scale
        _iters = valid_iters if valid_iters is not None else self.config.valid_iters
        _hierarchical = (
            use_hierarchical if use_hierarchical is not None else self.config.use_hierarchical
        )

        # Validate inputs
        assert left_img.ndim == 3 and left_img.shape[2] == 3, (
            f"Expected (H, W, 3) RGB image, got shape {left_img.shape}"
        )
        assert left_img.shape == right_img.shape, (
            f"Left/right shape mismatch: {left_img.shape} vs {right_img.shape}"
        )

        # Optional downscale
        if _scale != 1.0:
            assert 0 < _scale <= 1.0, f"Scale must be in (0, 1], got {_scale}"
            left_img = cv2.resize(
                left_img, None, fx=_scale, fy=_scale, interpolation=cv2.INTER_LINEAR
            )
            right_img = cv2.resize(
                right_img, None, fx=_scale, fy=_scale, interpolation=cv2.INTER_LINEAR
            )

        H, W = left_img.shape[:2]

        # Convert to torch tensors: (B, C, H, W), float32, values 0-255
        # NOTE: model.forward() normalizes internally — do NOT pre-normalize
        img0_t = (
            torch.as_tensor(left_img).to(self.config.device).float().permute(2, 0, 1).unsqueeze(0)
        )
        img1_t = (
            torch.as_tensor(right_img).to(self.config.device).float().permute(2, 0, 1).unsqueeze(0)
        )

        # Pad to be divisible by 32 (required by model architecture)
        padder = self._InputPadder(img0_t.shape, divis_by=32)
        img0_t, img1_t = padder.pad(img0_t, img1_t)

        # Forward pass with AMP
        with torch.cuda.amp.autocast(enabled=self.config.mixed_precision):
            if not _hierarchical:
                disp = self._model.forward(
                    img0_t,
                    img1_t,
                    iters=_iters,
                    test_mode=True,
                    low_memory=self.config.low_memory,
                )
            else:
                disp = self._model.run_hierachical(
                    img0_t,
                    img1_t,
                    iters=_iters,
                    test_mode=True,
                    low_memory=self.config.low_memory,
                    small_ratio=self.config.hierarchical_ratio,
                )

        # Unpad and convert to numpy
        disp = padder.unpad(disp.float())
        disp_np = disp.squeeze().cpu().numpy()  # (H, W)

        assert disp_np.shape == (H, W), (
            f"Output shape mismatch: {disp_np.shape} vs expected ({H}, {W})"
        )
        return disp_np

    def predict_batch(
        self,
        pairs: list[Tuple[str, str]],
        scale: Optional[float] = None,
        valid_iters: Optional[int] = None,
        use_hierarchical: Optional[bool] = None,
    ) -> list[np.ndarray]:
        """Run stereo matching on multiple pairs sequentially.

        Args:
            pairs: List of (left_path, right_path) tuples.
            scale: Override config scale.
            valid_iters: Override GRU iterations.
            use_hierarchical: Override hierarchical mode.

        Returns:
            List of disparity maps, each shape (H, W) float32.
        """
        results = []
        for left_path, right_path in pairs:
            disp = self.predict(
                left_path,
                right_path,
                scale=scale,
                valid_iters=valid_iters,
                use_hierarchical=use_hierarchical,
            )
            results.append(disp)
        return results

    # ──────────────────────────────────────────────────────────────────────
    # UTILITY METHODS
    # ──────────────────────────────────────────────────────────────────────

    @staticmethod
    def disparity_to_depth(
        disparity: np.ndarray,
        focal_length: float,
        baseline: float,
        min_depth: float = 0.01,
        max_depth: float = 100.0,
    ) -> np.ndarray:
        """Convert disparity map to metric depth map.

        Formula: depth = focal_length * baseline / disparity

        Args:
            disparity: Disparity in pixels, shape (H, W), float32.
            focal_length: Camera focal length in pixels (fx from intrinsic matrix K[0,0]).
                         If images were scaled, use: fx_original * scale.
            baseline: Stereo baseline in meters (distance between cameras).
            min_depth: Minimum valid depth (meters). Pixels below this are clipped.
            max_depth: Maximum valid depth (meters). Pixels above this are clipped.

        Returns:
            Depth map in meters, shape (H, W), float32.
            Invalid pixels (disparity <= 0) are set to 0.
        """
        depth = np.zeros_like(disparity)
        valid = disparity > 0
        depth[valid] = (focal_length * baseline) / disparity[valid]
        depth = np.clip(depth, min_depth, max_depth)
        depth[~valid] = 0.0
        return depth

    @staticmethod
    def disparity_to_depth_with_intrinsics(
        disparity: np.ndarray,
        K: np.ndarray,
        baseline: float,
        scale: float = 1.0,
        min_depth: float = 0.01,
        max_depth: float = 100.0,
    ) -> np.ndarray:
        """Convert disparity to depth using full intrinsic matrix.

        Args:
            disparity: Disparity in pixels, shape (H, W), float32.
            K: 3x3 camera intrinsic matrix (for original image resolution).
            baseline: Stereo baseline in meters.
            scale: Image scale factor that was applied (adjusts focal length).
            min_depth: Minimum valid depth (meters).
            max_depth: Maximum valid depth (meters).

        Returns:
            Depth map in meters, shape (H, W).
        """
        fx = K[0, 0] * scale
        return FoundationStereoInference.disparity_to_depth(
            disparity, fx, baseline, min_depth, max_depth
        )

    @staticmethod
    def depth_to_pointcloud(
        depth: np.ndarray,
        K: np.ndarray,
        rgb: Optional[np.ndarray] = None,
        scale: float = 1.0,
        max_depth: float = 50.0,
    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """Back-project depth map to 3D point cloud.

        Args:
            depth: Depth map in meters, shape (H, W).
            K: 3x3 camera intrinsic matrix (for original resolution).
            rgb: Optional RGB image for coloring points, shape (H, W, 3), uint8.
            scale: Scale factor applied to images (adjusts K accordingly).
            max_depth: Maximum depth to include in point cloud.

        Returns:
            Tuple of:
                - points: (N, 3) float32 array of 3D points [x, y, z] in meters.
                - colors: (N, 3) float32 array of RGB colors in [0, 1], or None.
        """
        H, W = depth.shape
        K_scaled = K.copy().astype(np.float64)
        K_scaled[:2] *= scale

        fx, fy = K_scaled[0, 0], K_scaled[1, 1]
        cx, cy = K_scaled[0, 2], K_scaled[1, 2]

        # Create pixel coordinate grid
        u, v = np.meshgrid(np.arange(W), np.arange(H))

        # Back-project: X = (u - cx) * Z / fx, Y = (v - cy) * Z / fy
        z = depth
        x = (u - cx) * z / fx
        y = (v - cy) * z / fy

        # Stack to (H, W, 3)
        xyz = np.stack([x, y, z], axis=-1)

        # Filter valid points
        valid = (depth > 0) & (depth < max_depth)
        points = xyz[valid].astype(np.float32)

        colors = None
        if rgb is not None:
            colors = rgb[valid].astype(np.float32) / 255.0

        return points, colors

    @staticmethod
    def visualize_disparity(
        disparity: np.ndarray,
        colormap: int = cv2.COLORMAP_INFERNO,
        max_disp: Optional[float] = None,
    ) -> np.ndarray:
        """Create a colored visualization of a disparity map.

        Args:
            disparity: Disparity map, shape (H, W), float32.
            colormap: OpenCV colormap constant. Options:
                      cv2.COLORMAP_INFERNO (default, good contrast)
                      cv2.COLORMAP_TURBO   (rainbow)
                      cv2.COLORMAP_MAGMA   (dark-to-bright)
                      cv2.COLORMAP_JET     (classic)
            max_disp: Maximum disparity for normalization. None = auto (99th percentile).

        Returns:
            Colored image, shape (H, W, 3), uint8, BGR format.
            Save with cv2.imwrite() or convert: cv2.cvtColor(..., cv2.COLOR_BGR2RGB).
        """
        disp_vis = disparity.copy()
        disp_vis[disp_vis <= 0] = 0

        if max_disp is None:
            valid_pixels = disp_vis[disp_vis > 0]
            max_disp = np.percentile(valid_pixels, 99) if len(valid_pixels) > 0 else 1.0

        disp_normalized = np.clip(disp_vis / max_disp, 0, 1)
        disp_uint8 = (disp_normalized * 255).astype(np.uint8)
        colored = cv2.applyColorMap(disp_uint8, colormap)

        return colored

    @staticmethod
    def load_intrinsics(intrinsics_path: str) -> Tuple[np.ndarray, float]:
        """Load camera intrinsics from a K.txt file (FoundationStereo format).

        File format:
            Line 1: 9 space-separated floats = 3x3 K matrix (row-major)
            Line 2: single float = baseline in meters

        Example K.txt:
            754.668 0.0 489.379 0.0 754.668 265.161 0.0 0.0 1.0
            0.063

        Args:
            intrinsics_path: Path to the intrinsics file.

        Returns:
            Tuple of (K, baseline):
                - K: 3x3 intrinsic matrix, float32.
                - baseline: Stereo baseline in meters (float).
        """
        with open(intrinsics_path, "r") as f:
            lines = f.readlines()

        K = np.array(
            list(map(float, lines[0].strip().split()))
        ).reshape(3, 3).astype(np.float32)
        baseline = float(lines[1].strip())
        return K, baseline

    @staticmethod
    def save_disparity(
        disparity: np.ndarray,
        output_path: str,
        format: str = "pfm",
    ) -> None:
        """Save disparity map to file.

        Args:
            disparity: Disparity map, shape (H, W), float32.
            output_path: Output file path.
            format: Output format:
                    - "pfm": Portable FloatMap (lossless, standard for stereo benchmarks)
                    - "npy": NumPy binary (fast, lossless)
                    - "png": 16-bit PNG (x256 for sub-pixel precision)
                    - "exr": OpenEXR float (requires imageio[openexr])
        """
        if format == "npy":
            np.save(output_path, disparity)

        elif format == "pfm":
            _write_pfm(output_path, disparity)

        elif format == "png":
            disp_16bit = (disparity * 256.0).astype(np.uint16)
            cv2.imwrite(output_path, disp_16bit)

        elif format == "exr":
            import imageio
            imageio.imwrite(output_path, disparity)

        else:
            raise ValueError(
                f"Unknown format: '{format}'. Use 'pfm', 'npy', 'png', or 'exr'."
            )

    # ──────────────────────────────────────────────────────────────────────
    # PRIVATE METHODS
    # ──────────────────────────────────────────────────────────────────────

    def _setup_imports(self):
        """Add FoundationStereo repo to sys.path so internal imports work."""
        repo_dir = str(Path(self.config.repo_dir).resolve())
        if repo_dir not in sys.path:
            sys.path.insert(0, repo_dir)

        # Verify critical files exist
        core_dir = os.path.join(repo_dir, "core")
        if not os.path.isdir(core_dir):
            raise FileNotFoundError(
                f"Cannot find 'core/' directory in repo_dir='{repo_dir}'. "
                f"Make sure repo_dir points to the cloned FoundationStereo repository root."
            )

    def _load_model(self):
        """Load the FoundationStereo model from checkpoint."""
        import torch
        from omegaconf import OmegaConf
        from core.foundation_stereo import FoundationStereo
        from core.utils.utils import InputPadder

        # Store InputPadder class for use in predict
        self._InputPadder = InputPadder

        # Resolve checkpoint path
        ckpt_path = self.config.ckpt_path
        if not os.path.isabs(ckpt_path):
            ckpt_path = os.path.join(self.config.repo_dir, ckpt_path)
        ckpt_path = str(Path(ckpt_path).resolve())

        if not os.path.isfile(ckpt_path):
            raise FileNotFoundError(
                f"Checkpoint not found at '{ckpt_path}'. "
                f"Download weights from Google Drive or HuggingFace mirror. "
                f"See: https://github.com/NVlabs/FoundationStereo#download-models"
            )

        # Load config from same directory as checkpoint
        cfg_path = os.path.join(os.path.dirname(ckpt_path), "cfg.yaml")
        if not os.path.isfile(cfg_path):
            raise FileNotFoundError(
                f"Config file not found at '{cfg_path}'. "
                f"The cfg.yaml must be in the same directory as the .pth checkpoint."
            )

        cfg = OmegaConf.load(cfg_path)

        # Inject vit_size if not in config (required for both released models)
        if "vit_size" not in cfg:
            cfg["vit_size"] = self.config.vit_size

        # Override with our runtime parameters
        cfg["valid_iters"] = self.config.valid_iters
        cfg["mixed_precision"] = self.config.mixed_precision
        cfg["low_memory"] = int(self.config.low_memory)
        cfg["max_disp"] = self.config.max_disp

        self._cfg = OmegaConf.create(cfg)

        # Set seed for reproducibility
        _set_seed(self.config.seed)

        # Build and load model
        model = FoundationStereo(self._cfg)
        ckpt = torch.load(ckpt_path, map_location="cpu")
        model.load_state_dict(ckpt["model"], strict=True)
        model.to(self.config.device)
        model.eval()

        # Disable gradient computation globally for inference
        torch.set_grad_enabled(False)

        self._model = model
        print(f"[FoundationStereo] Model loaded successfully from: {ckpt_path}")
        print(
            f"[FoundationStereo] ViT size: {self.config.vit_size} | "
            f"Iters: {self.config.valid_iters} | "
            f"Device: {self.config.device}"
        )


# ══════════════════════════════════════════════════════════════════════════════
# STANDALONE HELPER FUNCTIONS (can be used without the class)
# ══════════════════════════════════════════════════════════════════════════════


def load_stereo_model(
    repo_dir: str = "./FoundationStereo",
    ckpt_path: str = "pretrained_models/23-51-11/model_best_bp2.pth",
    vit_size: str = "vitl",
    valid_iters: int = 32,
    device: str = "cuda:0",
    **kwargs,
) -> FoundationStereoInference:
    """Convenience function to load FoundationStereo.

    Args:
        repo_dir: Path to cloned repo.
        ckpt_path: Relative or absolute path to .pth file.
        vit_size: "vitl" (best), "vits" (fast), or "vitb" (medium).
        valid_iters: GRU iterations (32=best, 16=fast).
        device: CUDA device.
        **kwargs: Any other FoundationStereoConfig fields.

    Returns:
        FoundationStereoInference instance ready for prediction.
    """
    return FoundationStereoInference(
        repo_dir=repo_dir,
        ckpt_path=ckpt_path,
        vit_size=vit_size,
        valid_iters=valid_iters,
        device=device,
        **kwargs,
    )


def estimate_disparity(
    left_path: str,
    right_path: str,
    repo_dir: str = "./FoundationStereo",
    ckpt_path: str = "pretrained_models/23-51-11/model_best_bp2.pth",
    vit_size: str = "vitl",
    valid_iters: int = 32,
    device: str = "cuda:0",
    **kwargs,
) -> np.ndarray:
    """One-shot convenience: load model and predict.

    WARNING: Loads the model every call. For multiple predictions,
    use load_stereo_model() once and call .predict() repeatedly.
    """
    model = load_stereo_model(
        repo_dir=repo_dir,
        ckpt_path=ckpt_path,
        vit_size=vit_size,
        valid_iters=valid_iters,
        device=device,
        **kwargs,
    )
    return model.predict(left_path, right_path)


# ══════════════════════════════════════════════════════════════════════════════
# INTERNAL UTILITIES
# ══════════════════════════════════════════════════════════════════════════════


def _set_seed(seed: int):
    """Set random seeds for reproducibility."""
    import torch
    import random

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _write_pfm(path: str, image: np.ndarray, scale: float = 1.0):
    """Write a PFM (Portable FloatMap) file."""
    if image.ndim == 2:
        color = False
    elif image.ndim == 3 and image.shape[2] == 3:
        color = True
    else:
        raise ValueError(f"Unsupported image shape for PFM: {image.shape}")

    with open(path, "wb") as f:
        header = "PF\n" if color else "Pf\n"
        f.write(header.encode())
        f.write(f"{image.shape[1]} {image.shape[0]}\n".encode())
        # PFM uses negative scale for little-endian
        endian = image.dtype.byteorder
        if endian == "<" or (endian == "=" and sys.byteorder == "little"):
            scale = -scale
        f.write(f"{scale}\n".encode())
        # PFM stores rows bottom-to-top
        image = np.flipud(image).astype(np.float32)
        f.write(image.tobytes())


def _save_ply(path: str, points: np.ndarray, colors: Optional[np.ndarray] = None):
    """Save point cloud as PLY file (simple ASCII format).

    Viewable in MeshLab, CloudCompare, Open3D, Blender, etc.
    """
    n = points.shape[0]
    has_color = colors is not None

    with open(path, "w") as f:
        f.write("ply\n")
        f.write("format ascii 1.0\n")
        f.write(f"element vertex {n}\n")
        f.write("property float x\n")
        f.write("property float y\n")
        f.write("property float z\n")
        if has_color:
            f.write("property uchar red\n")
            f.write("property uchar green\n")
            f.write("property uchar blue\n")
        f.write("end_header\n")

        for i in range(n):
            line = f"{points[i, 0]:.6f} {points[i, 1]:.6f} {points[i, 2]:.6f}"
            if has_color:
                r, g, b = (colors[i] * 255).astype(np.uint8)
                line += f" {r} {g} {b}"
            f.write(line + "\n")


# ══════════════════════════════════════════════════════════════════════════════
# EXAMPLE USAGE (runs when this file is executed directly)
# ══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    # ┌─────────────────────────────────────────────────────────────────────┐
    # │ CONFIGURATION — Edit these values for your setup                    │
    # └─────────────────────────────────────────────────────────────────────┘

    REPO_DIR = "./FoundationStereo"                                  # Path to cloned repo
    CKPT_PATH = "pretrained_models/23-51-11/model_best_bp2.pth"      # Checkpoint file
    VIT_SIZE = "vitl"                                                # "vitl" (best) or "vits" (fast)
    VALID_ITERS = 32                                                 # 32 = accurate, 16 = fast
    SCALE = 1.0                                                      # Downscale factor (1.0 = full res)
    USE_HIERARCHICAL = False                                         # True for high-res (>1K) images
    DEVICE = "cuda:0"                                                # GPU device
    MIXED_PRECISION = True                                           # AMP inference

    LEFT_IMAGE = "assets/left.png"                                   # Left stereo image
    RIGHT_IMAGE = "assets/right.png"                                 # Right stereo image
    INTRINSICS_FILE = "assets/K.txt"                                 # Camera intrinsics (optional)

    OUTPUT_DIR = "./output"                                          # Where to save results

    # ┌─────────────────────────────────────────────────────────────────────┐
    # │ RUN INFERENCE                                                       │
    # └─────────────────────────────────────────────────────────────────────┘

    # 1. Initialize model (loads weights, one-time cost ~5-10s)
    stereo = FoundationStereoInference(
        repo_dir=REPO_DIR,
        ckpt_path=CKPT_PATH,
        vit_size=VIT_SIZE,
        valid_iters=VALID_ITERS,
        scale=SCALE,
        mixed_precision=MIXED_PRECISION,
        use_hierarchical=USE_HIERARCHICAL,
        device=DEVICE,
    )

    # 2. Run stereo matching
    disparity = stereo.predict(
        os.path.join(REPO_DIR, LEFT_IMAGE),
        os.path.join(REPO_DIR, RIGHT_IMAGE),
    )
    print(f"Disparity shape: {disparity.shape}")
    print(f"Disparity range: [{disparity.min():.2f}, {disparity.max():.2f}] pixels")

    # 3. Visualize disparity
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    colored_disp = stereo.visualize_disparity(disparity)
    cv2.imwrite(os.path.join(OUTPUT_DIR, "disparity_colored.png"), colored_disp)
    print(f"Saved colored disparity to {OUTPUT_DIR}/disparity_colored.png")

    # 4. (Optional) Convert to depth if intrinsics are available
    intrinsics_path = os.path.join(REPO_DIR, INTRINSICS_FILE)
    if os.path.isfile(intrinsics_path):
        K, baseline = stereo.load_intrinsics(intrinsics_path)
        depth = stereo.disparity_to_depth_with_intrinsics(
            disparity, K, baseline, scale=SCALE
        )
        print(f"Depth range: [{depth[depth > 0].min():.3f}, {depth[depth > 0].max():.3f}] meters")

        # 5. (Optional) Generate point cloud
        import imageio

        left_rgb = imageio.imread(os.path.join(REPO_DIR, LEFT_IMAGE))
        if SCALE != 1.0:
            left_rgb = cv2.resize(left_rgb, None, fx=SCALE, fy=SCALE)

        points, colors = stereo.depth_to_pointcloud(
            depth, K, rgb=left_rgb, scale=SCALE, max_depth=50.0
        )
        print(f"Point cloud: {points.shape[0]} points")

        # Save as PLY
        _save_ply(os.path.join(OUTPUT_DIR, "pointcloud.ply"), points, colors)
        print(f"Saved point cloud to {OUTPUT_DIR}/pointcloud.ply")

    # 6. Save raw disparity
    stereo.save_disparity(
        disparity, os.path.join(OUTPUT_DIR, "disparity.npy"), format="npy"
    )
    print(f"Saved raw disparity to {OUTPUT_DIR}/disparity.npy")