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

2022-07-07林林王延杰孙海超

液晶与显示 2022年7期

林林王延杰*孙海超

林林1，2，王延杰1*，孙海超1

（1.中国科学院长春光学精密机械与物理研究所，吉林长春 130033；2.中国科学院大学，北京 100049）

针对传统热图回归使用的均方误差（MSE）损失函数训练热图回归网络的精度不高且训练缓慢的问题，本文提出了用于热图回归的损失函数Heatmap Wing Loss（HWing Loss）。该损失函数对于不同的像素值有不同的损失函数值，前景像素的损失函数梯度更大，可以使网络更加关注前景像素，使热图回归更加准确快速。同时根据热图分布特性，使用基于高斯分布的关键点推理方法减小热图推断关键点时的量化误差。以此两点为基础，构造新的基于关键点定位的单目标姿态估计的算法。实验结果表明，相比于使用MSE Loss的算法，使用HWing Loss的姿态估计算法有更高的ADD（-S）准确率，在LINEMOD数据集上达到了88.8%，性能优于近期其他的基于深度学习的姿态估计算法。本文算法在RTX3080 GPU上最快能以25 fps的速度运行，兼具速度与性能优势。

深度学习；姿态估计；损失函数；热图

1 引言

随着计算机视觉技术的不断发展，仅从图像中检测物体已经不能满足智能机器人对环境感知的需要，基于图像的物体6D姿态测量成为了新的研究热点。用RGB-D图像估计物体的姿态虽然精度高，但是计算复杂难以应用。使用单RGB图像配合目标的三维模型可以解决单目相机成像时的尺度不确定问题，也可以得到较好的结果，但在实际应用场景下，目标的背景较为复杂，而且会有很多遮挡现象出现，这都给姿态估计问题造成了很大的影响。因此，如何从单张图像中快速准确地在复杂场景下获得目标的6D姿态是一个非常值得深入研究的课题。

传统姿态估计算法通常使用手工特征［1-3］建立图像模板与图像之间的对应关系，很难处理无纹理对象，速度也较慢。随着深度学习的发展，基于CNN的姿态估计算法取得了优异的性能。Deep-6DPose［4］、AAE［5］等端到端算法是将图像输入到神经网络中直接输出目标姿态，但是这种方法的泛化能力并不好，网络无法学习到足够的特征以表达目标姿态。近期研究的两阶段算法如yolo-6d［6］、BetaPose［7］、Pix2Pose［8］、PVNet［9］等在精度上相比于端到端算法展现了较大的优势。其中BetaPose［7］为这类方法提供了一个简单高效的基础架构，即先使用神经网络回归热图定位关键点，再使用PP算法计算目标姿态。

热图是一种特殊的图像，图像上每个像素值代表关键点出现在该位置的概率值，可以用于神经网络回归关键点时的中间表示，利用预测得到的热图极大值及其邻域推断得到关键点位置。因此，热图的前景像素的预测精度对于关键点定位任务至关重要，这个部分即使出现了很小的预测误差，也会导致很大的关键点偏离。相反，预测热图的背景像素值并不重要，预测时只需使其逐步趋向于零值即可。而BetaPose并没有利用好热图的相关特性：首先，在回归热图时，使用传统的MSE Loss损失函数。MSE Loss对于不同位置的像素值没有区分，训练被大量无意义的背景像素主导，导致回归的热图精度不高。其次，在使用预测的热图定位关键点时，只简单地取最大值点，存在量化误差。这两点影响了算法最终的结果。

本文对BetaPose进行了改进。首先针对MSE Loss损失函数对像素没有区分的问题，提出适用于热图回归的损失函数Heatmap Wing Loss（HWing Loss），该损失函数更加关注前景像素误差，可以有效提高热图回归的质量。其次，使用基于高斯分布的关键点推理方法，减小热图推断关键点时的量化误差。以此两点为基础，构造了新的单目标姿态估计算法。实验结果表明，相比于其他姿态估计算法，本文算法在LINEMOD数据集上有更高的ADD（-S）准确率。

2 适用于热图回归的损失函数

2.1　热图以及传统损失函数

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

图1　关键点热图

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

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

.（2）

2.2　Heatmap Wing Loss

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

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

图3　HWing Loss示意图

3 目标姿态估计算法

本文姿态估计算法为两阶段算法：第一阶段预测热图，用于定位目标关键点；第二阶段由关键点计算目标姿态。具体流程如图4所示。首先将输入图像裁剪为×大小，输入热图回归网络，得到输出的热图，由预测热图推理关键点位置，最后使用PP算法计算得到姿态。

图4　姿态估计算法流程图

3.1　热图回归网络

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

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

图5　加权损失示意图

3.2　关键点推断

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

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

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

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

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

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

同理，对方向坐标有：

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

3.3　姿态计算

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

4 实验与结果

4.1　数据集

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

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

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

4.2　性能评估

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

4.3　实验参数

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

4.4　单目标姿态估计

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

为验证本文算法的有效性，我们将本文算法与近期算法进行对比。作为对比的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度量。加粗数字为本类别中准确率最高的结果。

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

表1单目标姿态估计ADD（-S）指标测试结果

Tab.1　Single object pose estimation ADD（-S） metric test results

*对称对象

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

表2　关键点定位误差与姿态估计误差

*对称对象

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

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

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

4.5　遮挡目标姿态估计

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

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

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

表3遮挡目标姿态估计ADD（-S）指标测试结果

Tab.3　Occluded object pose estimation ADD（-S） metric test results

*对称对象

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

4.6　运行时间实验

4.7　多目标姿态估计问题

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

4.8　损失函数对比实验

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

表4对比实验ADD（-S）指标测试对比结果

Tab.4　Comparison results of comparative experiment ADD （-S） metric test

*对称对象

可以看出在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更加适合回归热图的训练。

表5不同损失函数对预测热图误差的影响

Tab.5　Influence of different loss function on the error of predicting HeatMap

5 结论

本文分析了用于关键点回归热图的性质，并指出MSE Loss不利于热图回归的问题。为了解决这个问题，本文提出适用于热图回归的Heatmap Wing Loss。同时利用热图性质，改进热图推断关键点的方法。以此为基础，改进基于关键点定位的单目标姿态估计的算法，经实验验证，本文的单目标姿态估计算法在LINEMOD数据集上的ADD（-S）指标平均准确率达到了88.8%，相比于近期其他算法有更好的姿态估计准确率。算法运行速率最快可达到25 fps，适用于实时处理。

［1］ RUBLEE E， RABAUD V， KONOLIGE K，. ORB： an efficient alternative to SIFT or SURF［C］//2011Barcelona： IEEE， 2011： 2564-2571.

［2］ BAY H， ESS A， TUYTELAARS T，. Speeded-up robust features （SURF）［J］.， 2008， 110（3）： 346-359.

［3］丁南南，刘艳滢，朱明．尺度相互作用墨西哥帽小波提取图像特征点［J］．液晶与显示，2012，27（1）：125-129．

DING N N， LIU Y Y， ZHU M. Extracting image feature points using scale-interaction of mexican-hat wavelets［J］.， 2012， 27（1）： 125-129. （in Chinese）

［4］ DO T T， CAI M， PHAM T，. Deep-6DPose： recovering 6D object pose from a single RGB image［EB/OL］. （2018-02-28）. https：//arxiv.org/abs/1802.10367v1.

［5］ SUNDERMEYER M， MARTON Z C， DURNER M，. Augmented autoencoders： implicit 3D orientation learning for 6D object detection［J］.， 2020， 128（3）： 714-729.

［6］ TEKIN B， SINHA S N， FUA P，. Real-time seamless single shot 6D object pose prediction［C］//2018. Salt Lake City： IEEE， 2018： 292-301.

［7］ ZHAO Z L， PENG G， WANG H Y，. Estimating 6D pose from localizing designated surface keypoints［EB/OL］. （2018-12-04）. https：//arxiv.org/abs/1812.01387.

［8］ PARK K， PATTEN T， VINCZE M. Pix2Pose： pixel-wise coordinate regression of objects for 6D pose estimation［C］/2019. Seoul： IEEE， 2019： 7667-7676.

［9］ PENG S D， LIU Y， HUANG Q X，. PVNet： pixel-wise voting network for 6DoF pose estimation［C］//2019. Long Beach： IEEE， 2019： 4556-4565.

［10］ WANG X Y， BO L F， LI F X. Adaptive wing loss for robust face alignmentheatmap regression［C］//2019. Seoul： IEEE， 2019： 6970-6980.

［11］ SUN K， XIAO B， LIU D，. Deep high-resolution representation learning for human pose estimation［C］//Long Beach： IEEE， 2019： 5693-5703.

［12］ WANG J D， SUN K， CHENG T H，. Deep high-resolution representation learning for visual recognition［J］.， 2021， 43（10）： 3349-3364.

［13］ ZHANG F， ZHU X T， DAI H B，. Distribution-aware coordinate representation for human pose estimation［C］//2020. Seattle： IEEE， 2020： 7091-7100.

［14］ LEPETIT V， MORENO-NOGUER F， FUA P. EPP： an accurate（） solution to the PP problem［J］.， 2009， 81（2）： 155-166.

［15］ XIAO J X， HAYS J， EHINGER K A，. SUN database： large-scale scene recognition from abbey to zoo［C］//2010. San Francisco： IEEE， 2010： 3485-3492.

［16］ XIANG Y， SCHMIDT T， NARAYANAN V，. PoseCNN： a convolutional neural network for 6D object pose estimation in cluttered scenes［C］/14Pittsburgh：IEEE， 2018.

［17］ LI Y， WANG G， JI X Y，. DeepIM： deep iterative matching for 6D pose estimation［J］.， 2020， 128（3）： 657-678.

［18］ OBERWEGER M， RAD M， LEPETIT V. Making deep heatmaps robust to partial occlusions for 3D object pose estimation［C］//15Munich： Springer， 2018： 125-141.

Object 6D pose estimation algorithm based on improved heatmap loss function

LIN Lin1，2，WANG Yan-jie1*，SUN Hai-chao1

（1，，，130033，；2，100049，）

In view of the problem of low precision and slow training of heatmap regression network trained by mean square error （MSE） loss function used in traditional heatmap regression， the loss function Heatmap Wing Loss （HWing Loss） for heatmap regression is proposed in this thesis. In terms of different pixel values， the loss function has different loss function values， and the loss function gradient of foreground pixels is larger， which can make the network focus more on the foreground pixels and make the heatmap regression more accurate and faster. In line with the distribution characteristics of the heatmap， the keypoint inference method based on the Gaussian distribution is adopted in this thesis to reduce the quantization error when the heatmap infers the keypoints. By taking the two points as the basis， it constructs a new monocular pose estimation algorithm based on keypoint positioning. According to the experiments， in contrast with the algorithm using MSE Loss， the pose estimation algorithm using HWing Loss has a higher ADD（-S） accuracy rate， which reaches 88.8% on the LINEMOD dataset. Meanwhile， the performance is better than other recent pose estimation algorithms based on deep learning. The algorithm in this thesis can run at the fastest speed of 25 fps on RTX3080 GPU， in which the high speed and performance can be both embodied.

deep learning； pose estimation； loss function； heatmap

TP391.4

10.37188/CJLCD.2021-0317

1007-2780（2022）07-0913-11

2021-12-03；

2022-01-09.

吉林省科技发展计划（No.20210201132GX）

Supported by Department of Science and Technology of Jilin Province （No.20210201132GX）

，E-mail： wangyj@ciomp.ac.cn

林林（1997—），男，内蒙古赤峰人，硕士研究生，2019年于中国科学技术大学获得学士学位，主要从事计算机视觉方面的研究。E-mail：linlin19@mails.ucas.ac.cn