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:
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
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 pytorchYou can then proceed and install MinkowskiEngine library for sparse tensors:
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-depsOur 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.
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_datasetTo download all datasets simply run:
bash ./scripts/download_data.shThe data will be downloaded and extracted to ./data/name_of_the_dataset/.
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.shAdditionally, 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.shAll the models will be downloaded and extracted to ./logs/dataset_used_for_training/.
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.
To evaluate our backbone + scene flow head on FlyingThings3d please run:
python eval.py ./configs/eval/eval_flying_things_3d.yamlThis should recreate the results from the Table 1 of our paper (EPE3D: 0.052 m).
To evaluate our backbone + scene flow head on stereoKITTI please run:
python eval.py ./configs/eval/eval_stereo_kitti.yamlThis should again recreate the results from the Table 1 of our paper (EPE3D: 0.042 m).
To evaluate our full weakly supervised method on lidarKITTI please run:
python eval.py ./configs/eval/eval_lidar_kitti.yamlThis 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.
To evaluate our full weakly supervised method on semanticKITTI please run:
python eval.py ./configs/eval/eval_semantic_kitti.yamlThis 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.
To evaluate our fine-tuned model on waymo open please run:
python eval.py ./configs/eval/eval_waymo_open.yamlThis 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.
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/.
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.yamlThe 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.
To train our full method under weak supervision on semanticKITTI please run
python train.py ./configs/train/train_weakly_supervised.yamlThe 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.
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}
}If you run into any problems or have questions, please create an issue or contact Zan Gojcic.
In this project we use parts of the official implementations of:
We thank the respective authors for open sourcing their methods.