基于改进热图损失函数的目标6D姿态估计算法

如前文所述，热图绝大多数像素都是背景像素，对关键点至关重要的像素非常少，如图1所示，其中图1（a）为目标图像，红色部分为目标的关键点，图1（b）为该关键点对应的热图，图1（c）为图1（b）热图的局部放大图。因此，我们希望用合适的损失函数使网络训练可以对不同像素进行区分，使前景像素训练收敛速度更快，同时要对前景像素的小误差有更高的敏感度。

图1  关键点热图

Fig. 1  Keypoint and heatmap

在热图回归中，常使用损失函数MSE Loss即均方误差损失函数。但MSE Loss梯度随着误差增大而增大，大误差的损失梯度要高于小误差，这导致神经网络更关注于大误差部分，而忽略小误差部分。

Adaptive Wing Loss^［10］（AWing Loss）提高了小误差的敏感度，并对前景像素与背景像素做了区分。但AWing Loss前景像素梯度只在误差为0值附近大于背景像素，这导致在网络训练的大部分时期，网络都是向着背景像素误差减小方向前进，与热图前景像素更重要的特性背道而驰。MSE Loss和AWing Loss的表达式如式（1）、（2）所示。

根据热图本身特性，理想的热图损失函数应该有以下特性：（1）大误差部分具有恒定梯度，不会导致产生梯度爆炸的问题。（2）背景像素的小误差部分无需精准的预测，只需要逐步回归到零值，损失函数的特性与MSE Loss损失函数类似。（3）前景图像的小误差部分的回归精确程度是热图回归的关键，对小误差要有更大的梯度。损失函数特性与AWing Loss函数类似。并且在整个训练过程中前景像素的梯度都应该大于背景像素的梯度，这也有利于网络向着重要的前景像素损失降低方向训练。本文在AWing Loss的基础上提出适用于热图回归的损失函数Heatmap Wing Loss（HWing Loss），其表达式如公式（3）所示。

HWing Loss采用分段函数设计，p与pˆ分别代表真实像素值与预测像素值，p−pˆ为图像像素误差。θ为阈值，在0~1之间。 δ、ε、ω为参数，其中δ>1,ε>2。

为了提升前景像素损失函数的梯度，让整个网络训练过程能以前景像素主导，HWing Loss相比于AWing Loss做了两点改进。

首先是小误差部分，对像素误差p−pˆ增加系数1/(δ−p)，对误差进行缩放。此时将AWing Loss与HWing Loss函数看作误差p−pˆ的函数。在p−pˆ>0部分，其梯度分别如公式（4）、（5）所示。可以看出，当参数δ∈(1,2),ε>2时，对于背景像素p=0，公式（5）中(δ−p)ε−p>1，此时HWing Loss的函数梯度小于AWing Loss，即HWing Loss中1/(δ−p)部分减小了损失函数的梯度。而对前景像素p=1，公式（5）中(δ−p)ε−p<1，HWing Loss的函数梯度大于AWing Loss，即HWing Loss中1/(δ−p)部分放大了损失函数的梯度。因此，HWing Loss中增加的系数1/(δ−p)可以在前景像素处对损失函数梯度进行放大，同时缩小背景像素的损失函数，从而使损失函数在前景像素的梯度大于背景像素。其次，对大误差部分，取消AWing Loss中的系数A，从而避免在大误差部分背景像素梯度值大于前景像素的问题。

图2展示了HWing Loss分别在背景像素p=0（蓝色图像）和前景像素p=1（红色图像）时与MSE Loss（绿色图像）的比较，其中参数δ=1.8,ε=2.1。可以看出，HWing Loss在各点都是连续的，且梯度在0处连续。在p=1即前景像素部分，其特性与AWing Loss函数相像；在p=0即背景像素部分，其特性类似于MSE Loss。由图3（b）可以看出，本文提出的HWing Loss在小误差部分，误差p−pˆ一定的情况下，损失函数梯度会随着像素p增大而增大，大误差部分的损失函数梯度保持梯度恒定。这符合上文提到的理想损失函数特性，从而可以使整个网络训练过程都以前景像素为主。

图2  不同损失函数对比图

Fig.2  Diagram of loss functions

图3  HWing Loss示意图

Fig.3  Diagram of HWing Loss

使用HWing Loss损失函数训练热图回归网络时，训练开始所有像素都有较大误差，都应用损失函数的大误差部分，同时减小误差。当误差减小至θ后，背景像素随着误差减小，梯度逐渐减小，使得网络逐渐不再关注其收敛。而前景像素随着误差减小，梯度逐渐增大，在网络损失中占据主导位置，即网络会继续向着前景像素误差减小的方向训练，这符合热图前景像素更加重要的特点。

本文使用HRNet^［11-12］作为骨干网络构造热图回归网络。为减小网络运算量，使用步长为2的卷积将特征图由H×W降至H/2×W/2，输入到HRNet中。而为了减小特征网络输出计算关键点位置时，由于特征图与原图尺度不一致需要缩放产生的误差，将HRNet输出的特征进行上采样至原图大小，并与原图进行拼接，经过最后一个卷积块的运算，得到原图大小的热图。我们使用H×W×3大小的图像作为输入，每张图像定义C个关键点时，网络的输出为H×W×C的热图，在训练时使用上文提出的HWing Loss进行像素级的监督。

由于热图的前景像素部分只占整个图像的不足1%，即使我们使用了HWing Loss提高网络对前景像素的关注度，在训练时，整个损失仍然会以重要性不高的背景像素为主，导致网络不能更精确地得到前景像素的预测。因此，我们使用文献^［10］中的策略，在计算网络损失时，对前景像素与背景像素给予不同的权重。如图5所示，首先将热图进行灰度膨胀操作，在训练损失计算时，对膨胀热图像素值高于0.2的部分给予10倍的权重，使网络更加关注于前景像素的误差。

图5  加权损失示意图

Fig. 5  Schematic diagram of weighted loss

在传统方法中，通常使用热图的最大值位置作为关键点的预测值。但是由于图像是离散的，而关键点的位置很有可能不在像素点位置，这就导致了只用最大值位置作为关键点预测值会产生量化误差，影响最终姿态估计结果。

Zhang等^［13］提出了基于高斯分布的热图推理关键点的方法，但在二维图像上进行运算耗时较多。本文对其进行一维简化。

考虑到热图是由二维高斯函数生成，根据高斯函数的可分离性，二维高斯函数可以分解为两个一维高斯函数。

极值点在u处的一维高斯函数为：

为降低逼近难度，对高斯函数进行对数运算，将高斯函数变为二次函数，同时，可以保持极值点位置。

其中在极值点u处，L(x)以及其一阶二阶导数值分别为：

假设预测热图极大值在x方向坐标值为u0，易知u0位于极值点u附近，可以通过u0推断u位置。我们对L(x)在u0位置进行泰勒展开，忽略三阶小量：

联立式（9）、（11）、（12）可得：

其中，L'(u0)、L''(u0)分别可由预测热图在最大值位置在x方向上的一阶差分与二阶差分计算得到。

同理，对y方向坐标有：

我们使用公式（13）、（14）进行关键点优化。

通过预测热图推得关键点位置时，由于目标可能被遮挡，不可见的关键点会出现较大的定位误差从而影响姿态估计，所以需要对关键点进行筛选。考虑到热图每个像素值的本质是该关键点出现在该位置的概率，可以使用预测热图极值位置的像素值作为该关键点的置信系数，并对关键点进行筛选。为了保证姿态计算的稳定性，选择置信系数大的K个图像关键点。

获取图像关键点后，结合已知的三维关键点以及相机的内参，可以通过求解PnP问题计算得到姿态。我们使用EPnP算法^［14］进行PnP问题的求解。

实验中使用的数据集包括LINEMOD数据集和Occlusion LINEMOD数据集。

LINEMOD数据集是6D目标姿态估计的标准数据集，这个数据集包含多个姿态估计场景，比如复杂背景、无纹理目标的场景。每一幅图像的中心都有一个标记了平移、旋转和类别的目标，同时该数据集还提供了每个目标的3D模型。在LINEMOD 数据集中共有15 783张图像和13类目标，每个目标特征大约有1 200个实例。

Occlusion LINEMOD数据集是对LINEMOD数据集的扩展，每一幅图像包含一个被标记的目标，大部分图像中的目标被部分遮挡。在实验中，Occlusion LINEMOD数据集仅用来测试，只用LINEMOD 数据集进行训练。

我们使用ADD（-S）指标评估算法性能，其中ADD指标度量的是模型顶点之间3D平均距离。如果3D模型顶点的坐标与估计的坐标之间的平均距离小于目标直径的10%，预测就是正确的。对于对称对象，使用ADD-S指标度量，其平均距离是基于最近的点距离计算的。性能评估时以测试集中预测正确的图像数量与总数的百分比数计算。

本文使用PyTorch搭建训练环境。在LINEMOD数据集中每个类别随机选择30%的图像作为训练集，其余70%作为测试集。为了防止过拟合及增大训练空间，我们在训练集内添加合成图像。对于每个对象，我们渲染了10 000张视点均匀采样的图像。同时用剪切和粘贴策略合成了另外30 000张图像，每幅合成图像的背景随机采样自SUN397^［15］。

我们在目标三维模型上使用迭代最远点算法标注32个关键点并创建监督热图。输入图像裁剪至128×128，每一批次训练32张图像，学习率设置为0.01，训练30个轮次。对于损失函数HWing Loss，经过多次实验，选择参数ω=14,θ=0.5,δ=1.8,ε=2.1。在测试过程中，选择置信系数最大的24个关键点输入EPnP算法进行姿态估计。

我们在LINEMOD数据集上进行了单目标姿态估计测试，部分结果如图6所示。可以看出，无论是复杂背景还是无纹理目标，本文算法都有出色的姿态估计结果。

图6  部分单目标姿态估计结果（绿色框为真实姿态的目标三维边界框，蓝色框为估计姿态的目标三维边界框）

Fig.6  Qualitative results for single object pose estimation（green 3D bounding boxes represent the ground truth poses， and blue 3D bounding boxes represent our predictions）

为验证本文算法的有效性，我们将本文算法与近期算法进行对比。作为对比的BetaPose^［7］‍、yolo-6d^{［6］‍‍}、PVNet^［9］是两阶段算法，区别在于关键点间接表示形式。BetaPose^［7］与本文相同，使用热图，yolo-6d^［6］直接回归关键点坐标。PVNet^［9］使用图像像素点指向关键点的方向向量作为关键点的间接表达方式，该算法基于目标分割，在网络回归时需要增加单独的图像分割分支。PoseCNN^［16］是端到端算法，可以直接从图像中计算得到姿态。DeepIM^［17］在PoseCNN后增加了迭代优化的后处理部分。ADD（-S）指标测试结果如表1所示，其中eggbox和glue为对称对象，使用ADD（-S）度量，其余对象使用ADD度量。加粗数字为本类别中准确率最高的结果。

由测试结果可以看出，本文算法在无需后优化算法（细化算法，在粗姿态的基础上继续迭代优化）中有着最高的平均准确率，并且在13个类别中的8个类别准确率领先于其他算法。相比于同样使用热图回归关键点的BetaPose，本文算法的准确率提升了16.2%，主要原因在于我们使用的热图回归损失函数HWing Loss可以提升热图回归的质量。本文算法相比于PVNet的平均准确率虽然只领先2.5%，但是PVNet需要额外分割图像，训练收敛困难。并且PVNet使用像素级的方向向量作为关键点的间接表达方式，关键点推理阶段计算复杂度较高。同时，本文算法的准确率相比于PoseCNN提升了26.1%；而PoseCNN使用DeepIM优化后，算法准确率仍然不及本文算法。可以说本文算法有非常优秀的单目标姿态估计准确性。

表2给出了本文算法在LineMod数据集的不同类别下的目标直径dobj ，关键点定位误差ep以及姿态估计误差et、eR，其计算公式如式（15）~（17）所示。

其中p、t、R分别代表关键点定位、姿态平移向量、姿态旋转矩阵的测量值，pˆ、tˆ、Rˆ则分别代表关键点定位、姿态平移向量、姿态旋转矩阵的真实值。从表2中可看出在多数类别下，算法在13个类别的平均关键点定位误差为3.02个像素。不同类别的关键点定位误差主要取决于目标尺寸而与目标形状纹理关系不大，这体现了本文算法在复杂背景、无纹理目标的场景下优秀的关键点定位能力。对于类别内关键点，其定位误差与三维关键点与目标中心的距离相关。

以cat类别为例，图7展示了cat类别的32个图像关键点的定位误差均值、方差与三维关键点偏离目标中心的关系。可以看出随着三维关键点偏离目标中心，图像关键点定位误差的均值与方差都趋向于增大。但是三维更加分散的关键点选取有利于提升PnP算法计算姿态时的鲁棒性，因此从整体上看，由于三维关键点偏离目标中心导致的图像关键点定位误差增大是可以接受的。

图7  图像关键点误差与三维关键点偏离中心距离关系

Fig.7  Relationship between the keypoint error and the distance between the 3D keypoint and the object center

在姿态估计误差上，算法在13个类别的平均平移误差为1.02 cm，平均旋转误差为2.16°。姿态估计误差主要来源于关键点误差，包括关键点的绝对定位误差与关键点中的离群点。其次，数据集中的相机内参与姿态真值的不准确同样会影响最后测量的准确性。

本文在Occlusion LINEMOD 数据集上进行了遮挡目标的姿态估计测试。在目标被部分遮挡的情况下，本文算法可以正确估计目标的姿态，如图8（a）~（d）所示。但是当目标存在大面积遮挡或者或在极端角度时，过少的图像特征导致估计失败，如图8（e）~（h）所示。

图8  部分遮挡目标姿态估计结果（绿色框为真实姿态的目标三维边界框，蓝色框为估计姿态的目标三维边界框）

Fig.8  Qualitative results for occluded object pose estimation（green 3D bounding boxes represent the ground truth poses， and blue 3D bounding boxes represent our predictions）

表3展示了本文算法遮挡目标姿态估计ADD（-S）测试结果与近期其他算法的对比。其中Oberweger等^［18］使用热图作为关键点回归的中间形式，并且利用图像分块提升了算法对遮挡的鲁棒性。其余4个算法与单目标姿态估计对比算法相同。glue使用ADD（-S）度量，其余使用ADD度量。加粗数字为本类别中准确率最高的结果。

由表3可以看出，与单目标姿态估计结果类似，本文算法在无需后优化算法中有着最高的平均准确率。Oberweger等^［18］虽然对遮挡目标进行了优化，但是本文算法的平均准确率仍然高出9.4%。PVNet^［9］基于目标分割，对遮挡目标姿态估计有天然的优势，但其平均准确率比本文低了1.2%，可以看出本文算法的性能优越性。但是相比于后处理算法DeepIM^［17］，本文算法在准确率上有所不足。但是DeepIM在使用中要先利用PoseCNN算法获得姿态初始估计，再进行迭代后优化，复杂的处理流程导致算法无法达到实时处理速度。相比之下，本文算法可以直接通过PnP算法计算得到姿态，更加简洁高效。PoseCNN在使用后处理算法优化后，平均准确率可以提升27.2%，可以说本文算法拥有巨大的准确率提升潜力。

本文算法基于目标检测，因此整体算法运行时间与目标检测网络有关。对于640×480大小的图像，使用Inter I7-9700K CPU与NVIDIA RTX3080 GPU时，用yolo-v5进行目标检测的运行时间约为11 ms。本文算法的热图回归网络前向推理用时约为22 ms，热图推理关键点位置约为5 ms，使用EPnP算法求解姿态用时约为0.3 ms。算法整体用时约为40 ms，最快运行速率约为25 fps，可以用于实时处理。

本文算法同样可以处理多目标姿态估计问题，可以使用自上而下的姿态估计方法，将多目标姿态估计问题转化为多目标检测问题与多个单目标姿态估计。即首先通过yolo-v5网络同时检测图像中的不同目标，获取不同目标的类别标签及目标框。将不同目标分别裁剪后，输入相应类别的姿态估计网络中进行姿态估计。而得益于本文算法将目标检测与姿态估计分离的设计，本文算法在进行多目标姿态估计时可以具有与单目标姿态估计相同的准确性，不会使准确性下降。

为验证本文提出的HWing Loss的有效性，本文进行了不同损失函数的对比实验。使用MSE Loss代替HWing Loss训练热图回归网络，训练参数与4.3节中HWing Loss训练参数相同，同样训练30个轮次，测试结果如表4所示。

可以看出在ADD（-S）指标中，使用HWing Loss相比于使用MSE Loss，平均准确率提升7.1%，每个类别都有所提升。作为对比，表4同时给出了网络分别使用HWing Loss与MSE Loss训练第10轮的实验结果。可以看出，使用HWing Loss训练的网络在训练10轮后的平均准确率已经超越了使用MSE Loss训练30轮的网络。

究其原因是HWing Loss能在训练热图回归网络时，能更关注热图前景像素的误差，并且相比于MSE Loss，HWing Loss在小误差部分有更高的梯度，可以使网络更加快速收敛。

表5给出了使用不同损失函数训练网络预测得到的热图与预期热图的平均像素误差。可以看出，使用HWing Loss训练网络预测热图虽然全局平均误差更大，但是在数量更少且更加重要的前景像素上误差更小。而且误差下降速度更快，10轮训练的结果已经优于使用MSE Loss训练30轮结果，与表4中平均准确率结果一致。由此可以看出，本文提出的HWing Loss相比于MSE Loss更加适合回归热图的训练。