PV-RAFT

This repository contains the PyTorch implementation for paper "PV-RAFT: Point-Voxel Correlation Fields for Scene Flow Estimation of Point Clouds" (CVPR 2021)[arXiv]

Installation

Prerequisites

• Python 3.8
• PyTorch 1.8
• torch-scatter
• CUDA 10.2
• RTX 2080 Ti
• tqdm, tensorboard, scipy, imageio, png

conda create -n pvraft python=3.8
conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
conda install tqdm tensorboard scipy imageio
pip install pypng
pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu102.html

Usage

Data Preparation

We follow HPLFlowNet to prepare FlyingThings3D and KITTI datasets. Please refer to repo. Make sure the project structure look like this:

RAFT_SceneFlow/
    data/
        FlyingThings3D_subset_processed_35m/
        kitti_processed/
    data_preprocess/
    datasets/
    experiments/
    model/
    modules/
    tools/

After downloading datasets, we need to preprocess them.

FlyingThings3D Dataset

python process_flyingthings3d_subset.py --raw_data_path=path_src/FlyingThings3D_subset --save_path=path_dst/FlyingThings3D_subset_processed_35m --only_save_near_pts

You should replace raw_data_path and save_path with your own setting.

KITTI Dataset

python process_kitti.py --raw_data_path=path_src/kitti --save_path=path_dst/kitti_processed --calib_path=calib_folder_path

You should replace raw_data_path, save_path and calib_path with your own setting.

Train

python train.py --exp_path=pv_raft --batch_size=2 --gpus=0,1 --num_epochs=20 --max_points=8192 --iters=8  --root=./

where exp_path is the experiment folder name and root is the project root path. These 20 epochs take about 53 hours on two RTX 2080 Ti.

If you want to train the refine model, please add --refine and specify --weights parameter as the directory name of the pre-trained model. For example,

python train.py --refine --exp_path=pv_raft_finetune --batch_size=2 --gpus=0,1 --num_epochs=10 --max_points=8192 --iters=32 --root=./ --weights=./experiments/pv_raft/checkpoints/best_checkpoint.params

These 10 epochs take about 38 hours on two RTX 2080 Ti.

Test

python test.py --dataset=KITTI --exp_path=pv_raft --gpus=1 --max_points=8192 --iters=8 --root=./ --weights=./experiments/pv_raft/checkpoints/best_checkpoint.params

where dataset should be chosen from FT3D/KITTI, and weights is the absolute path of checkpoint file.

If you want to test the refine model, please add --refine. For example,

python test.py --refine --dataset=KITTI --exp_path=pv_raft_finetune --gpus=1 --max_points=8192 --iters=32 --root=./ --weights=./experiments/pv_raft_finetune/checkpoints/best_checkpoint.params

Reproduce results

You can download the checkpoint of refined model here.

Acknowledgement

Our code is based on FLOT. We also refer to RAFT and HPLFlowNet.

Citation

If you find our work useful in your research, please consider citing:

@inproceedings{wei2020pv,
  title={{PV-RAFT: Point-Voxel Correlation Fields for Scene Flow Estimation of Point Clouds}},
  author={Wei, Yi and Wang, Ziyi and Rao, Yongming and Lu, Jiwen and Zhou, Jie},
  booktitle={CVPR},
  year={2021}
}