Mask3D: Mask Transformer for 3D Instance Segmentation

Jonas Schult¹, Francis Engelmann^2,3, Alexander Hermans¹, Or Litany⁴, Siyu Tang³, Bastian Leibe¹

¹RWTH Aachen University ²ETH AI Center ³ETH Zurich ⁴NVIDIA

Mask3D predicts accurate 3D semantic instances achieving state-of-the-art on ScanNet, ScanNet200, S3DIS and STPLS3D.

[Project Webpage] [Paper] [Demo]

Code structure

We adapt the codebase of Mix3D which provides a highly modularized framework for 3D Semantic Segmentation based on the MinkowskiEngine.

├── mix3d
│   ├── main_instance_segmentation.py <- the main file
│   ├── conf                          <- hydra configuration files
│   ├── datasets
│   │   ├── preprocessing             <- folder with preprocessing scripts
│   │   ├── semseg.py                 <- indoor dataset
│   │   └── utils.py        
│   ├── models                        <- Mask3D modules
│   ├── trainer
│   │   ├── __init__.py
│   │   └── trainer.py                <- train loop
│   └── utils
├── data
│   ├── processed                     <- folder for preprocessed datasets
│   └── raw                           <- folder for raw datasets
├── scripts                           <- train scripts
├── docs
├── README.md
└── saved                             <- folder that stores models and logs

Dependencies 📝

The main dependencies of the project are the following:

python: 3.10.9
cuda: 11.3

You can set up a conda environment as follows

# Some users experienced issues on Ubuntu with an AMD CPU
# Install libopenblas-dev (issue #115, thanks WindWing)
# sudo apt-get install libopenblas-dev

export TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.2 7.5 8.0 8.6"

conda env create -f environment.yml

conda activate mask3d_cuda113

pip install pyyaml==5.4.1 --no-build-isolation

pip install pycocotools

pip3 install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip3 install torch-scatter -f https://data.pyg.org/whl/torch-1.12.1+cu113.html
pip3 install 'git+https://github.com/facebookresearch/detectron2.git@710e7795d0eeadf9def0e7ef957eea13532e34cf' --no-deps

mkdir third_party
cd third_party

git clone --recursive "https://github.com/NVIDIA/MinkowskiEngine"
cd MinkowskiEngine
git checkout 02fc608bea4c0549b0a7b00ca1bf15dee4a0b228
python setup.py install --force_cuda --blas=openblas

cd ..
git clone https://github.com/ScanNet/ScanNet.git
cd ScanNet/Segmentator
git checkout 3e5726500896748521a6ceb81271b0f5b2c0e7d2
make

cd ../../pointnet2
python setup.py install

cd ../../
pip3 install pytorch-lightning==1.7.2

pip install torchmetrics==0.11.4

Data preprocessing

After installing the dependencies, please copy the opmotion data to Mask3D/data/processed/opmotion, the files should be organized as follows:

Mask3D # project root
├── data
│   ├── processed
│   │   ├── opmotion
│   │   │   ├── opmotion.h5
│   │   │   ├── data.json

Training and testing

Train Mask3D on the opmotion dataset:

wandb login

# train with base part
python main_instance_segmentation.py data/datasets=opmotion model.num_queries=10 data.batch_size=32

# train without base part
python main_instance_segmentation.py data/datasets=opmotion model.num_queries=10 data.batch_size=32 data.train_dataset.num_labels=3 data.validation_dataset.num_labels=3 data.test_dataset.num_labels=3 general.num_targets=4

Test Mask3D on the opmotion dataset, the inference results will be saved out:

python main_instance_segmentation.py data/datasets=opmotion general.train_mode=false model.num_queries=10 general.checkpoint={path_to_ckpt_file} data.batch_size=32

BibTeX

@article{Schult23ICRA,
  title     = {{Mask3D: Mask Transformer for 3D Semantic Instance Segmentation}},
  author    = {Schult, Jonas and Engelmann, Francis and Hermans, Alexander and Litany, Or and Tang, Siyu and Leibe, Bastian},
  booktitle = {{International Conference on Robotics and Automation (ICRA)}},
  year      = {2023}
}