An Accurate and Lightweight Convolutional Neural Network for Stereo Estimation with Depth Completion
Dominik Hirner, Friedrich Fraundorfer
An pytorch implementation of our accurate and lightweight end-to-end trainable CNN for stereo estimation with depth completion. This method has been accepted and will be published at the ICPR 2022 conference. If you use our work please cite our paper
The whole project is in pure python 3 and pytorch 1.2.0
A demo of the whole end-to-end method is available online in Google Colab: Demo
Demo Video of the MPI dataset:
This repository contains
- jupyter notebooks for training and inference of disparity via a stereo-pair
- python3.6 code for training and inference
- trained weights for many publicly available datasets
Dataset | branch | simB | Incons |
---|---|---|---|
Middlebury | mb | mb_simB | Incons |
Kitti2012 | kitti2012 | kitti2012_simB | Icons_KITTI2012 |
Kitti2015 | kitti2015 | kitti2015_simB | Incons_KITTI2015 |
ETH3D | ETH | ETH_simB | Incons_ETH |
We use a trainable guided filter for the cost-volume (see project). This can be installed via pip.
pip install guided-filter-pytorch
If you want to train our method (either from scratch or continue from one of our provided weights), use the provided config files from the root of this repository and change them to fit your needs. If you want to train from scratch, first train the feature extractor and the similarity function as follows:
python FCDSN_train.py config/FCDSN-CONFIG-FILE.cfg
Afterwards, use the output created by this file to train the depth-completion part as follows:
python DC_train.py config/DC-CONFIG-FILE.cfg
Note that if you want to do transfer training on the depth-completion, the following files must be found for each sample in an individual folder:
- im0.png
- disp0GT.pfm
- disp_s.pfm
- keep_mask.png
- upd_mask.png
If you want to do inference on any rectified image-pair call the test.py function from the root of this repository as follows:
python test.py --weights_b path/to/branch_weight --weights_s path/to/sim_weights --weights_f path/to/fill_weights --left path/to/left_im.png --right /path/to/right_im.png --max_disp max_disp --out /path/to/out/out_name
Download the Middlebury weights from the link above and put it in the weights folder in the root of this repository. Then copy and paste the following:
python test.py --weights_b weights/branch/mb --weights_s weights/simb/mb_simB --weights_f weights/fill/Incons --left example/im0.png --right example/im1.png --max_disp 145 --out adirondack
If everything went ok this should produce the following output:
- adirondack.pfm: filtered disparity output of the FCDSN network
RGB | Disparity |
---|---|