Skip to content

Latest commit

 

History

History
123 lines (92 loc) · 5.72 KB

README.md

File metadata and controls

123 lines (92 loc) · 5.72 KB
Raw

Kitti support for CSPN

Fixed several bugs for CSPN and wrapped up for kitti dataset.

  • Create the datalist:
cd ~/CSPN/
cd ./datalist
python dataListGenerator.py your_h5_dataset_base_dir dataset_name

where dataset_name is kitti or nyudepth here.

  • run the bash train_cspn_kitti.sh to train on kitti dataset.
  • run the bash eval_kittidepth_cspn.sh to evaluate on kitti dataset.

Depth Estimation via Affinity Learned with Convolutional Spatial Propagation Network

By Xinjing Cheng*, Peng Wang* and Ruigang Yang (*Equal contribution)

Contents

  1. Introduction
  2. Requirements
  3. Models
  4. Testing
  5. Training
  6. Citation

Introduction

This repo contains the CNN models trained for depth completion from a RGB and sparse depth points, as described in the paper "Depth Estimation via Affinity Learned with Convolutional Spatial Propagation Network". The provided models are those that were used to obtain the results reported in the paper on the benchmark datasets NYU Depth v2 and KITTI for indoor and outdoor scenes respectively. Moreover, the provided code can be used for inference on arbitrary images.

Notice: there is a minor formulation error in the original paper, i.e. way to apply affinity kernel at center pixel (please check code for details, we will rectify it in our journal submission.)

Requirements

This code was tested with Python 3 and PyTorch 0.4.0.

  • Install PyTorch on a machine with CUDA GPU.
  • Install the HDF5 and other dependencies (files in our pre-processed datasets are in HDF5 formats). 
     sudo apt-get update
 sudo apt-get install -y libhdf5-serial-dev hdf5-tools
 pip3 install h5py pandas matplotlib imageio scikit-image opencv-python
  • Download the preprocessed NYU Depth V2 and/or KITTI Odometry dataset in HDF5 formats, and place them under the data folder. The downloading process might take an hour or so. The NYU dataset requires 32G of storage space, and KITTI requires 81G.
  • for NYU dataset 
     mkdir data; cd data
 wget http://datasets.lids.mit.edu/sparse-to-dense/data/nyudepthv2.tar.gz
 tar -xvf nyudepthv2.tar.gz && rm -f nyudepthv2.tar.gz
 mv nyudepthv2 nyudepth_hdf5
  • for KITTI dataset 
    wget http://datasets.lids.mit.edu/sparse-to-dense/data/kitti.tar.gz
tar -xvf kitti.tar.gz && rm -f kitti.tar.gz
mv kitti kitti_hdf5
cd ..

Models

The pretrained models, resnet18 and resnet50 can be downloaded here, and please put it into pretrained folder, if this folder does not exit, please run mkdir pretrained in root folder:

The trained models, namely +UNet+CSPN in the paper can be downloaded here:

  • NYU Depth V2: Pytorch model (Deprecated, only results images are available)
  • NYU Depth V2 (Fast Unpool, pos): Pytorch model
  • NYU Depth V2 (Fast Unpool, non-pos): Pytorch model
  • KITTI: Pytorch model(coming soon)

Download it under output/${dataset}_pretrain_cspn_${model_config}/, where dataset could be nyu or kitti, where model_config can be checked from eval_nyudepth_cspn.sh

Testing

  • For NYU Depth v2

Here we provide example for the model of NYU(Fast Unpool, non-pos affinity). Download the model from above link and put it under output/nyu_pretrain_cspn_1_net_cp500_bs8_adlr_ep40_8norm, then run,

    bash eval_nyudepth_cspn.sh

Run it multiple times and take mean, you should able obtain our depth results close here (5 time average due to randomness of sampled sparse points):

Data RMSE REL DELTA1.02 DELTA1.05 DELTA1.10 Results
NYU(Slow Unpool) 0.1165 0.0159 0.8331 0.9366 0.9716 Download
NYU(Fast Unpool, pos affinity) 0.1169 0.0161 0.8300 0.9347 0.9708 Download
NYU(Fast Unpool, non-pos affinity) 0.1172 0.0160 0.8344 0.9351 0.9707 Download

Here, the Slow Unpool means we originally loop over the image for unpooling. the Fast Unpool means we use adopt transpose conv to implement the unpooling. pos affinity means we enforce the affinity to be positive, i.e. affinities are normalized in [0, 1). non-pos affinity means we allow negative affinity, i.e. affinity are normalized in (-1, 1).

Training

  • For NYU Depth

You may set the configuration in train_cspn_nyu.sh and run,

    bash train_cspn_nyu.sh

We train with a Nvidia 1080Ti GPU, and within 40 epochs you should be able to get results close to what we reported above.

Citation

If you use this method in your research, please cite:

@inproceedings{cheng2018depth,
  title={Depth estimation via affinity learned with convolutional spatial propagation network},
  author={Cheng, Xinjing and Wang, Peng and Yang, Ruigang},
  booktitle={Proceedings of the European Conference on Computer Vision (ECCV)},
  pages={103--119},
  year={2018}
}