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

語言: CN / TW / HK

作者:順達

最近夸克端智慧小組在做端上的實時文件檢測,即輸入一張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 篇移動技術實踐&乾貨給你思考!