作者：顺达

最近夸克端智能小组在做端上的实时文档检测，即输入一张 RGB 图像，得到文档的四个角的关键点的坐标。整个 pipelines 属于关键点检测算法，因此最近对相关领域的论文进行阅读和进行了实验尝试。

将关键点检测算法按照不同模块进行拆分，可以分成以下几个部分，每个部分都有相关的方法可以进行优化：

1. 图片处理：包括数据光学增强，变换，resize，crop 等操作，扩充图片的多样性；

2. 编码：指的是在训练中，如何将坐标转换成所需要的 label，用于监督模型的输出；

3. 网络模型：指的是网络结构，可以有 backbone/FPN/detection head 等部分组成；

4. 解码：指的是如何将模型推理的结果转换成所需要的坐标形式，如笛卡尔坐标系下的坐标。

Related Works

关键点检测中主要有两条技术方案：

1. 类似人脸检测，模型输出的结果 tensor 通过 fc 层，直接得到一维的向量，通常是归一化后关键点坐标值；

2. 类似人体姿态估计，模型输出的结果 tensor 通过 argmax 等方式，获取 heatmap 中相应大的坐标，最后将此坐标恢复至原图坐标。

近年来，基于 heatmap 来进行关键点检测的方案居多，其主要原因是基于 heatmap 的效果要好于使用全连接层进行回归的方案。所以，我们采用的方案也是基于 heatmap 的，下面是近几年的一些相关论文工作。

DSNT

[1] Nibali A , He Z , Morgan S , et al. Numerical Coordinate Regression with Convolutional Neural Networks[J]. 2018.

思路

目前，在模型输出的 heatmap 到数值坐标的转换中，有两种方式：

1. 通过对 heatmap 中取 argmax，得到相应最大的点，以此来转换成数值坐标。此种方式具有较好的空间泛化性，但是由于在训练中 argmax 是不可导的，通常使用 heatmap 来逼近编码的高斯热例图，这会导致损失函数与最终评价指标的不一致。其次，在推理阶段，只使用到最大响应的坐标点来计算数值坐标，而在训练阶段，所有坐标点都对损失有贡献。第三，通过 heatmap 转换成数值坐标，是会存在理论误差下限的；

2. 通过在 heatmap 后接 fc 层，转换成数值坐标。此种方法让梯度从数值坐标回传到 input 中，但是结果严重依赖与数据分布（例如在训练集中，一个物体一直出现在坐标；而在测试集中，这个物体出现在右边，这样就会导致预测错误）。其次，通过 fc 转换，丢失了 heatmap 的空间信息。

针对上述的两种方案，作者兼容了这两种方案的优点（端到端优化和保持空间泛化性），提出一种可微分的方式来得到数值坐标。

具体步骤

1. 模型的输出 1_K_H*W 个 heatmaps，其中 K 表示关键点的数量；

2. 将每个通道的 heatmap 归一化，让其值都为非负且和为 1，从而得到 norm_heatmap 。这么做的目的是，使用归一化后的 heatmap 保证了预测的坐标位于 heatmap 的空间范围之内。同时， norm_heatmap 也可以理解成二维离散概率密度函数；

3. 生成 X 和 Y 矩阵，$X_{(i,j)}=\frac{2j-(w+1)}{w}$，$Y_{(i,j)}=\frac{2i-(h+1)}{h}$，分别表示 x 轴的索引和 y 轴的索引。可以理解成将图片的左上角缩放到 (-1,-1) 和右下角缩放到 (1,1) ；

4. 将 X 和 Y 矩阵分别与 norm_heatmap 点乘，从而得到最终的数值坐标。这么做的原因是， norm_heatmap 表示概率密度函数， X 矩阵表示索引，两者点成表示预测 x 的均值。通过均值来表示最终的预测的坐标，这样的好处是，a)可微分；b)理论误差下限小。

损失函数 loss

dsnt 模块的损失函数由 Euclidean loss 和 JS 正则约束组成。前者用于回归坐标，后者用于约束生成的热力图更加接近高斯分布。

优点

1. 整套模型是端到端训练的，损失函数与测试指标能对应；

2. 理论误差下限很小；

3. 引入 X 矩阵和 Y 矩阵，可以理解成引入先验，让模型的学习难度降低；

4. 低分辨率的效果依然不错。

缺点

在实验中，发现当关键点位于图片边缘时，预测结果不好。

DARK

[1] Zhang F , Zhu X , Dai H , et al. Distribution-Aware Coordinate Representation for Human Pose Estimation[J]. 2019.

思路

作者发现将 heatmaps 解码结果，对生成最终数值坐标存在较大影响。因此研究了标准的坐标解码方式的不足，提出一种分布已知的解码方式和编码方式，来提高模型的最终效果。

标准的坐标解码过程是，获得模型的 heatmaps 后，通过 argmax 找到最大响应点 m 和第二大响应点 s ，以此来计算最终的响应点 p :

这个公式意味着最大响应点向第二大响应点偏移 0.25 个像素，这么做的目的是补偿量化误差。然后把响应点映射回原图：

这也说明， heatmap 中最大响应点并不是与原图的关键点精确对应，只是大概位置。

基于上面的痛点，作者基于分布已知的前提（高斯分布），提出新的解码方式，解决如何从 heatmap 中得到精确的位置，最小化量化误差。同时，提出了配套的编码方式

具体步骤

解码

假设输出的 heatmap 符合高斯分布，那么 heatmap 就可以用下面函数表示

其中\(\mu\)表示关键点映射到 heatmap 的位置。我们需要求\(\mu\)的位置，因此将函数 g 转换成最大似然函数

对\(P(\mu )\)进行泰勒展开

其中，m 表示在热力图中最大响应的位置。而\(\mu\)在热力图对应的是极点，存在以下性质

结合上述公式，可以得到

因此，为了得到 heatmap 中\(\mu\)的位置，可以通过 heatmap 的一阶导数与二阶导数求得。这步的作用是通过数学的方法来说明该移动距离。

前面提及了假设输出的 heatmap 符合高斯分布，实际情况是不符合的，实际可能是多峰，因此需要对 heatmap 进行调制，让其尽量满足这个前提。具体做法是用高斯核函数来平滑 heatmap ，同时为了保证幅值一致，要进行归一化。

综上所述，步骤是：

1. 对 heatmap 使用高斯核来调制，并且缩放；

2. 求一阶导数和二阶导数，来得到$\mu$；

3. 将$\mu$映射回原图。

编码

编码指的是将关键点映射到 heatmap 上，并且生成高斯分布。

之前工作的做法是现对坐标进行下采样，然后将点进行量化（floor,ceil,round），最后使用量化后的坐标生成高斯函数。

因为量化是不可导的，存在量化误差，因此，作者提出不进行量化，使用 float 来生成高斯函数，这样就能生成无偏 heatmap 。

UDP

[1] Huang J , Zhu Z , Guo F , et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation[J]. 2019.

思路

作者从数据处理和坐标表示下手，以此来提高性能。作者发现，目前的数据处理方式是存在偏差的，特别是 flip 时，会与原数据不对齐；其次坐标表征也存在统计误差。这两个问题共同导致结果存在偏差。因此提出了一种数据处理方式 unbiased data processing，解决图像转换和坐标转换带来的误差。

具体步骤

Unbiased Coordinate System Transformation

在测试中，通常使用翻转后的$k_{o,flip}^{\prime }$与原始的$k_o^{\prime }$进行叠加，来得到最终的预测结果。但是$k_o^{\prime }$与${\hat{k}}_o^{\prime }$并不一致，存在偏差。可以看到翻转后的 heatmap 不与原来的 heatmap 对齐，会产生误差，与分辨率有关。

因此作者建议使用 unit length 来代替图片长度：$w=w^p-1$。这样翻转后的 heatmap 就对齐了。

Unbiased Keypoint Format Transformation

无偏的关键点转换方式应该是可逆的，即$k=Decoding(Enoding(k))$。因此，作者提出了两种方式：

• Combined classification and regression format

借鉴了目标检测中 anchor 的方式，假设需要预测的关键点$k=(m,n)$，则将其转换成如下。其中 C 表示关键点的位置范围，X 和 Y 表示需要预测的 offset。最终解码就是在热力图 C 上取到 argmax，然后对 X 与 Y 的热力图上拿到对应位置的 offset，最后进行相加得到数值坐标。

• Classification format

与 DARK 方式一致，即使用泰勒展开来逼近真实位置。

AID

[1] Huang J , Zhu Z , Huang G , et al. AID: Pushing the Performance Boundary of Human Pose Estimation with Information Dropping Augmentation[J]. 2020.

贡献点

对于关键点检测，外观信息与约束信息同样重要。而以往的工作通常是过拟合外观信息，而忽略了约束信息。因此，本文希望通过 information drop，可理解成掩膜，来强调约束信息。约束信息有利于在该关键点被遮挡时，预测出其准确位置。

而以往工作没有使用到 information drop 的原因是，使用该数据增强手段后指标下降。作者就通过实验，发现 information drop 是有助于提高模型精度的，但需要修改响应的训练策略：

1. 加倍训练次数；

2. 先使用没有 mask 的来训练，得到一个比较好的模型后，再把 mask 手段加入继续训练。

RSN

[1] Cai Y , Wang Z , Luo Z , et al. Learning Delicate Local Representations for Multi-Person Pose Estimation[J]. 2020.

贡献点

本文是 2019 年 coco 关键点检测冠军的方案。其本文的主要思想是，最大程度聚合具有相同空间尺寸的特征，以此来获得丰富的局部信息，局部信息有利于产生更加准确的位置。因此提出了 RSN 网络，如下图一所示。从图来看，即融合不同感受野特征。RSN 的输出包含了 low-level 准确的空间信息与 high-level 语义信息，空间信息有助于定位，语义信息有助于分类。但是这两类信息给最终预测带来的影响权重是不一致的，需要使用到 PRM 模块来平衡，RPM 模块本质就是一个通道注意力和空间注意力模块。

Lite-HRNet

[1] Yu C , Xiao B , Gao C , et al. Lite-HRNet: A Lightweight High-Resolution Network[J]. 2021.

贡献点

本文出了一个高效的高分辨率网络，是 HRNet 的轻量化版本，通过将 ShuffleNet 中的 shuffle block 引入到 HRNet 中。同时发现 shuffleNet 中大量使用了 pointwise convolution（1_1 卷积），是计算瓶颈，因此引入 contional channel weight 来取代 shuffle block 中的 1_1 卷积。网络的整体结构如下图所示。在模型中一致保留高分辨率特征，并不断融合 high-level 特征。

在前面提及到的 contional channel weight 如下所示。左边是 ShuffleNet 中的 shuffle block，右图是 contional channel weight。可以看到，采用新模块来取代了 1*1 的卷积，实现跨 stage 信息交流与局部信息交流。其具体做法包含了 Cross-resolution weight computation 和 Spatial weight computation。这两个模块的本质是注意力机制。

实验优化结果

模型结构

本次模型借鉴了 CenterNet/RetiaFace/DBFace 中的相关工作。本次的使用了 dsnt 的方案。主要原因是：需要运行在端上，实时性是首要考虑因素。dsnt 在低分辨率的优势明显。

MobileNet v3 使用 small 版本，FPN 中使用 Nearest Upsample + conv + bn + Relu 来进行上采样。在训练时使用了 keypoints ， mask 和 center 分支；而预测时，只使用到了 keypoints 分支。

优化策略

在本次实验中，使用到了以下几种优化策略：

1. 使用 mask 与 center 分支来辅助学习。其中 mask 表示文档的掩膜，center 表示文档的中心点；

2. 使用 deep Supervise。使用 4 倍下采样特征图与 8 倍下采样特征图来进行训练，使用相同的 loss 函数来监督这两层；

3. dsnt 中对边缘点的效果不佳，因此，对图片进行 padding，让点不再位于图片边缘；

4. 数据增强策略，除了常规的光学扰动增强外，还对图片进行 random crop、random erase 和 random flip 等操作；

5. 进行 loss 函数尝试工作，对于关键点分支的 loss，尝试过 euclidean loss ， l1 loss ， l2 loss 和 smoothl1 loss ，最终 smoothl1 loss 的效果最佳。

评价指标

• MSE

用于在训练中评价验证集的均方误差。

• oks-mAP

oks 用于评估预测与真实关键点之间的相似度，mAP 的评估方式类似 coco[0.5:0.05:0.95]的评价方式，这里取[0.99:0.001:0.999]。其中，oks 进行一定变换，\(d_{p,i}\)表示点的欧式距离，\(S_p\)表示该四边形的面积。

• 耗时

耗时指的是在红米 8 上，用 MNN 推理框架跑模型的平均时间。

实验结果

先构建一个 baseline，baseline 的模型为 moblieNet v3 + fpn + ssh module + keypoints 分支 + dsnt ，其中，都没有使用上述优化策略，使用 4 倍下采样特征图作为输出。

在 v2 版本中替换不同的 loss 函数。

此外，还尝试过其他无效的 tricks：

1. 辅助任务有利于提高模型的指标，因此还加入了 edge 的分支来辅助学习。实验下来，加入该分支反而损害模型的指标。可能原因是 edge 是利用 gt 关键点来生成的，可能某些 edge 并不是对应文档真正的边缘；

2. 现阶段是预测文档的 4 个角，因此在增加 4 个点来进行预测，分别是 4 条边的中心点，所以模型一共预测 8 个关键点。实验结果显示，指标也下降了。

Demo 演示

Demo 视频请点击此文章尾部查看。

总结

综上所述，在端上的文档关键点检测领域中，目前尝试下来，是基于 heatmap+dsnt 的方案较优，oks-mAP 的指标有提升空间。但是，对比使用 fc 层进行回归坐标的方式，基于 heatmap 的方案中存在一个不足是：无法根据文档的约束信息，来预测图片外的关键点坐标。此方案的不足，会导致文档内容的缺失，摆正效果不佳的情况。因此，后续需要弥补此不足。

如何高效开发端智能算法？MNN 工作台 Python 调试详解

关注我们，每周 3 篇移动技术实践 &干货给你思考！