夸克APP端智能：文档关键点检测实践与应用

作者：顺达

最近夸克端智能小组在做端上的实时文档检测，即输入一张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. 模型的输出1KH*W个heatmaps，其中K表示关键点的数量；
2. 将每个通道的heatmap归一化，让其值都为非负且和为1，从而得到 norm_heatmap 。这么做的目的是，使用归一化后的heatmap保证了预测的坐标位于heatmap的空间范围之内。同时， norm_heatmap 也可以理解成二维离散概率密度函数；
3. 生成 X 和 Y 矩阵，，，分别表示x轴的索引和y轴的索引。可以理解成将图片的左上角缩放到 (-1,-1) 和右下角缩放到 (1,1) ；
4. 将X 和 Y 矩阵分别与 norm_heatmap 点乘，从而得到最终的数值坐标。这么做的原因是， norm_heatmap 表示概率密度函数， X 矩阵表示索引，两者点成表示预测x的均值。通过均值来表示最终的预测的坐标，这样的好处是，a)可微分；b)理论误差下限小。