StereoMamba: Real-time and Robust Intraoperative Stereo Disparity Estimation via Long-range Spatial Dependencies

Overview

• [UPDATE Feb. 2026] We uploaded weights pretrained on Sceneflow.
• [UPDATE Feb. 2026] We fixed some errors in the code.
• [UPDATE Oct. 2025] This work has been accepted in IEEE RAL! Paper link: https://ieeexplore.ieee.org/document/11146458

Architecture

Requirements

• Python 3.9+
• PyTorch 2.0+
• CUDA compatible GPU
• VMamba package
• Mamba package

Setup

All experiments are conducted in docker container, docker iamge: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

1. Create and activate conda environment:

conda create -n stereomamba python=3.10
conda activate stereomamba

2. Clone and install dependencies:

git clone https://github.com/MichaelWangGo/StereoMamba.git
cd StereoMamba
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121
# you can try other cuda versions that are compatible to your machine
pip install wandb

# Install VMamba
git clone https://github.com/MzeroMiko/VMamba.git
cd VMamba
pip install -r requirements.txt
cd kernels/selective_scan && pip install .
# if you got problem when installing kernels/selective_scan, you can try: pip install --no-build-isolation .

# Install Mamba2
cd /path/to/StereoMamba/models
git clone https://github.com/state-spaces/mamba.git
cd mamba
git fetch --tags
git checkout v2.2.4
pip install --no-build-isolation .
cd /path/to/StereoMamba

Training Pipeline

1. Pretraining on Sceneflow

1. Download the Sceneflow dataset
2. Generate filename paths to .txt file (see ./datasets/filenames/)
3. Configure dataset paths in ./training_configs/pretrain/config_sceneflow.json:
   • Set data_path
   • Update trainlist and testlist
4. Start training:

python -m training_scripts.train_disparity \
    --default_config ./training_configs/pretrain/config_sceneflow.json \
    --cuda_id 0

2. Finetuning on SCARED

We provide pretrained weights on SceneFlow. Please download it here.

1. Prepare the SCARED dataset using the Unofficial SCARED toolkit
2. Configure dataset paths in ./training_configs/fine_tune/config_scared.json
3. Start finetuning:

python -m training_scripts.train_disparity \
    --default_config ./training_configs/fine_tune/config_scared.json \
    --cuda_id 0

Evaluation

Generate Disparity Maps

python -m evaluation_scripts.save_disp \
    --restore_ckpt <path/to/model> \
    --default_config <path/to/config_file> \
    --output_directory <output_path> \
    --images_filename <path/to/image_list.txt>

Evaluate on SCARED Dataset

python -m evaluation_scripts.evaluate_scared \
    --predictions <path/to/predictions> \
    --ground_truth <path/to/ground_truth> \
    --csv <output_results.csv>

Evaluate on RIS/StereoMIS

1. Modify the following paths in the script:
   • left_img_path
   • right_img_path
   • disp_path
2. Run evaluation:

python -m evaluation_scripts.reprojection_error

Additional Resources

For dataset preparation and tools:

• Unofficial SCARED toolkit - For generating disparity samples and stereo rectification