TD3D: Top-Down Beats Bottom-Up in 3D Instance Segmentation

News:

• 🔥 February 6, 2023. We achieved SOTA results on the ScanNet test subset (mAP@25).
• 🔥 February 2023. The source code has been published.

This repository contains an implementation of TD3D, a 3D instance segmentation method introduced in our paper:

Top-Down Beats Bottom-Up in 3D Instance Segmentation
Maksim Kolodiazhnyi, Danila Rukhovich, Anna Vorontsova, Anton Konushin
Samsung Research
https://arxiv.org/abs/2302.02871

Installation

For convenience, we provide a Dockerfile.

Alternatively, you can install all required packages manually. This implementation is based on mmdetection3d framework.

Please refer to the original installation guide getting_started.md, including MinkowskiEngine installation, replacing open-mmlab/mmdetection3d with samsunglabs/td3d.

Most of the TD3D-related code locates in the following files: detectors/td3d_instance_segmentor.py, necks/ngfc_neck.py, decode_heads/td3d_instance_head.py.

Getting Started

Please see getting_started.md for basic usage examples. We follow the mmdetection3d data preparation protocol described in s3dis for S3DIS and in scannet for ScanNet and ScanNet200.

Training

To start training, run train with TD3D configs. To avoid gpu memory problems during validation callback, set score_thr to 0.15 and nms_pre to 100 in configs before training (then return them to their original values during testing):

python tools/train.py configs/td3d_is/td3d_is_scannet-3d-18class.py

For training on S3DIS with pretrained on ScanNet weights, download ScanNet model and put it into your working directory. Then use configs/td3d_is/td3d_is_s3dis-3d-5class_pretrain.py according to the previous instructions.

Testing

Test pre-trained model using test with TD3D configs:

python tools/test.py configs/td3d_is/td3d_is_scannet-3d-18class.py \
    work_dirs/td3d_is_scannet-3d-18class/latest.pth --eval mAP

Visualization

Visualizations can be created with test script. For better visualizations, you may set score_thr to 0.20 and nms_pre to 200 in configs:

python tools/test.py configs/td3d_is/td3d_is_scannet-3d-18class.py \
    work_dirs/td3d_is_scannet-3d-18class/latest.pth --eval mAP --show \
    --show-dir work_dirs/td3d_is_scannet-3d-18class

Models (quality on validation subset)

Dataset	mAP@0.25	mAP@0.5	mAP	Download
ScanNet	81.9	71.2	47.3	model | config
S3DIS (5 area)	73.8	65.1	48.6	model | config
S3DIS (5 area) (ScanNet pretrain)	75.0	67.2	52.1	model | config
ScanNet200	40.4	34.8	23.1	model | config
STPLS3D	74.0	69.6	54.5	model | config

Examples

Citation

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

@misc{kolodiazhnyi2023topdown,
  doi = {10.48550/ARXIV.2302.02871},
  url = {https://arxiv.org/abs/2302.02871},
  author = {Kolodiazhnyi, Maksim and Rukhovich, Danila and Vorontsova, Anna and Konushin, Anton},
  title = {Top-Down Beats Bottom-Up in 3D Instance Segmentation},
  publisher = {arXiv},
  year = {2023}
}