Mask3D:用于3D实例分割的掩码变换器
Jonas Schult1, Francis Engelmann2,3, Alexander Hermans1, Or Litany4, Siyu Tang3, Bastian Leibe1
1亚琛工业大学 2苏黎世联邦理工学院人工智能中心 3苏黎世联邦理工学院 4NVIDIA
Mask3D预测精确的3D语义实例,在ScanNet、ScanNet200、S3DIS和STPLS3D数据集上达到了最先进水平。
新闻
- 2023年10月29日:查看Mask3D的简易设置。
- 2023年1月17日:Mask3D被ICRA 2023接收。🔥
- 2022年10月14日:添加了STPLS3D支持。
- 2022年10月10日:Mask3D在ECCV 2022 Urban3D Workshop举办的STPLS3D挑战赛中排名第二。
- 2022年10月6日:Mask3D预印本在arXiv上发布。
- 2022年9月25日:代码发布。
代码结构
我们采用了Mix3D的代码库,该代码库基于MinkowskiEngine为3D语义分割提供了高度模块化的框架。
├── mix3d
│ ├── main_instance_segmentation.py <- 主文件
│ ├── conf <- hydra配置文件
│ ├── datasets
│ │ ├── preprocessing <- 预处理脚本文件夹
│ │ ├── semseg.py <- 室内数据集
│ │ └── utils.py
│ ├── models <- Mask3D模块
│ ├── trainer
│ │ ├── __init__.py
│ │ └── trainer.py <- 训练循环
│ └── utils
├── data
│ ├── processed <- 预处理数据集文件夹
│ └── raw <- 原始数据集文件夹
├── scripts <- 训练脚本
├── docs
├── README.md
└── saved <- 存储模型和日志的文件夹
依赖 :memo:
项目的主要依赖如下:
python: 3.10.9
cuda: 11.3
您可以按如下方式设置conda环境
# 一些用户在使用AMD CPU的Ubuntu系统上遇到了问题
# 安装libopenblas-dev(问题#115,感谢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
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
数据预处理 :hammer:
安装完依赖后,我们对数据集进行预处理。
ScanNet / ScanNet200
首先,我们使用默认参数对测试场景应用Felzenswalb和Huttenlocher的基于图的图像分割算法。
详细信息请参考原始仓库。
将生成的分割结果放在./data/raw/scannet_test_segments
目录下。
python -m datasets.preprocessing.scannet_preprocessing preprocess \
--data_dir="原始SCANNET数据集路径" \
--save_dir="data/processed/scannet" \
--git_repo="SCANNET_GIT仓库路径" \
--scannet200=false/true
S3DIS
S3DIS数据集包含一些小错误,我们最初是手动修复的。我们将很快发布一个预处理脚本,可以直接预处理原始数据集。目前,请按照这里的说明手动修复数据集。之后,按如下方式调用预处理脚本:
python -m datasets.preprocessing.s3dis_preprocessing preprocess \
--data_dir="Stanford3dDataset_v1.2的路径" \
--save_dir="data/processed/s3dis"
STPLS3D
python -m datasets.preprocessing.stpls3d_preprocessing preprocess \
--data_dir="STPLS3D路径" \
--save_dir="data/processed/stpls3d"
训练和测试 :train2:
在ScanNet数据集上训练Mask3D:
python main_instance_segmentation.py
请参考配置脚本(例如这里)以获取如何复现我们结果的详细说明。 最简单的推理命令如下:
python main_instance_segmentation.py \
general.checkpoint='检查点路径.ckpt' \
general.train_mode=false
训练好的检查点 :floppy_disk:
我们提供了详细的得分和网络配置以及训练好的检查点。
S3DIS(在ScanNet train+val上预训练)
按照PointGroup、HAIS和SoftGroup的做法,我们对在ScanNet上预训练的模型进行微调(配置和检查点)。
数据集 | AP | AP_50 | AP_25 | 配置 | 检查点 :floppy_disk: | 得分 :chart_with_upwards_trend: | 可视化 :telescope: |
---|---|---|---|---|---|---|---|
Area 1 | 69.3 | 81.9 | 87.7 | 配置 | 检查点 | 得分 | 可视化 |
Area 2 | 44.0 | 59.5 | 66.5 | 配置 | 检查点 | 得分 | 可视化 |
Area 3 | 73.4 | 83.2 | 88.2 | 配置 | 检查点 | 得分 | 可视化 |
Area 4 | 58.0 | 69.5 | 74.9 | 配置 | 检查点 | 得分 | 可视化 |
Area 5 | 57.8 | 71.9 | 77.2 | 配置 | 检查点 | 得分 | 可视化 |
Area 6 | 68.4 | 79.9 | 85.2 | 配置 | 检查点 | 得分 | 可视化 |
S3DIS(从头训练)
数据集 | AP | AP_50 | AP_25 | 配置 | 检查点 :floppy_disk: | 得分 :chart_with_upwards_trend: | 可视化 :telescope: |
---|---|---|---|---|---|---|---|
Area 1 | 74.1 | 85.1 | 89.6 | 配置 | 检查点 | 得分 | 可视化 |
Area 2 | 44.9 | 57.1 | 67.9 | 配置 | 检查点 | 得分 | 可视化 |
Area 3 | 74.4 | 84.4 | 88.1 | 配置 | 检查点 | 得分 | 可视化 |
Area 4 | 63.8 | 74.7 | 81.1 | 配置 | 检查点 | 得分 | 可视化 |
Area 5 | 56.6 | 68.4 | 75.2 | 配置 | 检查点 | 得分 | 可视化 |
Area 6 | 73.3 | 83.4 | 87.8 | 配置 | 检查点 | 得分 | 可视化 |
ScanNet v2
数据集 | AP | AP_50 | AP_25 | 配置 | 检查点 :floppy_disk: | 分数 :chart_with_upwards_trend: | 可视化 :telescope: |
---|---|---|---|---|---|---|---|
ScanNet 验证集 | 55.2 | 73.7 | 83.5 | 配置 | 检查点 | 分数 | 可视化 |
ScanNet 测试集 | 56.6 | 78.0 | 87.0 | 配置 | 检查点 | 分数 | 可视化 |
ScanNet 200
数据集 | AP | AP_50 | AP_25 | 配置 | 检查点 :floppy_disk: | 分数 :chart_with_upwards_trend: | 可视化 :telescope: |
---|---|---|---|---|---|---|---|
ScanNet200 验证集 | 27.4 | 37.0 | 42.3 | 配置 | 检查点 | 分数 | 可视化 |
ScanNet200 测试集 | 27.8 | 38.8 | 44.5 | 配置 | 检查点 | 分数 | 可视化 |
STPLS3D
数据集 | AP | AP_50 | AP_25 | 配置 | 检查点 :floppy_disk: | 分数 :chart_with_upwards_trend: | 可视化 :telescope: |
---|---|---|---|---|---|---|---|
STPLS3D 验证集 | 57.3 | 74.3 | 81.6 | 配置 | 检查点 | 分数 | 可视化 |
STPLS3D 测试集 | 63.4 | 79.2 | 85.6 | 配置 | 检查点 | 分数 | 可视化 |
BibTeX :pray:
@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}
}