EPro-PnP

📢 新闻: 我们发布了 EPro-PnP-v2。新的更新预印本可在 arXiv 上找到。

EPro-PnP：用于单目物体位姿估计的广义端到端概率视角n点法
发表于 CVPR 2022（口头报告，最佳学生论文）。[论文][视频]
Hansheng Chen*^1,2, Pichao Wang†², Fan Wang², Wei Tian†¹, Lu Xiong¹, Hao Li²

¹同济大学，²阿里巴巴集团
*部分工作在阿里巴巴集团实习期间完成。
†通讯作者：Pichao Wang, Wei Tian。

引言

EPro-PnP 是一种用于端到端 6DoF 位姿估计网络的概率视角n点法（PnP）层。广义上讲，它本质上是广泛使用的分类 Softmax 层的连续对应物，并且在理论上可以推广到其他具有嵌套 优化的学习模型中。

给定该层的输入：一个 点对应集 ，由 3D 物体坐标 、2D 图像坐标 ，以及 2D 权重 组成。传统的 PnP 求解器会寻找一个最优位姿 （SE(3) 中的刚性变换），以使加权重投影误差最小化。以往的工作试图对 PnP 操作进行反向传播，然而 由于内部的 操作而本质上不可微分。这会导致当 中的所有组件都需要由网络学习时出现收敛问题。

相比之下，我们的概率 PnP 层输出位姿的后验分布，其概率密度 可以推导出来以便进行适当的反向传播。该分布通过蒙特卡洛采样来近似。借助 EPro-PnP，对应关系 可以从头开始全部学习，方法是通过最小化预测与目标位姿分布之间的 KL 散度。

模型

本仓库中的 V1 模型

[EPro-PnP-6DoF] 用于 6DoF 位姿估计

[EPro-PnP-Det] 用于 3D 物体检测

新的 V2 模型

EPro-PnP-Det v2：最先进的单目 3D 目标检测器

与 v1b 的主要区别：

• 使用 GaussianMixtureNLLLoss 作为辅助坐标回归损失
• 增加了深度和边界框的辅助损失

在提交时（2022年8月30日），EPro-PnP-Det v2 在 nuScenes 官方基准测试 上的所有基于相机的单帧目标检测模型中 排名第一（测试集，不使用额外数据）。

方法	TTA	主干网络	NDS	mAP	mATE	mASE	mAOE	mAVE	mAAE	训练轮数
EPro-PnP-Det v2 (我们的方法)	Y	R101	0.490	0.423	0.547	0.236	0.302	1.071	0.123	12 轮
PETR	N	Swin-B	0.483	0.445	0.627	0.249	0.449	0.927	0.141	24 轮
BEVDet-Base	Y	Swin-B	0.482	0.422	0.529	0.236	0.395	0.979	0.152	20 轮
EPro-PnP-Det v2 (我们的方法)	N	R101	0.481	0.409	0.559	0.239	0.325	1.090	0.115	12 轮
PolarFormer	N	R101	0.470	0.415	0.657	0.263	0.405	0.911	0.139	24 轮
BEVFormer-S	N	R101	0.462	0.409	0.650	0.261	0.439	0.925	0.147	24 轮
PETR	N	R101	0.455	0.391	0.647	0.251	0.433	0.933	0.143	24 轮
EPro-PnP-Det v1	Y	R101	0.453	0.373	0.605	0.243	0.359	1.067	0.124	12 轮
PGD	Y	R101	0.448	0.386	0.626	0.245	0.451	1.509	0.127	24+24 轮
FCOS3D	Y	R101	0.428	0.358	0.690	0.249	0.452	1.434	0.124	-

EPro-PnP-6DoF v2 用于 6DoF 姿态估计

与 v1b 的主要区别：

• 修复 w2d 缩放处理 (非常重要)
• 改进网络初始化
• 调整损失权重

通过这些改进，v2 模型可以在 无需 3D 模型 的情况下进行训练，并取得比 GDRNet 更好的性能（ADD 0.1d = 93.83），从而充分发挥简单端到端训练的潜力。

在您自己的模型中使用 EPro-PnP

我们提供了一个关于如何使用 EPro-PnP 层的 示例。

引用

如果您在研究中发现本项目有用，请考虑引用以下内容：

@inproceedings{epropnp, 
  author = {Hansheng Chen and Pichao Wang and Fan Wang and Wei Tian and Lu Xiong and Hao Li, 
  title = {EPro-PnP: Generalized End-to-End Probabilistic Perspective-n-Points for Monocular Object Pose Estimation}, 
  booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 
  year = {2022}
}

@article{epropnpv2,
  author={Chen, Hansheng and Tian, Wei and Wang, Pichao and Wang, Fan and Xiong, Lu and Li, Hao},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
  title={EPro-PnP: Generalized End-to-End Probabilistic Perspective-N-Points for Monocular Object Pose Estimation}, 
  year={2024},
  doi={10.1109/TPAMI.2024.3354997}
}

EPro-PnP 快速上手指南

EPro-PnP 是一个用于端到端 6DoF 姿态估计网络的概率 Perspective-n-Points (PnP) 层。它解决了传统 PnP 求解器不可微的问题，允许网络从头学习 3D-2D 对应关系，广泛应用于单目物体姿态估计和 3D 目标检测任务。

环境准备

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

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04)
• Python: 3.7 或更高版本
• 深度学习框架: PyTorch 1.8+ (需包含 CUDA 支持)
• 硬件: NVIDIA GPU (建议显存 8GB 以上，具体取决于模型版本和数据集)
• 编译器: GCC/G++ (用于编译自定义 CUDA 算子)

前置依赖安装： 建议使用 conda 创建独立环境，并优先使用国内镜像源加速依赖下载。

# 创建虚拟环境
conda create -n epro_pnp python=3.8 -y
conda activate epro_pnp

# 安装 PyTorch (以 CUDA 11.3 为例，可根据实际显卡版本调整，使用清华源)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu113

# 安装基础依赖
pip install numpy opencv-python scipy tqdm matplotlib -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目包含自定义 CUDA 算子，需要从源码进行安装。

1. 克隆代码仓库

   git clone https://github.com/tjiiv-cprg/EPro-PnP.git
   cd EPro-PnP
   

   (注：如需体验最新的 SOTA 模型，可参考 README 中的 EPro-PnP-v2 仓库链接)

2. 安装核心库 执行以下命令安装 epro_pnp 包及其 CUDA 扩展：

   pip install -e .
   

   如果在编译过程中报错，请确保已安装 ninja 构建工具 (pip install ninja) 且 CUDA 环境变量配置正确。

3. 验证安装 运行官方提供的 Demo 脚本验证层是否可用：

   python demo/fit_identity.py
   

   若无报错并输出收敛结果，则安装成功。

基本使用

EPro-PnP 的核心是作为一个可微层嵌入到神经网络中。以下是将其集成到自定义模型中的最小化示例。

1. 导入模块

import torch
from epro_pnp import EProPnP

2. 初始化层

初始化 PnP 层，通常需要指定采样次数（num_samples）以通过蒙特卡洛采样近似后验分布。

# num_samples: 蒙特卡洛采样点数，越大越精确但速度越慢
pnp_layer = EProPnP(num_samples=100).cuda()

3. 前向传播

假设网络已经输出了对应点集数据，直接调用层进行姿态解算。

输入数据格式说明：

• pts_3d: 3D 物体坐标 $[B, N, 3]$
• pts_2d: 2D 图像坐标 $[B, N, 2]$
• weights: 2D 权重 $[B, N, 2]$ (表示置信度)

# 模拟批量数据 (Batch=2, 点数=50)
B, N = 2, 50
pts_3d = torch.randn(B, N, 3).cuda()       # 3D 坐标
pts_2d = torch.randn(B, N, 2).cuda()       # 2D 投影坐标
weights = torch.rand(B, N, 2).cuda() + 0.1 # 权重 (必须为正)

# 执行 PnP 解算
# 输出 pose_dist 包含姿态的分布参数 (如均值和协方差)，可用于计算损失
pose_dist = pnp_layer(pts_3d, pts_2d, weights)

# 获取姿态均值 (旋转和平移)
pose_mean = pose_dist.mean
print(f"Estimated Pose Shape: {pose_mean.shape}") 
# 通常输出为 [B, 6] 或 [B, 3, 4]，具体取决于实现细节，代表 SE(3) 变换

4. 损失计算

由于输出是概率分布，通常使用 KL 散度（Kullback-Leibler Divergence）来衡量预测分布与真实姿态分布之间的差异，从而实现端到端训练。

# 假设 target_dist 是根据 Ground Truth 构建的目标分布
# loss = torch.distributions.kl_divergence(pose_dist, target_dist).mean()
# loss.backward()

更多详细用法（如结合检测头、可视化等）请参考仓库中的 demo/fit_identity.ipynb 笔记本及各子模型目录下的训练脚本。