OneFormer3D: One Transformer for Unified Point Cloud Segmentation
News:
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
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.
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:
model.criterion.sem_criterion.loss_weight
in config file leads to better semantic metrics, and decreasing improve instance metrics.model.test_cfg.topk_insts
in config file.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
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
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
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 |
If you find this work useful for your research, please cite our paper:
@misc{kolodiazhnyi2023oneformer3d,
url = {https://arxiv.org/abs/2311.14405},
author = {Kolodiazhnyi, Maxim and Vorontsova, Anna and Konushin, Anton and Rukhovich, Danila},
title = {OneFormer3D: One Transformer for Unified Point Cloud Segmentation},
publisher = {arXiv},
year = {2023}
}