amina193 / mobilestereonet

Lightweight stereo matching network based on MobileNetV1 and MobileNetV2

MobileStereoNet: Towards Lightweight Deep Networks for Stereo Matching

This repository contains the code for

2D-MobileStereoNet prediction	Error map

3D-MobileStereoNet prediction	Error map

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