Weakly Supervised Learning of Rigid 3D Scene Flow
This repository provides code and data to train and evaluate a weakly supervised method for rigid 3D scene flow estimation. It represents the official implementation of the paper:
Weakly Supervised Learning of Rigid 3D Scene Flow
Zan Gojcic, Or Litany, Andreas Wieser, Leonidas J. Guibas, Tolga Birdal
| IGP ETH Zurich | Nvidia Toronto AI Lab | Guibas Lab Stanford University |
For more information, please see the project webpage
Environment Setup
Note: the code in this repo has been tested on Ubuntu 16.04/20.04 with Python 3.7, CUDA 10.1/10.2, PyTorch 1.7.1 and MinkowskiEngine 0.5.1. It may work for other setups, but has not been tested.
Before proceding, make sure CUDA is installed and set up correctly.
After cloning this reposiory you can proceed by setting up and activating a virual environment with Python 3.7. If you are using a different version of cuda (10.1) change the pytorch installation instruction accordingly.
export CXX=g++-7
conda config --append channels conda-forge
conda create --name rigid_3dsf python=3.7
source activate rigid_3dsf
conda install --file requirements.txt
conda install -c open3d-admin open3d=0.9.0.0
conda install -c intel scikit-learn
conda install pytorch==1.7.1 torchvision cudatoolkit=10.1 -c pytorch
You can then proceed and install MinkowskiEngine library for sparse tensors:
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps
Our repository also includes a pytorch implementation of Chamfer Distance in ./utils/chamfer_distance
which will be compiled on the first run.
In order to test if Pytorch and MinkwoskiEngine are installed correctly please run
python -c "import torch, MinkowskiEngine"
which should run without an error message.
Data
We provide the preprocessed data of flying_things_3d (108GB), stereo_kitti (500MB), lidar_kitti (~160MB), semantic_kitti (78GB), and waymo_open (50GB) used for training and evaluating our model.
To download a single dataset please run:
bash ./scripts/download_data.sh name_of_the_dataset
To download all datasets simply run:
bash ./scripts/download_data.sh
The data will be downloaded and extracted to ./data/name_of_the_dataset/
.
Pretrained models
We provide the checkpoints of the models trained on flying_things_3d or semantic_kitti, which we use in our main evaluations.
To download these models please run:
bash ./scripts/download_pretrained_models.sh
Additionally, we provide all the models used in the ablation studies and the model fine tuned on waymo_open.
To download these models please run:
bash ./scripts/download_pretrained_models_ablations.sh
All the models will be downloaded and extracted to ./logs/dataset_used_for_training/
.
Evaluation with pretrained models
Our method with pretrained weights can be evaluated using the ./eval.py
script. The configuration parameters of the evaluation can be set with the *.yaml
configuration files located in ./configs/eval/
. We provide a configuration file for each dataset used in our paper. For all evaluations please first download the pretrained weights and the corresponding data. Note, if the data or pretrained models are saved to a non-default path the config files also has to be adapted accordingly.
FlyingThings3D
To evaluate our backbone + scene flow head on FlyingThings3d please run:
python eval.py ./configs/eval/eval_flying_things_3d.yaml
This should recreate the results from the Table 1 of our paper (EPE3D: 0.052 m).
stereoKITTI
To evaluate our backbone + scene flow head on stereoKITTI please run:
python eval.py ./configs/eval/eval_stereo_kitti.yaml
This should again recreate the results from the Table 1 of our paper (EPE3D: 0.042 m).
lidarKITTI
To evaluate our full weakly supervised method on lidarKITTI please run:
python eval.py ./configs/eval/eval_lidar_kitti.yaml
This should recreate the results for Ours++ on lidarKITTI (w/o ground) from the Table 2 of our paper (EPE3D: 0.094 m). To recreate other results on lidarKITTI please change the ./configs/eval/eval_lidar_kitti.yaml
file accordingly.
semanticKITTI
To evaluate our full weakly supervised method on semanticKITTI please run:
python eval.py ./configs/eval/eval_semantic_kitti.yaml
This should recreate the results of our full model on semanticKITTI (w/o ground) from the Table 4 of our paper. To recreate other results on semanticKITTI please change the ./configs/eval/eval_semantic_kitti.yaml
file accordingly.
waymo open
To evaluate our fine-tuned model on waymo open please run:
python eval.py ./configs/eval/eval_waymo_open.yaml
This should recreate the results for Ours++ (fine-tuned) from the Table 9 of the appendix. To recreate other results on waymo open please change the ./configs/eval/eval_waymo_open.yaml
file accordingly.
Training our method from scratch
Our method can be trained using the ./train.py
script. The configuration parameters of the training process can be set using the config files located in ./configs/train/
.
Training our backbone with full supervision on FlyingThings3D
To train our backbone network and scene flow head under full supervision (corresponds to Sec. 4.3 of our paper) please run:
python train.py ./configs/train/train_fully_supervised.yaml
The checkpoints and tensorboard data will be saved to ./logs/logs_FlyingThings3D_ME
. If you run out of GPU memory with the default setting please adapt the batch_size
and acc_iter_size
in the ./configs/default.yaml
to e.g. 4 and 2, respectively.
Training under weak supervision on semanticKITTI
To train our full method under weak supervision on semanticKITTI please run
python train.py ./configs/train/train_weakly_supervised.yaml
The checkpoints and tensorboard data will be saved to ./logs/logs_SemanticKITTI_ME
. If you run out of GPU memory with the default setting please adapt the batch_size
and acc_iter_size
in the ./configs/default.yaml
to e.g. 4 and 2, respectively.
Citation
If you found this code or paper useful, please consider citing:
@misc{gojcic2021weakly3dsf,
title = {Weakly {S}upervised {L}earning of {R}igid {3D} {S}cene {F}low},
author = {Gojcic, Zan and Litany, Or and Wieser, Andreas and Guibas, Leonidas J and Birdal, Tolga},
year = {2021},
eprint={2102.08945},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
Contact
If you run into any problems or have questions, please create an issue or contact Zan Gojcic.
Acknowledgments
In this project we use parts of the official implementations of:
We thank the respective authors for open sourcing their methods.