This repo provides an official PyTorch implementation of "Learning a Geometric Representation for Data-Efficient Depth Estimation via Gradient Field and Contrastive Loss" (ICRA 2021). [paper]
It is recommended to create a new Anaconda virtual environment for reproduction or evaluation with pretrained models.
conda create -n grmc python==3.5.6
conda activate grmc
conda install pytorch=1.6.0 torchvision=0.5.0 -c pytorch
We ran our experiments with PyTorch 1.6.0, CUDA 10.2, Python 3.5.6 and Ubuntu 18.04. Usage of higher or lower version of PyTorch seems to be incompatible to our pre-trained model.
pip install pillow==5.2.0 opencv-contrib-python
If you do not use Anaconda environment, please use pip3 rather than pip for dependencies with Python3.
Download the subset of preprocessed NYU Depth v2 (50k) here for pretraining the encoder and training the depth estimation network (source). Extract the zip file and do not delete it after extraction. The directory is then be like
└── NYU_Depth_v2_DATASET_ROOT
|
├── data
| ├── nyu2_train
| ├── nyu2_train.csv
| ├── nyu2_test
| └── nyu2_test.csv
|
└── nyu_data.zip
Any parametric model can be trained with our proposed self-supervised algorithm, and we provide pretrained models of ResNet-50 and DenseNet-161 as the encoder of depth estimation network. Download the pretrained weights and place them in the checkpoints folder to train the depth estimation network with pretrained encoders.
Encoder | batch_size |
---|---|
DenseNet-161 | 16 |
ResNet-50 | 64 |
We use RTX 2080Ti (11GB) for training the encoder and its batch_size can be resized as you use GPU with different memory size.
python encoder_pretrain.py --encoder_type=densenet --layers=161 -b=8
We provide weights of two monocular depth estimation networks, DenseDepth and FCRN with pretrained encoders by our proposed algorithm. Download the pretrained weights and place them on the checkpoints for the evaluation.
Model | Encoder | batch_size |
---|---|---|
DenseDepth | DenseNet-161 | 8 |
FCRN | ResNet-50 | 8 |
python train.py --model_type=densedepth --layers=161 --bs=8
Donwload the test data here and place it on the NYU_Depth_v2_DATASET_ROOT without any extraction. Then, run evaluate_pytorch.py to evaluate the performance of the network on NYU Depth v2. All the pretrained weights above must be downloaded in the checkpoints/ directory for the evaluation.
python evaluate_pytorch.py --model_type=densdepth --layers=161
Some codes are sourced and modified from DenseDepth(github), FCRN(github), and MoCo(github).