Skip to content

[CVPR2024] OneFormer3D: One Transformer for Unified Point Cloud Segmentation

License

Notifications You must be signed in to change notification settings

filaPro/oneformer3d

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

OneFormer3D: One Transformer for Unified Point Cloud Segmentation

News:

  • September, 2024. We release state-of-the-art 3D object detector UniDet3D based on OneFormer3D.
  • 🔥 February, 2024. OneFormer3D is now accepted at CVPR 2024.
  • 🔥 November, 2023. OneFormer3D achieves state-of-the-art in
    • 3D instance segmentation on ScanNet (hidden test) PWC
      leaderboard screenshot ScanNet leaderboard
    • 3D instance segmentation on S3DIS (6-Fold) PWC
    • 3D panoptic segmentation on ScanNet PWC
    • 3D object detection on ScanNet (w/o TTA) PWC
    • 3D semantic segmentation on ScanNet (val, w/o extra training data) PWC

This repository contains an implementation of OneFormer3D, a 3D (instance, semantic, and panoptic) segmentation method introduced in our paper:

OneFormer3D: One Transformer for Unified Point Cloud Segmentation
Maksim Kolodiazhnyi, Anna Vorontsova, Anton Konushin, Danila Rukhovich
Samsung Research
https://arxiv.org/abs/2311.14405

Installation

For convenience, we provide a Dockerfile. This implementation is based on mmdetection3d framework v1.1.0. If installing without docker please follow their getting_started.md.

Getting Started

Please see test_train.md for basic usage examples. For ScanNet and ScanNet200 datasets preprocessing please follow our instruction. It differs from original mmdetection3d only by adding superpoint clustering. For S3DIS preprocessing we follow original instruction from mmdetection3d. We also support Structured3D dataset for pre-training.

Important notes:

  • The metrics from our paper can be achieved in several ways, we just choose the most stable one for each dataset in this repository.
  • If you are interested in only one of three segmentation tasks, it is possible to achieve slightly better metrics, than declared in our paper. Specifically, increasing model.criterion.sem_criterion.loss_weight in config file leads to better semantic metrics, and decreasing improve instance metrics.
  • All models can be trained with a single GPU with 32 Gb memory (or even 24 Gb for ScanNet dataset). If you face issues with RAM during instance segmentation evaluation at validation or test stages feel free to decrease model.test_cfg.topk_insts in config file.
  • Due to the bug in SpConv we reshape backbone weights between train and test stages.

ScanNet

For ScanNet we present the model with SpConv backbone, superpoint pooling, selecting all queries, and predicting semantics directly from instance queries. Backbone is initialized from SSTNet checkpoint. It should be downloaded and put to work_dirs/tmp before training.

# train (with validation)
python tools/train.py configs/oneformer3d_1xb4_scannet.py
# test
python tools/fix_spconv_checkpoint.py \
    --in-path work_dirs/oneformer3d_1xb4_scannet/epoch_512.pth \
    --out-path work_dirs/oneformer3d_1xb4_scannet/epoch_512.pth
python tools/test.py configs/oneformer3d_1xb4_scannet.py \
    work_dirs/oneformer3d_1xb4_scannet/epoch_512.pth

ScanNet200

For ScanNet200 we present the model with MinkowskiEngine backbone, superpoint pooling, selecting all queries, and predicting semantics directly from instance queries. Backbone is initialized from Mask3D checkpoint. It should be downloaded and put to work_dirs/tmp before training.

# train (with validation)
python tools/train.py configs/oneformer3d_1xb4_scannet200.py
# test
python tools/test.py configs/oneformer3d_1xb4_scannet200.py \
    work_dirs/oneformer3d_1xb4_scannet/epoch_512.pth

S3DIS

For S3DIS we present the model with SpConv backbone, w/o superpoint pooling, w/o query selection, and with separate semantic queries. Backbone is pretrained on Structured3D and ScanNet. It can be downloaded and put to work_dirs/tmp before training or trained with our code. We train the model on Areas 1, 2, 3, 4, 6 and test on Area 5. To change this split feel free to modify train_area and test_area parameters in config.

# pre-train
python tools/train.py configs/instance-only-oneformer3d_1xb2_scannet-and-structured3d.py
python tools/fix_spconv_checkpoint.py \
    --in-path work_dirs/instance-only-oneformer3d_1xb2_scannet-and-structured3d/iter_600000.pth \
    --out-path work_dirs/tmp/instance-only-oneformer3d_1xb2_scannet-and-structured3d.pth
# train (with validation)
python tools/train.py configs/oneformer3d_1xb2_s3dis-area-5.py
# test
python tools/fix_spconv_checkpoint.py \
    --in-path work_dirs/oneformer3d_1xb2_s3dis-area-5/epoch_512.pth \
    --out-path work_dirs/oneformer3d_1xb2_s3dis-area-5/epoch_512.pth
python tools/test.py configs/oneformer3d_1xb2_s3dis-area-5.py \
    work_dirs/oneformer3d_1xb2_s3dis-area-5/epoch_512.pth

Models

Metric values in the table are given for the provided checkpoints and may vary a little from the ones in our paper. Due to randomness it may be needed to run training with the same config for several times to achieve the best metrics.

Dataset mAP25 mAP50 mAP mIoU PQ Download
ScanNet 86.7 78.8 59.3 76.4 70.7 model | log | config
ScanNet200 44.6 40.9 30.2 29.4 29.7 model | log | config
S3DIS 80.6 72.7 58.0 71.9 64.6 model | log | config

Example Predictions

ScanNet predictions

Citation

If you find this work useful for your research, please cite our paper:

@inproceedings{kolodiazhnyi2024oneformer3d,
  title={Oneformer3d: One transformer for unified point cloud segmentation},
  author={Kolodiazhnyi, Maxim and Vorontsova, Anna and Konushin, Anton and Rukhovich, Danila},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={20943--20953},
  year={2024}
}