Pix2Pose: Pixel-Wise Coordinate Regression of Objects for 6D Pose Estimation

  本文介绍了一种两阶段的6D位姿估计方法，第一阶段通过CNN预测图像中物体像素集对应的三维点在物体坐标系中的归一化坐标；第二阶段基于二维像素点和预测三维点的对应关系通过RANSAC-based PnP方法求解物体的位姿。
  论文中还提出了一种新的损失函数Transformer loss，该损失函数可以较好地解决对称物体地位姿估计问题。

1 网络结构

  网络采用生成对抗式学习方式训练。其中，Generator部分采用了encoder-decoder式的结构，输入是固定大小为$128\times 128$的彩色图像，输出是$4\times 128\times 128$的张量(前三个通道是该位置像素对应的三维点的归一化世界坐标值，最后一个通道是三维点坐标的误差值)。
  需要注意的是，预测得到的三维点坐标是基于物体坐标系的，而且三个维度的坐标值通道Tanh激活函数归一化到$[-1,1]$范围内；而三维坐标的误差值通过Sigmoid激活函数被归一化到$[0,1]$范围内。
  顺便提一下，三维坐标对应的标签在给定真实位姿数据下通过相机模型逆运算得到，之后再对得到的三维坐标值进行归一化处理。

2 损失函数

  对于非对称物体，三维坐标回归部分的损失函数如下计算：
$$L_r=\frac{1}{n}[\beta \sum _{i\in M}||I_{3D}^i-I_{gt}^i|{|}_1+\sum _{i\notin M}||I_{3D}^i-I_{gt}^i|{|}_1]\text{\hspace{1em}(1)}$$
式中，$n$是像素个数，$M$表示图像中物体对应的Mask，注意，这里的mask包含了物体被遮挡部分，这样就使得训练出来的网络对遮挡有一定的鲁棒性；对于属于物体的像素，损失项前面乘以一个大于1的$\beta$，这是因为物体像素对于训练更重要。
  而对于对称物体，三维坐标回归部分的损失函数如下计算：
$$L_{3D}=\underset{p\in sym}{\min }{L}_r(I_{3D},R_p{I}_{gt})\text{\hspace{1em}(2)}$$
式中，$R_p$是一个$3\times 3$的旋转向量，它取自于集合$sym$中，该集合包含了若干个旋转向量，每一个旋转向量的旋转轴为物体对称轴、旋转角度为物体对称角度或其整数倍，当然该集合中还包括了单位向量，表示物体不旋转。
  误差项的损失函数如下计算：
$$L_e=\frac{1}{n}\sum _i||I_e^i-\min [L_r^i,1]|{|}_2^2,\beta =1.\text{\hspace{1em}(3)}$$

3 位姿计算

  位姿结算分为以下两步：

1. 给定一张输入图像，首先通过2D目标检测网络crop出保持长宽比的图像块，并将其resize成$128\times 128$大小，然后将其输入到训练好的网络中，将得到的三维坐标图像(也即$I_{3D}$)中的非零值标记出来，再将误差预测值小于异常阈值${\theta }_o$对应的像素标记出来，去上述两部分标记的并集制作mask，再以该mask重新从原图像crop$128\times 128$的图像块，注意重新得到的图像块的未标记部分用[0, 0, 0]填充，也即黑色。
2. 将第一阶段得到的图像块再次输入到网络中，将输出中误差值大于${\theta }_i$的像素剔除到，则可以用剩余的像素集合通过RANSAC-based PnP求解位姿。