This is the PyTorch code for Self-Supervised 3D Scene Flow Estimation and Motion Prediction using Local Rigidity Prior (T-PAMI 2024). You can also check out the arXiv version at RigidFlowPP-arXiv.
In this repository, we apply RigidFlow++ for self-supervised motion prediction. For the codes in self-supervised 3D scene flow estimation, please refer to RigidFlowPP. The code is created by Ruibo Li ([email protected]).
- Python 3.7
- NVIDIA GPU + CUDA CuDNN
- PyTorch (torch == 1.9.0)
Create a conda environment for RigidPPMotion:
conda create -n RigidPPMotion python=3.7
conda activate RigidPPMotion
conda install pytorch==1.9.0 torchvision==0.10.0 cudatoolkit=10.2 -c pytorch
pip install numpy tqdm scikit-learn opencv-python matplotlib pyquaternion
Compile the furthest point sampling, grouping and gathering operation for PyTorch. We use the operation from this repo.
cd lib
python setup.py install
cd ../Install & complie supervoxel segmentation method:
cd Supervoxel_utils
g++ -std=c++11 -fPIC -shared -o main.so main.cc
cd ../More details about the supervoxel segmentation method, please refer to Supervoxel-for-3D-point-clouds.
-
Prepare the input data and the motion ground truth:
- Download the nuScenes data, and then follow MotionNet to process the training, validation, and test data.
-
Prepare raw point clouds and perform ground segmentation for self-supervision:
- Extract raw point clouds and backward BEV maps for the training samples in nuScenes:
python gen_data/gen_back_BEV_raw_point.py --root /path_to/nuScenes/nuScenes-data/ --split train --savepath /path_to/nuScenes/self-data/ - Perform ground segmentation for raw point clouds:
The data for self-supervised learning will be saved in
python gen_data/gen_ground_point.py --savepath /path_to/nuScenes/self-data//path_to/nuScenes/self-data/. Please referencesgen_data/README.mdfor more details.
- Extract raw point clouds and backward BEV maps for the training samples in nuScenes:
The trained model can be downloaded from model_nuScenes.pth.
Run the command:
python eval_SelfMotionNet.py --evaldata /path_to/nuScenes/input-data/test/ --pretrained pretrained/model_nuScenes.pth
set evaldata to the directory of the test data (e.g., /path_to/nuScenes/input-data/test/).
set pretrained to the trained model (e.g., pretrained/model_nuScenes.pth).
Run the command:
python train_SelfMotionNet.py --motiondata /path_to/nuScenes/input-data/train/ --selfdata /path_to/nuScenes/self-data/ --evaldata /path_to/nuScenes/input-data/val/
set motiondata to the directory of the input training data (e.g., /path_to/nuScenes/input-data/train/).
set selfdata to the directory of the self-supervised data (e.g., /path_to/nuScenes/self-data/).
set evaldata to the directory of the input validation or test data (e.g., /path_to/nuScenes/input-data/val/).
If you find this code useful, please cite our paper:
@article{li2024self,
title={Self-Supervised 3D Scene Flow Estimation and Motion Prediction using Local Rigidity Prior},
author={Li, Ruibo and Zhang, Chi and Wang, Zhe and Shen, Chunhua and Lin, Guosheng},
journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
year={2024},
publisher={IEEE}
}
Our project references the codes in the following repos.