PREDATOR：低重叠度三维点云配准（CVPR 2021，口头报告）

本仓库是论文的官方实现：

PREDATOR：低重叠度三维点云配准

*黄圣宇， *赞·戈伊契奇，米哈伊尔·乌斯维亚佐夫，安德烈亚斯·维瑟，康拉德·辛德勒
|苏黎世联邦理工学院 | * 共同第一作者

如需使用 MinkowskiEngine 作为骨干网络的实现，请查看 此链接。

更多信息请访问 项目官网。

联系方式

如有任何问题，请联系我们：

• 黄圣宇 {shengyu.huang@geod.baug.ethz.ch}
• 赞·戈伊契奇 {zan.gojcic@geod.baug.ethz.ch}

最新消息

• 2021-08-09：我们更新了 arXiv 上的 论文，性能得到进一步提升！
• 2021-06-02：修复了 k-nn 图中的特征提取 bug，请参阅此 issue 中的改进效果。敬请关注其他实验的更新！
• 2021-05-31：请在 项目页面 上观看我们的视频和海报！
• 2021-03-25：最终版已发布于 arXiv！我还就 Predator 进行了中文演讲，录像可在以下平台观看：Bilibili、Youtube。
• 2021-02-28：基于 MinkowskiEngine 的 PREDATOR 发布。
• 2020-11-30：代码与论文正式发布。

使用说明

该代码已在以下环境中测试通过：

• Python 3.8.5，PyTorch 1.7.1，CUDA 11.2，gcc 9.3.0，GeForce RTX 3090/GeForce GTX 1080Ti

注意：我们观察到数据加载器因内存问题而随机崩溃。若遇到类似情况，请尝试减少工作进程数或增加 CPU 内存。目前我们已发布基于稀疏卷积的 Predator 实现，详情请见 此处！

环境要求

请在您的工作目录中运行以下命令以创建虚拟环境并安装所需依赖：

git clone https://github.com/overlappredator/OverlapPredator.git
virtualenv predator; source predator/bin/activate
cd OverlapPredator; pip install -r requirements.txt
cd cpp_wrappers; sh compile_wrappers.sh; cd ..

数据集与预训练模型

对于 KITTI 数据集，请按照 KITTI Odometry 官网 的说明下载 KITTI 奥多米测训练集。

我们提供：

• 预处理过的 3DMatch 成对数据集（体素网格下采样后的片段及其真实变换矩阵）
• 原始密集型 3DMatch 数据集
• ModelNet 数据集
• 在 3DMatch、KITTI 和 ModelNet 上的预训练模型

可通过运行以下命令下载预处理数据及模型：

sh scripts/download_data_weight.sh

若需下载原始密集型 3DMatch 数据，请执行：

wget --no-check-certificate --show-progress https://share.phys.ethz.ch/~gsg/pairwise_reg/3dmatch.zip
unzip 3dmatch.zip

文件夹组织如下：

• 3dmatch
  • train
    • 7-scenes-chess
      • fragments
        • cloud_bin_*.ply
        • …
      • poses
        • cloud_bin_*.txt
        • …
    • …
  • test

3DMatch（室内）

训练

创建虚拟环境并下载数据集后，可使用以下命令训练 Predator：

python main.py configs/train/indoor.yaml

评估

对于 3DMatch，为复现主论文中的表 2，我们首先通过运行以下命令提取特征及重叠/匹配度分数：

python main.py configs/test/indoor.yaml

特征及分数将保存至 snapshot/indoor/3DMatch。随后可使用 RANSAC 估计变换参数，具体操作如下：

for N_POINTS in 250 500 1000 2500 5000
do
  python scripts/evaluate_predator.py --source_path snapshot/indoor/3DMatch --n_points $N_POINTS --benchmark 3DMatch --exp_dir snapshot/indoor/est_traj --sampling prob
done

根据 RANSAC 使用的点数不同，此过程可能需要几分钟。最终结果将存储在 snapshot/indoor/est_traj/{benchmark}_{n_points}_prob/result 中。若要在 3DLoMatch 基准上评估 PREDATOR，请在 configs/test/indoor.yaml 中将 3DMatch 更改为 3DLoMatch。

演示

我们准备了一个小型演示，展示了使用 3DMatch 数据集中两个随机片段的完整 Predator 流程。要运行演示，请执行：

python scripts/demo.py configs/test/indoor.yaml

演示脚本将可视化输入点云、推断出的重叠区域以及根据估计变换参数对齐后的点云：

ModelNet（合成数据）

训练

要在 ModelNet 上训练 PREDATOR，请运行：

python main.py configs/train/modelnet.yaml

我们提供一个小型脚本来评估 Predator 在 ModelNet 测试集上的表现，请运行：

python main.py configs/test/modelnet.yaml

由于 RANSAC 的随机性，旋转和平移误差可能会比报告的结果更好或更差。

KITTI（室外）

我们提供一个小型脚本来评估 Predator 在 KITTI 测试集上的表现。配置好 KITTI 数据集后，请运行：

python main.py configs/test/kitti.yaml

结果将保存到日志文件中。

自定义数据集

针对自定义数据集的训练和测试，我们有以下几点建议：

• 若数据集场景与室内相似，请先运行 demo.py 检查模型的泛化能力，再进行重新训练。
• 请务必在数据加载器中对数据进行体素下采样，可参考 kitti.py。

引用

若您在工作中使用了本代码或将其应用于项目中，请考虑引用以下内容：

@InProceedings{Huang_2021_CVPR,
    author    = {Huang, Shengyu and Gojcic, Zan and Usvyatsov, Mikhail and Wieser, Andreas and Schindler, Konrad},
    title     = {Predator: Registration of 3D Point Clouds With Low Overlap},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2021},
    pages     = {4267-4276}
}

致谢

在本项目中，我们使用了以下工作的官方实现（部分）：

• FCGF（KITTI 数据预处理）
• D3Feat（KPConv 主干网络）
• 3DSmoothNet（3DMatch 数据准备）
• MultiviewReg（3DMatch 基准测试）
• SuperGlue（Transformer 部分）
• DGCNN（自注意力图神经网络）
• RPMNet（ModelNet 数据预处理与评估）

我们感谢各位作者开源他们的方法。同时，我们也感谢审稿人，尤其是第 2 位审稿人提供的宝贵意见。

OverlapPredator 快速上手指南

OverlapPredator 是 CVPR 2021 Oral 论文《PREDATOR: Registration of 3D Point Clouds with Low Overlap》的官方实现，专为解决低重叠率 3D 点云配准问题而设计。

1. 环境准备

本代码已在以下环境中测试通过，建议尽量匹配该配置以避免兼容性问题：

• 系统: Linux (推荐)
• Python: 3.8.5
• PyTorch: 1.7.1
• CUDA: 11.2
• 编译器: gcc 9.3.0
• 显卡: GeForce RTX 3090 / GTX 1080Ti

前置依赖注意：

• 需要安装 virtualenv 用于创建隔离环境。
• 编译 C++ 扩展需要 gcc 和对应的 CUDA Toolkit。
• 内存提示：数据加载器可能因内存不足崩溃，如遇此问题请减少 workers 数量或增加物理内存。

2. 安装步骤

请按顺序执行以下命令完成环境搭建和依赖安装：

# 1. 克隆仓库
git clone https://github.com/overlappredator/OverlapPredator.git
cd OverlapPredator

# 2. 创建并激活虚拟环境
virtualenv predator
source predator/bin/activate

# 3. 安装 Python 依赖
# 国内用户可使用清华源加速：pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -r requirements.txt

# 4. 编译 C++ 封装器
cd cpp_wrappers
sh compile_wrappers.sh
cd ..

3. 基本使用

3.1 下载数据与预训练模型

项目提供了脚本一键下载预处理后的 3DMatch 数据集及预训练权重（包含 3DMatch, KITTI, ModelNet）：

sh scripts/download_data_weight.sh

注：如需原始稠密 3DMatch 数据，需单独使用 wget 下载（见原文 README），通常直接使用上述脚本下载的预处理数据即可开始训练或测试。

3.2 快速演示 (Demo)

最简单的使用方式是运行官方提供的 Demo 脚本。它将加载两个随机的 3DMatch 点云片段，展示重叠区域推断及配准对齐效果：

python scripts/demo.py configs/test/indoor.yaml

运行后将在可视化窗口中看到输入点云、推断的重叠区域以及配准后的结果。

3.3 模型评估 (Evaluate)

使用预训练模型在 3DMatch 测试集上进行评估（复现论文 Table 2 结果）：

第一步：提取特征与重叠分数

python main.py configs/test/indoor.yaml

特征将保存至 snapshot/indoor/3DMatch。

第二步：执行 RANSAC 配准评估

for N_POINTS in 250 500 1000 2500 5000
do
  python scripts/evaluate_predator.py --source_path snapshot/indoor/3DMatch --n_points $N_POINTS --benchmark 3DMatch --exp_dir snapshot/indoor/est_traj --sampling prob
done

最终结果将保存在 snapshot/indoor/est_traj/{benchmark}_{n_points}_prob/result 目录下。

3.4 模型训练 (Train)

如需从头训练模型（以室内场景 3DMatch 为例）：

python main.py configs/train/indoor.yaml

其他数据集（如 ModelNet 或 KITTI）只需将配置文件替换为 configs/train/modelnet.yaml 或 configs/test/kitti.yaml 即可。