antabangun / coex

GNU General Public License v3.0
142 stars 19 forks source link

CoEx

PyTorch implementation of our paper:

Correlate-and-Excite: Real-Time Stereo Matching via Guided Cost Volume Excitation
Authors: Antyanta Bangunharcana1, Jae Won Cho2, Seokju Lee2, In So Kweon2, Kyung-Soo Kim1, Soohyun Kim1
1MSC Lab, 2RVC Lab, Korea Advanced Institute of Science and Technology (KAIST)
IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2021

[Project page] | [Paper]

We propose a Guided Cost volume Excitation (GCE) and top-k soft-argmax disparity regression for real-time and accurate stereo matching.

Contents

Installation

We recommend using conda for installation:

conda env create -f environment.yml
conda activate coex

Update: SceneFlow model

Model Weights

Our pre-trained SceneFlow weights can be downloaded via the following link:

Performance

Our model achieves a new SceneFlow EPE (End-Point-Error) of 0.596, improving upon the previous EPE of 0.69 reported in the original paper.

Training Details

Datasets

Data for demo

For a demo of our code on the KITTI dataset, download the "[synced+rectified data]" from raw KITTI data. Unzip and place the extracted folders following the directory tree below.

If you want to re-train the models

Sceneflow dataset
Download the finalpass data of the Sceneflow dataset as well as the Disparity data.

KITTI 2015
Download kitti15 dataset, and unzip data_scene_flow.zip, rename it as kitti15, and move it into SceneFlow directory as shown in the tree below.

KITTI 2012
Download kitti12 dataset. Unzip data_stereo_flow.zip, rename it as kitti12, and move it into SceneFlow directory as shown in the tree below.

Make sure the directory names matches the tree below so that the dataloaders can locate the files.

Data directories

In our setup, the dataset is organized as follows

../../data
└── datasets
    ├── KITTI_raw
    |   ├── 2011_09_26
    |   │   ├── 2011_09_26_drive_0001_sync
    |   │   ├── 2011_09_26_drive_0002_sync
    |   |       :
    |   |
    |   ├── 2011_09_28
    |   │   ├── 2011_09_28_drive_0001_sync
    |   │   └── 2011_09_28_drive_0002_sync
    |   |       :
    |   |   :    
    |
    └── SceneFlow
        ├── driving
        │   ├── disparity
        │   └── frames_finalpass
        ├── flyingthings3d_final
        │   ├── disparity
        │   └── frames_finalpass
        ├── monkaa
        │   ├── disparity
        │   └── frames_finalpass
        ├── kitti12
        │   ├── testing
        │   └── training
        └── kitti15
            ├── testing
            └── training

Demo on KITTI raw data

The pretrained KITTI model is already included in './logs'. Run

python demo.py

to perform stereo matching on raw kitti sequence. Here is an example result on our system with RTX 2080Ti on Ubuntu 18.04.

For more demo results, check out our Project page

Re-training the model

To re-train the model, configure './configs/stereo/cfg_yaml', e.g., batch_size, paths, device num, precision, etc. Then run

python stereo.py

Citation

If you find our work useful in your research, please consider citing our paper

@inproceedings{bangunharcana2021correlate,
  title={Correlate-and-Excite: Real-Time Stereo Matching via Guided Cost Volume Excitation},
  author={Bangunharcana, Antyanta and Cho, Jae Won and Lee, Seokju and Kweon, In So and Kim, Kyung-Soo and Kim, Soohyun},
  booktitle={2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
  pages={3542--3548},
  year={2021},
  organization={IEEE}
}

Acknowledgements

Part of the code is adopted from previous works: PSMNet, AANet, GANet, SpixelFCN