夸克APP端智慧：文件關鍵點檢測實踐與應用

阿里巴巴移動技術發表於2021-09-27
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_{(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正則約束組成。前者用於迴歸座標，後者用於約束生成的熱力圖更加接近高斯分佈。

$$ L_{euc}(u,p) = ||p-u||_2 $$

$$ L_D(Z,p) =JS(p(c)||N(p,I))) $$

優點

  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 :

$$ p=m+0.25\frac{s-m}{\left \| s-m \right \|_2} $$

這個公式意味著最大響應點向第二大響應點偏移0.25個畫素，這麼做的目的是補償量化誤差。然後把響應點對映回原圖：

$$ \hat{p} = \lambda p $$

這也說明， heatmap 中最大響應點並不是與原圖的關鍵點精確對應，只是大概位置。

基於上面的痛點，作者基於分佈已知的前提（高斯分佈），提出新的解碼方式，解決如何從 heatmap 中得到精確的位置，最小化量化誤差。同時，提出了配套的編碼方式

具體步驟

解碼

假設輸出的 heatmap 符合高斯分佈，那麼 heatmap 就可以用下面函式表示

其中\( \mu \)表示關鍵點對映到 heatmap 的位置。我們需要求\( \mu \)的位置，因此將函式g轉換成最大似然函式

對\( P(\mu ) \)進行泰勒展開

其中，m表示在熱力圖中最大響應的位置。而\( \mu \)在熱力圖對應的是極點，存在以下性質

結合上述公式，可以得到

因此，為了得到 heatmap 中\( \mu \)的位置，可以通過 heatmap 的一階導數與二階導數求得。這步的作用是通過數學的方法來說明該移動距離。

前面提及了假設輸出的 heatmap 符合高斯分佈，實際情況是不符合的，實際可能是多峰，因此需要對 heatmap 進行調製，讓其儘量滿足這個前提。具體做法是用高斯核函式來平滑 heatmap ，同時為了保證幅值一致，要進行歸一化。

$$ {h}'=K\circledast h $$

$$ {h}'=\frac{{h}'-min({h}')}{max({h}')-min({h}')}*max(h) $$

綜上所述，步驟是：

  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} \)與原始的\( {k}'_o \)進行疊加，來得到最終的預測結果。但是\( {k}'_o \)與\( {\hat{k}}'_o \)並不一致，存在偏差。可以看到翻轉後的 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（11卷積），是計算瓶頸，因此引入contional channel weight來取代shuffle block中的11卷積。網路的整體結構如下圖所示。在模型中一致保留高解析度特徵，並不斷融合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

用於在訓練中評價驗證集的均方誤差。

$$ mse = \frac{\sum |d_i - \hat{d_i}|_2^2}{N} $$

  • oks-mAP

oks用於評估預測與真實關鍵點之間的相似度，mAP的評估方式類似coco[0.5:0.05:0.95]的評價方式，這裡取[0.99:0.001:0.999]。其中，oks進行一定變換，\( d_{p,i} \)表示點的歐式距離，$S_p$表示該四邊形的面積。

$$ oks_{p,i} =e^{-\frac{d_{p,i}^2}{2S_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 篇移動技術實踐&乾貨給你思考！

相關文章