MobileStereoNet

This repository contains the code for "MobileStereoNet: Towards Lightweight Deep Networks for Stereo Matching" [arXiv] [project].

Input image

2D-MobileStereoNet prediction

3D-MobileStereoNet prediction

Evaluation Results

	SceneFlow: EPE (px)	KITTI2015 val: EPE (px)
2D-MobileStereoNet	1.14	0.79
3D-MobileStereoNet	0.80	0.66

Installation

Requirements

The code is tested on:

Ubuntu 18.04
Python 3.6
PyTorch 1.4.0
Torchvision 0.5.0
CUDA 10.0

Setting up the environment

conda env create --file mobilestereonet.yml
conda activate mobilestereonet

Training

Set a variable (e.g. DATAPATH) for the dataset directory DATAPATH="/Datasets/SceneFlow/" or DATAPATH="/Datasets/KITTI2015/". Then, you can run the train.py file as below:

Pretraining on SceneFlow

python train.py --dataset sceneflow --datapath $DATAPATH --trainlist ./filenames/sceneflow_train.txt --testlist ./filenames/sceneflow_test.txt --epochs 20 --lrepochs "10,12,14,16:2" --batch_size 8 --test_batch_size 8 --model MSNet2D

Finetuning on KITTI

python train.py --dataset kitti --datapath $DATAPATH --trainlist ./filenames/kitti15_train.txt --testlist ./filenames/kitti15_val.txt --epochs 400 --lrepochs "200:10" --batch_size 8 --test_batch_size 8 --loadckpt ./checkpoints/pretrained.ckpt --model MSNet2D

The arguments in both cases can be set differently depending on the model and the system.

Prediction

The following script creates disparity maps for a specified model:

python prediction.py --datapath $DATAPATH --testlist ./filenames/kitti15_test.txt --loadckpt ./checkpoints/finetuned.ckpt --dataset kitti --colored True --model MSNet2D

Credits

The implementation of this code is based on PSMNet and GwcNet. Also, thanks to Matteo Poggi for the KITTI python utils.

License

This project is released under the Apache 2.0 license.

Citation

If you use this code, please cite this project.

@article{shamsafar2021mobilestereonet,
  title={MobileStereoNet: Towards Lightweight Deep Networks for Stereo Matching},
  author={Shamsafar, Faranak and Woerz, Samuel and Rahim, Rafia and Zell, Andreas},
  journal={arXiv preprint arXiv:2108.09770},
  year={2021}
}

scott-mao/mobilestereonet