Update ML Intern artifact metadata

24210ae verified about 1 month ago

10 kB

	---
	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).

	<!-- ml-intern-provenance -->
	## 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.