3DETR：用于3D目标检测的端到端Transformer模型

3DETR 的 PyTorch 实现及预训练模型。

3DETR（3D DEtection TRansformer）是复杂的手工设计3D检测流水线的一种更简单的替代方案。 它不依赖于 PointNet++ 等3D骨干网络，且使用的3D专用算子较少。 3DETR 在性能上可与 VoteNet 等3D检测方法相媲美，甚至有所超越。 其编码器还可用于其他3D任务，例如形状分类。 更多详情请参阅论文《用于3D目标检测的端到端Transformer模型》（http://arxiv.org/abs/2109.08141）。

[网站] [arXiv] [BibTeX]

代码说明。 我们的代码基于 DETR 和 VoteNet 等先前的工作，并在实现中力求简洁。我们希望这能为3D检测领域的研究提供便利。

预训练模型

我们提供了预训练模型权重以及在验证集上的相应指标（按类别计算的AP和召回率）。 我们还提供了一个Python脚本 utils/download_weights.py，方便用户下载权重和指标文件。

架构	数据集	训练轮数	AP25	AP50	模型权重	评估指标
3DETR-m	SUN RGB-D	1080	59.1	30.3	权重	指标
3DETR	SUN RGB-D	1080	58.0	30.3	权重	指标
3DETR-m	ScanNet	1080	65.0	47.0	权重	指标
3DETR	ScanNet	1080	62.1	37.9	权重	指标

模型 zoo

为方便起见，我们提供了针对不同训练轮数的 3DETR 模型权重。

架构	数据集	训练轮数	AP25	AP50	模型权重	评估指标
3DETR-m	SUN RGB-D	90	51.0	22.0	权重	指标
3DETR-m	SUN RGB-D	180	55.6	27.5	权重	指标
3DETR-m	SUN RGB-D	360	58.2	30.6	权重	指标
3DETR-m	SUN RGB-D	720	58.1	30.4	权重	指标
3DETR	SUN RGB-D	90	43.7	16.2	权重	指标
3DETR	SUN RGB-D	180	52.1	25.8	权重	指标
3DETR	SUN RGB-D	360	56.3	29.6	权重	指标
3DETR	SUN RGB-D	720	56.0	27.8	权重	指标
3DETR-m	ScanNet	90	47.1	19.5	权重	指标
3DETR-m	ScanNet	180	58.7	33.6	权重	指标
3DETR-m	ScanNet	360	62.4	37.7	权重	指标
3DETR-m	ScanNet	720	63.7	44.5	权重	指标
3DETR	ScanNet	90	42.8	15.3	权重	指标
3DETR	ScanNet	180	54.5	28.8	权重	指标
3DETR	ScanNet	360	59.0	35.4	权重	指标
3DETR	ScanNet	720	61.1	40.2	权重	指标

运行 3DETR

安装

我们的代码已在 PyTorch 1.9.0、CUDA 10.2 和 Python 3.6 环境下测试通过，也可能在其他版本中正常运行。

您需要通过以下命令安装 pointnet2 层：

cd third_party/pointnet2 && python setup.py install

此外，您还需要安装以下 Python 依赖项（可通过 conda install 或 pip install 安装）：

matplotlib
opencv-python
plyfile
'trimesh>=2.35.39,<2.35.40'
'networkx>=2.2,<2.3'
scipy

部分用户在使用 CUDA 11 或更高版本时遇到了问题。如果遇到 CUDA 相关问题，请尝试使用 CUDA 10.2。

可选地，您可以安装 gIOU 的 Cython 化实现以加快训练速度。

conda install cython
cd utils && python cython_compile.py build_ext --inplace

基准测试

数据集准备

我们沿用了 VoteNet 代码库中的数据预处理流程。 SUN RGB-D 的预处理说明请参见 这里，ScanNet 的预处理说明请参见 这里。

您可以在 datasets/sunrgbd.py 和 datasets/scannet.py 中编辑数据集路径，也可以选择在运行时指定路径。

测试

准备好数据集后，您可以按照以下方式测试预训练模型：

python main.py --dataset_name <数据集名称> --nqueries <查询数量> --test_ckpt <检查点路径> --test_only [--enc_type masked]

对于 SUN RGB-D 数据集，我们使用 128 个查询；对于 ScanNet 数据集，使用 256 个查询。 在测试 3DETR-m 检查点时，您需要添加 --enc_type masked 标志。 请注意，测试过程具有随机性（由于点云采样和查询采样的随机性），因此每次运行的 AP25 结果可能会有 1% 左右的波动。这种推理过程的随机性在 VoteNet 等方法中也很常见。

如果您尚未编辑 datasets 文件夹中文件的数据集路径，则可以使用 --dataset_root_dir 标志来指定数据集的根目录路径。

训练

只需运行 main.py 即可简单地训练模型：

python main.py --dataset_name <数据集名称> --checkpoint_dir <保存输出的路径>

为了复现论文中的结果，我们在 scripts 文件夹中提供了相应的参数设置。不同训练运行之间预计会出现 AP25 约 1% 的差异。

您可以通过按照 scripts/scannet_quick.sh 文件中的说明，在 ScanNet 数据集上训练一个 3DETR 模型 90 个 epoch 来快速验证安装是否正确，并将其与 Model Zoo 中的预训练检查点进行比较。

许可证

3DETR 的大部分代码采用 Apache 2.0 许可证，具体见 LICENSE 文件。然而，项目中部分组件使用了单独的许可证条款：PointNet2 的许可信息可在 https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/UNLICENSE 找到。

贡献

我们欢迎您的拉取请求！有关更多信息，请参阅 CONTRIBUTING 和 CODE_OF_CONDUCT。

引用

如果您觉得本仓库对您有所帮助，请考虑给它加星 :star: 并引用以下文献：

@inproceedings{misra2021-3detr,
    title={{An End-to-End Transformer Model for 3D Object Detection}},
    author={Misra, Ishan and Girdhar, Rohit and Joulin, Armand},
    booktitle={{ICCV}},
    year={2021},
}

3DETR 快速上手指南

3DETR 是一个基于 Transformer 的端到端 3D 物体检测模型。它摒弃了复杂的手工设计流程和特定的 3D 骨干网络（如 PointNet++），仅使用少量 3D 专用算子，即可在 SUN RGB-D 和 ScanNet 数据集上获得媲美或超越 VoteNet 等方法的性能。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.6+
• PyTorch: 1.9.0
• CUDA: 10.2 (注意：部分用户在使用 CUDA 11 及以上版本时遇到问题，建议优先使用 CUDA 10.2)
• 编译器: 支持 C++14 的 GCC

安装步骤

1. 克隆代码库

git clone https://github.com/facebookresearch/3detr.git
cd 3detr

2. 安装 Python 依赖

建议使用 pip 安装所需依赖。国内用户可添加清华源加速下载：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple matplotlib opencv-python plyfile scipy
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple 'trimesh>=2.35.39,<2.35.40'
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple 'networkx>=2.2,<2.3'

3. 编译 PointNet++ 模块

3DETR 依赖自定义的 PointNet++ 层，需要手动编译安装：

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

4. (可选) 安装加速模块

为了加快训练速度，可以编译安装 Cython 版本的 gIOU：

conda install cython
cd utils && python cython_compile.py build_ext --inplace
cd ..

基本使用

1. 数据集准备

3DETR 沿用了 VoteNet 的数据预处理流程。您需要先下载并预处理 SUN RGB-D 或 ScanNet 数据集。

• SUN RGB-D 预处理指南：VoteNet SUN RGB-D
• ScanNet 预处理指南：VoteNet ScanNet

预处理完成后，您可以在代码中修改默认路径，或在运行时通过 --dataset_root_dir 指定。

2. 下载预训练模型

您可以使用官方提供的脚本下载预训练权重和评估指标：

python utils/download_weights.py

或者直接从 Model Zoo 表格中手动下载 .pth 权重文件。

3. 推理测试 (Testing)

使用预训练模型进行测试。以下命令以 ScanNet 数据集为例（使用 256 个查询向量）：

python main.py --dataset_name scannet --nqueries 256 --test_ckpt <path_to_checkpoint> --test_only

如果是测试 3DETR-m (masked encoder) 版本的模型，请务必添加 --enc_type masked 标志：

python main.py --dataset_name sunrgbd --nqueries 128 --test_ckpt <path_to_masked_checkpoint> --test_only --enc_type masked

注意：由于点云采样和查询采样的随机性，测试结果可能会有约 1% AP25 的波动，这属于正常现象。

4. 模型训练 (Training)

从头开始训练模型。以下命令将模型输出保存到指定目录：

python main.py --dataset_name scannet --checkpoint_dir ./outputs/scannet_train

若要复现论文中的结果，请参考 scripts/ 文件夹中提供的具体脚本（如 scripts/scannet_quick.sh 可用于快速验证安装）。