PoseDiffusion：通过扩散辅助光束平差法解决姿态估计问题

[论文] [项目页面]

更新：

• [2024年8月6日] 若要获取图5、图7和图10的原始数据，请从arXiv下载我们的提交文件。您可以通过点击 https://arxiv.org/format/2306.15667 上的“Download source”按钮来完成此操作。提交文件中包含以下文件：plots_co3dv2.tex、plots_re10k.tex 和 plots_co3dv2_suppl.tex。这些文件使用来自 csvs.tex 和 csvs_suppl.tex 的数据生成图表，而这两个文件也包含在LaTeX提交源代码中。我们在Re10K数据集上使用的划分可在 re10k_test_1800.txt 中找到。

• [2024年4月24日] 您可能也会对 VGGSfM 感兴趣，在该工具中，类似于PoseDiffusion的模型被用作相机预测器。它还支持通过光束平差法优化相机参数。

• [2024年4月24日] 更新了RealEstate10K数据集的检查点。

安装

我们提供了一个简单的安装脚本，默认会设置一个包含Python 3.9、PyTorch 1.13和CUDA 11.6的conda环境。

source install.sh

快速入门

1. 下载检查点

您可以下载在Co3D或RealEstate10K上训练的模型检查点。请注意，RealEstate10K的检查点是在336×336的图像尺寸上重新训练的（您需要在相应的配置文件中更改image_size）。预测的相机位姿和焦距是按照NDC坐标系定义的。

2. 运行演示

python demo.py image_folder="samples/apple" ckpt="/PATH/TO/DOWNLOADED/CKPT"

您可以通过指定不同的image_folder来尝试自己的数据。

在Quadro GP100 GPU上，对于一个20帧序列，不启用GGS时的推理时间约为0.8秒，启用GGS时则约为80秒（其中匹配提取耗时20秒）。

您可以在./cfgs/default.yaml中选择是否启用GGS或其他设置。

我们默认使用Visdom进行可视化。请确保您的Visdom设置正确配置，以便准确地可视化结果。不过，运行模型并不需要Visdom。

训练

1. 预处理标注

首先，请按照这里的说明预处理Co3D V2数据集的标注。这将显著减少训练过程中的数据处理时间。

2. 指定路径

接下来，在./cfgs/default_train.yaml中指定CO3D_DIR和CO3D_ANNOTATION_DIR的路径。CO3D_DIR应设置为您下载的Co3D数据集所在的路径，而CO3D_ANNOTATION_DIR则应指向步骤1中完成预处理后生成的标注文件所在位置。

3. 开始训练

• 单GPU训练：

  python train.py
  

• 多GPU训练时，请使用accelerate启动训练脚本，例如在1台机器上的8个GPU（进程）上训练：

  accelerate launch --num_processes=8 --multi_gpu --num_machines=1 train.py 
  

所有配置均在./cfgs/default_train.yaml中指定。请注意，我们使用Visdom记录日志。

每次迭代的训练时间大约为13秒，具体取决于设备的不同。您可以通过查看日志中的sec/it来确认。在8块A100 GPU上，整个训练大约需要23天。

注意：在某些集群中，我们发现公开发布的多GPU训练代码运行速度非常慢。这似乎是因为accelerate在某些设置下表现不佳，导致数据加载非常缓慢（甚至卡死）。最简单的解决方案是移除代码中的accelerate (accelerator)部分，改用PyTorch自带的分布式训练工具或PyTorch Lightning来启动训练。这个问题不会影响单GPU训练。如果您遇到类似情况，请提交一个问题或在此处报告您的案例 here（这很可能与accelerate的数据加载有关，简单解决方案是使用PyTorch的分布式训练功能）。

测试

1. 指定路径

请在./cfgs/default_test.yaml中指定CO3D_DIR、CO3D_ANNOTATION_DIR和resume_ckpt的路径。resume_ckpt指您下载的模型检查点。

2. 运行测试

python test.py

您可以通过调整./cfgs/default_test.yaml中的num_frames、GGS.enable等参数来尝试不同的测试设置。

致谢

感谢denoising-diffusion-pytorch、guided-diffusion、hloc和relpose的出色实现。

许可证

有关本代码所采用许可证的详细信息，请参阅LICENSE文件。

PoseDiffusion 快速上手指南

PoseDiffusion 是一个基于扩散模型辅助光束法平差（Bundle Adjustment）的相机姿态估计工具。本指南将帮助您快速完成环境配置并运行演示。

环境准备

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.9
• PyTorch: 1.13
• CUDA: 11.6
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (单卡或多卡均可，多卡训练需注意数据加载优化)
• 可选依赖: Visdom (用于可视化结果，非运行必需)

注意：官方安装脚本默认创建包含上述版本要求的 Conda 环境。如果您需要使用更新的 PyTorch 或 CUDA 版本，可能需要手动修改 install.sh 或自行创建环境。

安装步骤

1. 克隆仓库并进入目录（假设您已获取源码）。
2. 运行官方提供的安装脚本，它将自动设置 Conda 环境及依赖：

source install.sh

基本使用

1. 下载预训练模型

根据您的需求下载在 Co3D 或 RealEstate10K 数据集上训练的权重文件：

• Co3D 模型: 下载链接
• RealEstate10K 模型: 下载链接
  • 注：RealEstate10K 模型是在 336 分辨率图像上重新训练的，使用时需在配置文件中调整 image_size。

2. 运行演示

准备好您的图像文件夹（或使用自带的 samples/apple），执行以下命令进行推理：

python demo.py image_folder="samples/apple" ckpt="/PATH/TO/DOWNLOADED/CKPT"

• 请将 /PATH/TO/DOWNLOADED/CKPT 替换为您实际下载的模型权重路径。
• 您可以将 image_folder 指向包含自定义图像序列的文件夹进行测试。

3. 查看结果

• 可视化: 默认使用 Visdom 展示结果。请确保已启动 Visdom 服务 (visdom) 并在浏览器访问对应端口。
• 性能参考: 在 Quadro GP100 GPU 上，处理 20 帧序列的耗时约为：
  • 不开启 GGS: ~0.8 秒
  • 开启 GGS (含匹配提取): ~80 秒
• 配置调整: 如需启用/禁用 GGS 或其他参数，请编辑 ./cfgs/default.yaml 文件。