This is the pytorch implementation of our paper Bilateral Propagation Network for Depth Completion.
You can directly build the environment by running the following command if you use conda as the environment management tool.
conda env create -f environment.yml
Then you should have an env named bp.
Compile the C++ and CUDA code:
cd exts
python setup.py install
We train and evaluate on KITTI and NYUV2 dataset.
Please download KITTI depth completion dataset. The structure of data directory:
└── datas
└── kitti
├── data_depth_annotated
│ ├── train
│ └── val
├── data_depth_velodyne
│ ├── train
│ └── val
├── raw
│ ├── 2011_09_26
│ ├── 2011_09_28
│ ├── 2011_09_29
│ ├── 2011_09_30
│ └── 2011_10_03
├── test_depth_completion_anonymous
│ ├── image
│ ├── intrinsics
│ └── velodyne_raw
└── val_selection_cropped
├── groundtruth_depth
├── image
├── intrinsics
└── velodyne_raw
We used preprocessed NYUv2 HDF5 dataset provided by Fangchang Ma. Note, the original full NYUv2 dataset is available at the official website.
The structure of data directory:
└── datas
└── nyudepthv2
├── train
│ ├── basement_0001a
│ │ ├── 00001.h5
│ │ └── ...
│ ├── basement_0001b
│ │ ├── 00001.h5
│ │ └── ...
│ └── ...
└── val
└── official
├── 00001.h5
└── ...
You can directly download the trained models and extract them in checkpoints directory. Our models are trained on a GPU workstation with 4 Nvidia GTX 3090.
You can also train by yourself.
train on KITTI
torchrun --nproc_per_node=4 --master_port 4321 train.py \
gpus=[0,1,2,3] num_workers=4 name=BP_KITTI \
net=PMP data=KITTI \
lr=1e-3 train_batch_size=2 test_batch_size=2 \
sched/lr=NoiseOneCycleCosMo sched.lr.policy.max_momentum=0.90 \
nepoch=30 test_epoch=25 ++net.sbn=true
train on NYUV2
torchrun --nproc_per_node=2 --master_port 1100 train.py \
gpus=[0,1] num_workers=4 name=BP_NYU \
net=PMP data=NYU data.num_mask=1 \
lr=2e-3 train_batch_size=8 test_batch_size=1 \
nepoch=100 test_epoch=80 ++net.sbn=true
With the trained model, you can test and save results.
test on KITTI selval set
python test.py gpus=[0] name=BP_KITTI ++chpt=BP_KITTI \
net=PMP num_workers=4 \
data=KITTI data.testset.mode=selval \
test_batch_size=1 metric=RMSE ++net.compile=true
test on KITTI test set and save for submission
python test.py gpus=[0] name=BP_KITTI ++chpt=BP_KITTI \
net=PMP num_workers=4 \
data=KITTI data.testset.mode=test data.testset.height=352 \
test_batch_size=1 metric=RMSE ++save=true
test on NYUV2 test set
python test.py gpus=[0] name=BP_NYU ++chpt=BP_NYU \
net=PMP num_workers=4 \
data=NYU test_batch_size=1 metric=MetricALL
If you find this work useful in your research, please consider citing:
@article{BP-Net,
title={Bilateral Propagation Network for Depth Completion},
author={Tang, Jie and Tian, Fei-Peng and An, Boshi and Li, Jian and Tan, Ping},
journal={CVPR},
year={2024}
}