--- tags: - stereo-matching - depth-estimation - 3d-reconstruction - computer-vision - zero-shot - ml-intern license: apache-2.0 library_name: pytorch pipeline_tag: depth-estimation --- # FoundationStereo — Clean Python Wrapper A clean, single-file Python wrapper for [NVlabs/FoundationStereo](https://github.com/NVlabs/FoundationStereo) (CVPR 2025 Best Paper Nomination). **Zero-shot stereo matching** — no fine-tuning needed. Works on any stereo pair out of the box. > No argparse, no CLI. Just import and call functions from your own code. ## Quick Start ```python from foundation_stereo import FoundationStereoInference # Load model (one-time, ~5-10s) stereo = FoundationStereoInference( repo_dir="./FoundationStereo", ckpt_path="pretrained_models/23-51-11/model_best_bp2.pth", ) # Predict disparity from file paths disparity = stereo.predict("left.png", "right.png") # Or from numpy arrays (RGB uint8, shape H,W,3) disparity = stereo.predict_arrays(left_rgb, right_rgb) # Convert to metric depth depth = stereo.disparity_to_depth(disparity, focal_length=754.668, baseline=0.063) # Visualize colored = stereo.visualize_disparity(disparity) cv2.imwrite("disparity.png", colored) ``` ## Installation ### 1. Clone FoundationStereo ```bash git clone https://github.com/NVlabs/FoundationStereo.git cd FoundationStereo ``` ### 2. Create Environment ```bash conda env create -f environment.yml conda activate foundation_stereo ``` Or install manually: ```bash pip install torch torchvision omegaconf opencv-python imageio timm scipy einops xformers pip install flash-attn # optional, requires GPU compute >= 8.0 ``` ### 3. Download This Wrapper Place `foundation_stereo.py` in your project (or anywhere on your Python path): ```bash # Download from this repo wget https://huggingface.co/bdck/foundation-stereo/resolve/main/foundation_stereo.py ``` ### 4. Download Pretrained Weights The model weights are **not** included in the GitHub repo. Download them separately: #### Option A: Google Drive (Official) See the [FoundationStereo README](https://github.com/NVlabs/FoundationStereo#download-models) for Google Drive links. Download and extract to: ``` FoundationStereo/ └── pretrained_models/ ├── 23-51-11/ ← ViT-Large (best quality) │ ├── cfg.yaml │ └── model_best_bp2.pth └── 11-33-40/ ← ViT-Small (faster) ├── cfg.yaml └── model_best_bp2.pth ``` #### Option B: HuggingFace Mirror ```bash pip install huggingface-hub huggingface-cli download vitaebin/foundation-stereo-model --local-dir pretrained_models ``` ## API Reference ### `FoundationStereoInference` The main class. Initialize once, call `predict()` many times. ```python stereo = FoundationStereoInference( repo_dir="./FoundationStereo", # Path to cloned repo ckpt_path="pretrained_models/23-51-11/model_best_bp2.pth", # Checkpoint vit_size="vitl", # "vitl" (best) | "vits" (fast) | "vitb" (medium) valid_iters=32, # GRU iterations: 32=accurate, 16=fast scale=1.0, # Input downscale factor (0 < scale <= 1.0) mixed_precision=True, # AMP inference (faster, less VRAM) low_memory=False, # Trade speed for lower VRAM use_hierarchical=False, # Coarse-to-fine (recommended for >1K images) hierarchical_ratio=0.5, # First pass resolution ratio max_disp=416, # Maximum disparity search range (pixels) device="cuda:0", # GPU device seed=0, # Random seed ) ``` ### Methods | Method | Description | Returns | |--------|-------------|---------| | `predict(left_path, right_path)` | Stereo matching from image files | `np.ndarray (H, W)` float32 disparity | | `predict_arrays(left_rgb, right_rgb)` | Stereo matching from numpy arrays | `np.ndarray (H, W)` float32 disparity | | `predict_batch([(l1,r1), ...])` | Process multiple pairs | `list[np.ndarray]` | | `disparity_to_depth(disp, fx, baseline)` | Disparity → metric depth (meters) | `np.ndarray (H, W)` | | `disparity_to_depth_with_intrinsics(disp, K, baseline)` | Using full 3x3 K matrix | `np.ndarray (H, W)` | | `depth_to_pointcloud(depth, K, rgb)` | Depth → 3D points | `(points, colors)` | | `visualize_disparity(disp)` | Colored heatmap | `np.ndarray (H, W, 3)` uint8 BGR | | `load_intrinsics(path)` | Load K.txt file | `(K, baseline)` | | `save_disparity(disp, path, format)` | Save as pfm/npy/png/exr | None | ### Convenience Functions ```python from foundation_stereo import load_stereo_model, estimate_disparity # Load once, predict many stereo = load_stereo_model(repo_dir="./FoundationStereo", vit_size="vitl") disp = stereo.predict("left.png", "right.png") # One-shot (reloads model each call — avoid for batch processing) disp = estimate_disparity("left.png", "right.png", repo_dir="./FoundationStereo") ``` ### Config Dataclass ```python from foundation_stereo import FoundationStereoConfig, FoundationStereoInference config = FoundationStereoConfig( repo_dir="./FoundationStereo", vit_size="vits", # Use small model for speed valid_iters=16, # Fewer iterations for speed use_hierarchical=True, # Better for high-res ) stereo = FoundationStereoInference.from_config(config) ``` ## Models | Model | Folder | `vit_size` | Speed | Quality | VRAM | |-------|--------|-----------|-------|---------|------| | **Large** (recommended) | `23-51-11` | `"vitl"` | ~2s/pair | Best | ~8-12 GB | | **Small** | `11-33-40` | `"vits"` | ~0.5s/pair | Good | ~4-6 GB | > **Important:** The `vit_size` parameter is NOT in the released `cfg.yaml` files. This wrapper automatically injects it — just pass the correct value matching your downloaded model. ## Full Pipeline Example ```python import cv2 import numpy as np from foundation_stereo import FoundationStereoInference # ─── Configuration ─────────────────────────────────────────────── REPO_DIR = "./FoundationStereo" CKPT = "pretrained_models/23-51-11/model_best_bp2.pth" # ─── Initialize ───────────────────────────────────────────────── stereo = FoundationStereoInference(repo_dir=REPO_DIR, ckpt_path=CKPT) # ─── Predict disparity ────────────────────────────────────────── disp = stereo.predict("scene/left.png", "scene/right.png") print(f"Disparity: shape={disp.shape}, range=[{disp.min():.1f}, {disp.max():.1f}] px") # ─── Convert to depth ─────────────────────────────────────────── K, baseline = stereo.load_intrinsics("scene/K.txt") depth = stereo.disparity_to_depth_with_intrinsics(disp, K, baseline) print(f"Depth: range=[{depth[depth>0].min():.2f}, {depth[depth>0].max():.2f}] meters") # ─── Generate point cloud ─────────────────────────────────────── import imageio rgb = imageio.imread("scene/left.png") points, colors = stereo.depth_to_pointcloud(depth, K, rgb=rgb, max_depth=50.0) print(f"Point cloud: {points.shape[0]:,} points") # ─── Save outputs ─────────────────────────────────────────────── cv2.imwrite("disparity.png", stereo.visualize_disparity(disp)) stereo.save_disparity(disp, "disparity.pfm", format="pfm") stereo.save_disparity(disp, "disparity.npy", format="npy") ``` ## Intrinsics File Format (K.txt) ``` fx 0 cx 0 fy cy 0 0 1 baseline_meters ``` Example: ``` 754.668 0.0 489.379 0.0 754.668 265.161 0.0 0.0 1.0 0.063 ``` - Line 1: 3x3 intrinsic matrix K, row-major, space-separated (9 values) - Line 2: Stereo baseline in meters ## Tips - **High-res images (>1K pixels):** Use `use_hierarchical=True` for better results - **Low VRAM (<12GB):** Use `vit_size="vits"` + `low_memory=True` + `scale=0.5` - **Speed priority:** `valid_iters=16` + `vit_size="vits"` (3-4x faster) - **Best quality:** `valid_iters=32` + `vit_size="vitl"` + `use_hierarchical=True` - **Multiple GPU:** Set `device="cuda:1"` etc. for different instances - **Per-call overrides:** `stereo.predict(l, r, valid_iters=16)` without changing defaults ## Requirements - Python 3.10+ - PyTorch 2.0+ with CUDA - NVIDIA GPU (8+ GB VRAM for small model, 12+ GB for large) - Key packages: `torch`, `torchvision`, `omegaconf`, `opencv-python`, `imageio`, `timm`, `einops`, `scipy` - Optional: `flash-attn` (GPU compute >= 8.0), `open3d` (point cloud viz) ## Citation ```bibtex @inproceedings{wen2025foundationstereo, title={FoundationStereo: Zero-Shot Stereo Matching}, author={Wen, Bowen and Trepte, Matthew and Aribido, Joseph and Kautz, Jan and Birchfield, Stan and Okatani, Takayuki}, booktitle={CVPR}, year={2025} } ``` ## License This wrapper is provided under Apache-2.0. The underlying FoundationStereo model and weights are subject to [NVIDIA's license](https://github.com/NVlabs/FoundationStereo/blob/master/LICENSE). ## Generated by ML Intern This model repository was generated by [ML Intern](https://github.com/huggingface/ml-intern), an agent for machine learning research and development on the Hugging Face Hub. - Try ML Intern: https://smolagents-ml-intern.hf.space - Source code: https://github.com/huggingface/ml-intern ## Usage ```python from transformers import AutoModelForCausalLM, AutoTokenizer model_id = "bdck/foundation-stereo" tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained(model_id) ``` For non-causal architectures, replace `AutoModelForCausalLM` with the appropriate `AutoModel` class.