MBNSF

This repository is the official open-source implementation of the paper MBNSF - 3DV'2024 (oral):

Multi-Body Neural Scene Flow
Kavisha Vidanapathirana, Shin-Fang Ch'ng, Xueqian Li, Simon Lucey
International Conference on 3D Vision 2024 arXiv

This repository contains the code for:

Evaluation of 3D scene flow and 4D trajectory estimation (via forward-Euler integration) on the Argoverse dataset using MBNSF (Ours) and NSFP - NeurIPS'2021 (Tab. 1 and 2 in the paper).
Our implementation of NSFP++ - ECCV'2022 (Tab. 4 in the paper).
Our implementation of NTP - CVPR'2022 and MBNT (ours) for direct 4D trajectory prediction (Tab. 2 in the paper).

Method overview

MBNSF is a regularizor for incorporating multi-body rigidity into neural scene flow in a manner that (1) doesn't require explicit estimaiton of SE(3) parameters for each rigid body, (2) is robust to the inherent noise in real-world LiDAR datasets, (3) maintains the fully unsupervised nature, and (4) maintains the estimation of continuous flow fields.

Prerequisites

Environment dependencies

Set up base environment

Create conda environment with python:

conda create --name mbnsf python=3.9.4
conda activate mbnsf

Install PyTorch with suitable cudatoolkit version. See here:

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 -c pytorch
# Make sure the pytorch cuda version matches your output of 'nvcc --version'

Install Pytorch3D:

conda install -c fvcore -c iopath -c conda-forge fvcore iopath
conda install -c bottler nvidiacub
conda install pytorch3d -c pytorch3d

Install Open3D:

pip install open3d

Test installation using:

python -c "import torch ; import pytorch3d ; import open3d ; print(torch.cuda.is_available())"

Datasets

This repository provides evaluation scripts for the Argoverse dataset.

Download Argoverse test set:

Our test set for Argoverse consists of 18 sequences with 25 consecutive frames for evaluating long-term trajectories (which also results in 450 pairs for scene flow evaluation)

Download the Argoverse test set from here (~1.5 GB).
(Optional) The code for preparing this test set is provided in utils/get_gt_traj_argoverse.py.

For the Waymo test set:

For Waymo, we use the same test set as provided in FNSF.

Download the Waymo test set from here.
(Optional) Instructions for preparing this test set are provided here.

Evaluation

This section re-creates the results of Tab. 1 and Tab. 2 in our paper on the Argoverse dataset.

3D Scene Flow Prediction

In this section we optimize a scene flow field for a given pair of point clouds.

cd scene_flow_estimation/

Scene flow optimization using NSFP (baseline):

python nsfp.py --dataset_path </path/to/data>

Scene flow optimization using MBNSF (ours):

python mbnsf.py --dataset_path </path/to/data>

4D Trajectory Prediction

In this section we optmize a trajectory field for a sequence of point clouds.

cd trajectory_estimation/

Long-Tem trajectory optimization using NSFP (baseline) + Forward Euler integration:

python nsfp_fe.py --exp_name nsfp_fe_test --dataset_path </path/to/data>
python compute_traj_metrics.py --exp_name nsfp_fe_test --dataset_path </path/to/data>

Long-Tem trajectory optimization using MBNSF (ours) + Forward Euler integration:

python mbnsf_fe.py --exp_name mbnsf_fe_test --dataset_path </path/to/data>
python compute_traj_metrics.py --exp_name mbnsf_fe_test --dataset_path </path/to/data>

Long-Tem trajectory optimization using NTP (baseline, our implementation):

python ntp.py --exp_name ntp_test --dataset_path </path/to/data>
python compute_traj_metrics.py --exp_name ntp_test --dataset_path </path/to/data>

Long-Tem trajectory optimization using MBNT (ours): NTP + our regularizor:

python mbnt.py --exp_name mbnt_test --dataset_path </path/to/data>
python compute_traj_metrics.py --exp_name mbnt_test --dataset_path </path/to/data>

Citation

If you find this work useful in your research, please cite:

@inproceedings{vidanapathirana2023mbnsf,
  title={Multi-Body Neural Scene Flow},
  author={Vidanapathirana, Kavisha and Chng, Shin-Fang and Li, Xueqian and Lucey, Simon},
  booktitle={2024 International Conference on 3D Vision (3DV)},
  year={2024},
  organization={IEEE}
}

Acknowledgement

Functions from 3rd party have been acknowledged at the respective function definitions or readme files.

This project was mainly inspired by NSFP and SpectralMatching.

kavisha725/MBNSF