Mask3D：用于3D实例分割的掩码Transformer

[项目主页] [论文] [演示]

新闻

• 2023年10月29日：请查看这个针对Mask3D的简易安装指南。
• 2023年1月17日：Mask3D已被ICRA 2023接收。 :fire:
• 2022年10月14日：新增了STPLS3D支持。
• 2022年10月10日：在ECCV 2022期间由Urbn3D Workshop举办的STPLS3D挑战赛中，Mask3D位列第二。
• 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环境：

# 一些用户在Ubuntu系统上使用AMD CPU时遇到了问题
# 请安装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="RAW_SCANNET_DATASET路径" \
--save_dir="data/processed/scannet" \
--git_repo="SCANNET_GIT_REPO路径" \
--scannet200=false/true

S3DIS

S3DIS数据集存在一些小错误，我们最初是手动修复的。我们很快会发布一个可以直接预处理原始数据集的脚本。目前，请按照此处的说明手动修复数据集。之后，按如下方式调用预处理脚本：

python -m datasets.preprocessing.s3dis_preprocessing preprocess \
--data_dir="PATH_TO_Stanford3dDataset_v1.2" \
--save_dir="data/processed/s3dis"

STPLS3D

python -m datasets.preprocessing.stpls3d_preprocessing preprocess \
--data_dir="PATH_TO_STPLS3D" \
--save_dir="data/processed/stpls3d"

训练与测试 :train2:

在 ScanNet 数据集上训练 Mask3D：

python main_instance_segmentation.py

请参考 配置脚本（例如 这里），以获取重现我们结果的详细说明。 最简单的推理命令如下：

python main_instance_segmentation.py \
general.checkpoint='PATH_TO_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: 用于3D语义实例分割的掩码Transformer}},
  author    = {Schult, Jonas and Engelmann, Francis and Hermans, Alexander and Litany, Or and Tang, Siyu and Leibe, Bastian},
  booktitle = {{国际机器人与自动化会议 (ICRA)}},
  year      = {2023}
}

Mask3D 快速上手指南

Mask3D 是一个基于 Mask Transformer 的 3D 实例分割模型，在 ScanNet、S3DIS 等数据集上达到了业界领先（SOTA）的性能。本指南将帮助你快速搭建环境并运行模型。

环境准备

在开始之前，请确保你的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
  • 注意: 部分用户在 AMD CPU 的 Ubuntu 系统上遇到过问题，建议预先安装 libopenblas-dev (sudo apt-get install libopenblas-dev)。
• Python: 3.10.9
• CUDA: 11.3
• GPU: 支持 CUDA 的 NVIDIA 显卡

安装步骤

1. 创建 Conda 环境并安装基础依赖

首先设置 CUDA 架构列表，然后创建并激活虚拟环境：

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

2. 安装 PyTorch 及相关库

按照官方源安装指定版本的 PyTorch、torch-scatter 和 detectron2：

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

3. 编译第三方依赖 (MinkowskiEngine, ScanNet Segmentator, PointNet2)

创建目录并依次编译必要的底层库：

mkdir third_party
cd third_party

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

cd ..

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

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

cd ../../

4. 安装 PyTorch Lightning

最后安装训练框架：

pip3 install pytorch-lightning==1.7.2

数据预处理

在训练或测试前，需对原始数据集进行预处理。假设你已下载原始数据，运行以下命令（以 ScanNet 为例）：

python -m datasets.preprocessing.scannet_preprocessing preprocess \
--data_dir="PATH_TO_RAW_SCANNET_DATASET" \
--save_dir="data/processed/scannet" \
--git_repo="PATH_TO_SCANNET_GIT_REPO" \
--scannet200=false

注：对于 S3DIS 和 STPLS3D 数据集，请参考项目原文中的特定预处理脚本指令。

基本使用

训练模型

使用默认配置在 ScanNet 数据集上开始训练：

python main_instance_segmentation.py

提示：如需复现特定论文结果或针对不同数据集训练，请参考 scripts/ 目录下的具体 shell 脚本配置。

推理测试 (使用预训练权重)

下载预训练模型 checkpoint 后，运行以下命令进行推理：

python main_instance_segmentation.py \
general.checkpoint='PATH_TO_CHECKPOINT.ckpt' \
general.train_mode=false

将 PATH_TO_CHECKPOINT.ckpt 替换为你实际的模型文件路径即可得到分割结果。