Awesome

OGC: Unsupervised 3D Object Segmentation from Rigid Dynamics of Point Clouds (NeurIPS 2022)

!Updates:

New contents in the extension of OGC have been included in this repo:

Support for Waymo Open dataset.
A multi-frame co-segmentation algorithm to leverage multi-frame inputs in testing.
Support for SparseConv & Point Transformer backbones in the object segmentation network.

Overview

We propose the first unsupervised 3D object segmentation method, learning from dynamic motion patterns in point cloud sequences.

Our method demonstrates promising results on various scenarios:

Object part instance segmentation

Object segmentation in indoor scenes

Object segmentation in outdoor scenes

Full demo (Youtube)

1. Environment

Please first install a GPU-supported pytorch version which fits your machine. We have tested with pytorch 1.9.0.

Install PointNet2 CPP lib:

cd pointnet2
python setup.py install
cd ..

Install other dependencies:

pip install -r requirements

(Optional) Install Open3D for the visualization of point cloud segmentation:

pip install open3d

2. Data preparation

(1) SAPIEN

Please download from links provided by MultibodySync:

Train+Val (mbs-shapepart): Google Drive
Test (mbs-sapien): Google Drive

Then put them into your ${SAPIEN} path.

(2) OGC-DR (Dynamic Room) & OGC-DRSV (Single-View Dynamic Room)

Please download the complete datasets from links below:

OGC-DR: Dropbox
OGC-DRSV: Dropbox

Alternatively, you can generate the dataset by yourself.

OGC-DR: Please first download the ShapeNet Core v1. Select the archives according to object categories specified in data_prepare/ogcdr/meta.yaml and unzip them into your ${OGC_DR}/ShapeNet_mesh path. Then run the following script to generate the dataset.

python data_prepare/ogcdr/build_ogcdr.py ${OGC_DR}

OGC-DRSV: Run the following script to collect single depth scans on OGC-DR mesh models and generate incomplete point cloud dataset OGC-DRSV.

python data_prepare/ogcdrsv/build_ogcdrsv.py --src_root ${OGC_DR} --dest_root ${OGC_DRSV}

Collect groundtruth segmentation for OGC-DRSV and downsample the point clouds:

python data_prepare/ogcdrsv/collect_segm.py --src_root ${OGC_DR} --dest_root ${OGC_DRSV}

(3) KITTI-SF (Scene Flow)

Please first download:

Merge the training folder of them in your ${KITTI_SF} path. Then run the following script to unproject disparity, optical flow, 2D segmentation into point cloud, scene flow, 3D segmentation:

python data_prepare/kittisf/process_kittisf.py ${KITTI_SF}

Finally, downsample all point clouds to 8192-point:

python data_prepare/kittisf/downsample_kittisf.py ${KITTI_SF} --save_root ${KITTI_SF}_downsampled
# After extracting flow estimations in the following, come back here to downsample flow estimations
python data_prepare/kittisf/downsample_kittisf.py ${KITTI_SF} --save_root ${KITTI_SF}_downsampled --predflow_path flowstep3d

${KITTI_SF}_downsampled will be the path for the downsampled dataset.

(4) KITTI-Det (Detection)

Please first download the following items from KITTI 3D Object Detection Evaluation 2017:

Merge the training folder of them in your ${KITTI_DET} path. Then run the following script to extract 8192-point front-view point cloud, and obtain segmentation from bounding box annotations.

python data_prepare/kittidet/process_kittidet.py ${KITTI_DET}

(5) SemanticKITTI

Please first download the following iterms from SemanticKITTI:

Merge the velodyne, labels and calib.txt of each sequence. The organized dataset should be as follows:

SemanticKITTI
└── sequences
    └── 00
    │   ├── velodyne
    │   ├── labels
    │   └── calib.txt
    └── 01
    ...

Then run the following script to extract 8192-point front-view point cloud, and obtain segmentation from panoptic annotations.

python data_prepare/semantickitti/process_semantickitti.py ${SEMANTIC_KITTI}

(6) Waymo Open

Please first register at the Waymo Open Dataset, then download the raw data.

Process raw data into point clouds. Please carefully follow the instructions in OpenPCDet. You need to install waymo-open-dataset and OpenPCDet here.
Extract groundtruth scene flow. Please follow the instructions in DCA-SRSFE.
Generate groundtruth segmentations from bounding box annotations:

python data_prepare/waymo/process_waymo.py --data_root $WAYMO_RAW --save_root $WAYMO --split $SPLIT

Downsample all point clouds to 8192-point:

python data_prepare/waymo/downsample_waymo.py --data_root $WAYMO --save_root ${WAYMO}_downsampled --split $SPLIT
# After extracting flow estimations in the following, come back here to downsample flow estimations
python data_prepare/waymo/downsample_waymo.py --data_root $WAYMO --save_root ${WAYMO}_downsampled --split $SPLIT --predflow_path flowstep3d

3. Pre-trained models

You can download all our pre-trained models from Dropbox (including self-supervised scene flow networks, and unsupervised/supervised segmentation networks) and extract them to ./ckpt.

4. Scene flow estimation

Train

Train the self-supervised scene flow networks:

# SAPIEN 
python train_flow.py config/flow/sapien/sapien_unsup.yaml
# OGC-DR 
python train_flow.py config/flow/ogcdr/ogcdr_unsup.yaml
# OGC-DRSV 
python train_flow.py config/flow/ogcdrsv/ogcdrsv_unsup.yaml

For KITTI-SF dataset, we directly employ the pre-trained model released by FlowStep3D.

Test

Evaluate and save the scene flow estimations.

# SAPIEN 
python test_flow.py config/flow/sapien/sapien_unsup.yaml --split ${SPLIT} --save
# OGC-DR 
python test_flow.py config/flow/ogcdr/ogcdr_unsup.yaml --split ${SPLIT} --test_batch_size 12 --test_model_iters 5 --save
# OGC-DRSV 
python test_flow.py config/flow/ogcdrsv/ogcdrsv_unsup.yaml --split ${SPLIT} --test_batch_size 12 --test_model_iters 5 --save
# KITTI-SF 
python test_flow_kittisf.py config/flow/kittisf/kittisf_unsup.yaml --split ${SPLIT} --test_model_iters 5 --save
# Waymo Open
python test_flow_waymo.py config/flow/waymo/kittisf_unsup.yaml --split ${SPLIT} --use_odometry --bound --test_model_iters 5 --save

${SPLIT} can be train/val/test for SAPIEN & OGC-DR/OGC-DRSV, train/val for KITTI-SF & Waymo Open.

5. Unsupervised segmentation

Train

Alternate the segmentation network training and scene flow improvement for R rounds. In each ${ROUND} (starting from 1):

# SAPIEN: first R-1 rounds
python train_seg.py config/seg/sapien/sapien_unsup_woinv.yaml --round ${ROUND}
python oa_icp.py config/seg/sapien/sapien_unsup_woinv.yaml --split ${SPLIT} --round ${ROUND} --save
# SAPIEN: the last round
python train_seg.py config/seg/sapien/sapien_unsup.yaml --round ${ROUND}

# OGC-DR: first R-1 rounds
python train_seg.py config/seg/ogcdr/ogcdr_unsup_woinv.yaml --round ${ROUND}
python oa_icp.py config/seg/ogcdr/ogcdr_unsup_woinv.yaml --split ${SPLIT} --round ${ROUND} --test_batch_size 24 --save
# OGC-DR: the last round
python train_seg.py config/seg/ogcdr/ogcdr_unsup.yaml --round ${ROUND}

# OGC-DRSV: first R-1 rounds
python train_seg.py config/seg/ogcdrsv/ogcdrsv_unsup_woinv.yaml --round ${ROUND}
python oa_icp.py config/seg/ogcdrsv/ogcdrsv_unsup_woinv.yaml --split ${SPLIT} --round ${ROUND} --test_batch_size 24 --save
# OGC-DRSV: the last round
python train_seg.py config/seg/ogcdrsv/ogcdrsv_unsup.yaml --round ${ROUND}

# KITTI-SF: first R-1 rounds
python train_seg.py config/seg/kittisf/kittisf_unsup_woinv.yaml --round ${ROUND}
python oa_icp.py config/seg/kittisf/kittisf_unsup_woinv.yaml --split ${SPLIT} --round ${ROUND} --test_batch_size 4 --save
# KITTI-SF: the last round
python train_seg.py config/seg/kittisf/kittisf_unsup.yaml --round ${ROUND}

# Waymo Open: only 1 round
python train_seg_waymo.py config/seg/waymo/waymo_unsup.yaml --round 1

When performing scene flow improvement, ${SPLIT} needs to traverse train/val/test for SAPIEN & OGC-DR/OGC-DRSV, train/val for KITTI-SF.

Test

# SAPIEN 
python test_seg.py config/seg/sapien/sapien_unsup.yaml --split test --round ${ROUND}
# OGC-DR 
python test_seg.py config/seg/ogcdr/ogcdr_unsup.yaml --split test --round ${ROUND} --test_batch_size 16
# OGC-DRSV 
python test_seg.py config/seg/ogcdrsv/ogcdrsv_unsup.yaml --split test --round ${ROUND} --test_batch_size 16
# KITTI-SF 
python test_seg.py config/seg/kittisf/kittisf_unsup.yaml --split val --round ${ROUND} --test_batch_size 8
# KITTI-Det 
python test_seg.py config/seg/kittidet/kittisf_unsup.yaml --split val --round ${ROUND} --test_batch_size 8
# SemanticKITTI 
python test_seg.py config/seg/semantickitti/kittisf_unsup.yaml --round ${ROUND} --test_batch_size 8
# Waymo Open
python test_seg_waymo.py config/seg/waymo/waymo_unsup.yaml --split val --round 1 --test_batch_size 8

${ROUND} can be 1/2/3/..., and we take 2 rounds as default in our experiments. Specify --save to save the estimations. Specify --visualize for qualitative evaluation mode.

Test the scene flow improvement

Your can follow the evaluation settings of FlowStep3D to test the improved flow, and see how our method push the boundaries of unsupervised scene flow estimation:

# Refine the scene flow estimations
python oa_icp.py config/seg/kittisf/kittisf_unsup.yaml --split train --round 2 --test_batch_size 4 --save --saveflow_path flowstep3d_for-benchmark
python oa_icp.py config/seg/kittisf/kittisf_unsup.yaml --split val --round 2 --test_batch_size 4 --save --saveflow_path flowstep3d_for-benchmark
# Evaluate
python test_flow_kittisf_benchmark.py config/flow/kittisf/kittisf_unsup.yaml

6. Supervised segmentation

You can train the segmentation network with full annotations.

Train

# SAPIEN 
python train_seg_sup.py config/seg/sapien/sapien_sup.yaml
# OGC-DR
python train_seg_sup.py config/seg/ogcdr/ogcdr_sup.yaml
# OGC-DRSV
python train_seg_sup.py config/seg/ogcdrsv/ogcdrsv_sup.yaml
# KITTI-SF 
python train_seg_sup.py config/seg/kittisf/kittisf_sup.yaml
# KITTI-Det 
python train_seg_sup.py config/seg/kittidet/kittidet_sup.yaml
# Waymo Open
python train_seg_waymo_sup.py config/seg/waymo/waymo_sup.yaml

Test

# SAPIEN 
python test_seg.py config/seg/sapien/sapien_sup.yaml --split test
# OGC-DR
python test_seg.py config/seg/ogcdr/ogcdr_sup.yaml --split test --test_batch_size 16
# OGC-DRSV
python test_seg.py config/seg/ogcdrsv/ogcdrsv_sup.yaml --split test --test_batch_size 16
# KITTI-SF 
python test_seg.py config/seg/kittisf/kittisf_sup.yaml --split val --test_batch_size 8
# KITTI-Det 
python test_seg.py config/seg/kittidet/kittisf_sup.yaml --split val --test_batch_size 8
python test_seg.py config/seg/kittidet/kittidet_sup.yaml --split val --test_batch_size 8
# SemanticKITTI 
python test_seg.py config/seg/semantickitti/kittisf_sup.yaml --test_batch_size 8
# Waymo Open
python test_seg_waymo.py config/seg/waymo/waymo_sup.yaml --split val --test_batch_size 8

7. Multi-frame co-segmentation

# SAPIEN dataset
python vote.py config/seg/sapien/sapien_unsup.yaml --split test --round $ROUND
# OGC-DR dataset
python vote.py config/seg/ogcdr/ogcdr_unsup.yaml --split test --round $ROUND --test_batch_size 16
# OGC-DRSV dataset
python vote.py config/seg/ogcdrsv/ogcdrsv_unsup.yaml --split test --round $ROUND --test_batch_size 16

Use --time_window_size to control the frame range used for multi-frame co-segmentation.

Citation

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

@article{song2022,
  title={{OGC: Unsupervised 3D Object Segmentation from Rigid Dynamics of Point Clouds}},
  author={Song, Ziyang and Yang, Bo},
  journal={NeurIPS},
  year={2022}
}

Acknowledgements

Some code is borrowed from: