聊聊人像摳圖背後的演算法技術

語言: CN / TW / HK
本文分享自華為雲社群《 人像摳圖:演算法概述及工程實現(一)》,原文作者:杜甫蓋房子 。

本文將從演算法概述、工程實現、優化改進三個方面闡述如何實現一個實時、優雅、精確的視訊人像摳圖專案。

什麼是摳圖

對於一張圖I, 我們感興趣的人像部分稱為前景F,其餘部分為背景B,則影象I可以視為F與B的加權融合:I = alpha * F + (1 - alpha) * BI=alphaF+(1−alpha)∗B,而摳圖任務就是找到合適的權重alpha。值得一提的是,如圖,檢視摳圖ground truth可以看到,alpha是[0, 1]之間的連續值,可以理解為畫素屬於前景的概率,這與人像分割是不同的。如圖,在人像分割任務中,alpha只能取0或1,本質上是分類任務,而摳圖是迴歸任務。

摳圖ground truth:

分割ground truth:

相關工作

我們主要關注比較有代表性的基於深度學習的摳圖演算法。目前流行的摳圖演算法大致可以分為兩類,一種是需要先驗資訊的Trimap-based的方法,寬泛的先驗資訊包括Trimap、粗糙mask、無人的背景影象、Pose資訊等,網路使用先驗資訊與圖片資訊共同預測alpha;另一種則是Trimap-free的方法,僅根據圖片資訊預測alpha​,對實際應用更友好,但效果普遍不如Trimap-based的方法。

Trimap-based

Trimap是最常用的先驗知識,顧名思義Trimap是一個三元圖,每個畫素取值為{0,128,255}其中之一,分別代表前景、未知與背景,如圖。

Deep Image Matting

多數摳圖演算法採用了Trimap作為先驗知識。Adobe在17年提出了Deep Image Matting[^1],這是首個端到端預測alpha的演算法,整個模型分Matting encoder-decoder stage與Matting refinement stage兩個部分,Matting encoder-decoder stage是第一部分,根據輸入影象與對應的Trimap,得到較為粗略的alpha matte。Matting refinement stage是一個小的卷積網路,用來提升alpha matte的精度與邊緣表現。

本文在當時達到了state-of-the-art,後續很多文章都沿用了這種“粗略-精細”的摳圖思路,此外,由於標註成本高,過去摳圖任務的資料是非常有限的。本文還通過合成提出了一個大資料集Composition-1K,將精細標註的前景與不同背景融合,得到了45500訓練影象和1000測試影象,大大豐富了摳圖任務的資料。

Background Matting

Background Matting[^2]是華盛頓大學提出的摳圖演算法,後續釋出了Backgroun MattingV2,方法比較有創新點,並且在實際工程應用中取得了不錯的效果。

同時,由於Adobe的資料都是基於合成的,為了更好的適應真實輸入,文中提出一個自監督網路訓練G_{Real}GReal​來對未標註的真實輸入進行學習。G_{Real}GReal​輸入與G_{Adobe}GAdobe​相同,用G_{Adobe}GAdobe​輸出的alpha matte與F來監督G_{Real}GReal​的輸出得到loss,此外,G_{Real}GReal​的輸出合成得到的RGB還將通過一個鑑別器來判斷真偽得到第二個loss,共同訓練G_{Real}GReal​。

文中列舉了一些使用手機拍攝得到的測試結果,可以看到大部分情況結果還是很不錯的。

Background Matting V2

Background Matting得到了不錯的效果,但該專案無法實時執行,也無法很好的處理高解析度輸入。所以專案團隊又推出了Background Matting V2[^3],該專案可以以30fps的速度在4k輸入上得到不錯的結果。

文章實現高效高解析度摳圖的一個重要想法是,alpha matte中大部分畫素是0或1,只有少量的區域包含過渡畫素。因此文章將網路分為base網路和refine網路,base網路對低解析度影象進行處理,refine網路根據base網路的處理結果選擇原始高解析度影象上特定影象塊進行處理。

base網路輸入為c倍下采樣的影象與背景,通過encoder-decoder輸出粗略的alpha matte、F、error map與hidden features。將取樣c倍得到的error map E_cEc​上取樣到原始解析度的\frac{1}{4}41​為E_4E4​,則E_4E4​每個畫素對應原圖4x4影象塊,從E_4E4​選擇topk error畫素,即為原始topk error 4x4影象塊。在選擇出的畫素周圍裁剪出多個8x8影象塊送入refine網路。refine網路是一個two-stage網路,首先將輸入通過部分CBR操作得到第一階段輸出,與原始輸入中提取的8x8影象塊cat後輸入第二階段,最後將refine後的影象塊與base得到的結果交換得到最終的alpha matte和F。

此外文章還發布了兩個資料集:視訊摳圖資料集VideoMatte240K與影象摳圖資料集PhotoMatte13K/85。VideoMatte240K收集了484個高解析度視訊,使用Chroma-key軟體生成了240000+前景和alpha matte對。PhotoMatte13K/85則是在良好光照下拍攝照片使用軟體和手工調整的方法得到13000+前景與alpha matte資料對。大型資料集同樣是本文的重要貢獻之一。

此外還有一些文章如Inductive Guided Filter[^4]、MGMatting[^5]等,使用粗略的mask作為先驗資訊預測alpha matte,在應用時也比trimap友好很多。MGMatting同時也提出了一個有636張精確標註人像的摳圖資料集RealWorldPortrait-636,可以通過合成等資料增廣方法擴充套件使用。

Trimap-free

實際應用中先驗資訊獲取起來是很不方便的,一些文章將先驗資訊獲取的部分也放在網路中進行。

Semantic Human Matting

阿里巴巴提出的Semantic Human Matting[^6]同樣分解了摳圖任務,網路分為三個部分,T-Net對畫素三分類得到Trimap,與影象concat得到六通道輸入送入M-Net,M-Net通過encoder-decoder得到較為粗糙的alpha matte,最後將T-Net與M-Net的輸出送入融合模組Fusion Module,最終得到更精確的alpha matte。

網路訓練時的alpha loss分為alpha loss與compositional loss,與DIM類似,此外還加入了畫素分類lossL_tLt​,最終loss為:L = L_p + L_t=L_\alpha + L_c + L_tL=Lp​+Lt​=​+Lc​+Lt​。文章實現了端到端Trimap-free的摳圖演算法,但較為臃腫。此外文章提出Fashion Model資料集,從電商網站收集整理了35000+標註的圖片,但並沒有開放。

Modnet

modnet[^7]認為神經網路更擅長學習單一任務,所以將摳圖任務分為三個子任務,分別進行顯式監督訓練和同步優化,最終可以以63fps在512x512輸入下達到soft結果,因此在後續的工程實現中我也選擇了modnet作為Baseline。

網路的三個子任務分別是Semantic Estimation、Detail Prediction和Semantic-Detail Fusion,Semantic Estimation部分由backbone與decoder組成,輸出相對於輸入下采樣16倍的semantics,用來提供語義資訊,此任務的ground truth是標註的alpha經過下采樣與高斯濾波得到的。 Detail Prediction任務輸入有三個:原始影象、semantic分支的中間特徵以及S分支的輸出S_pSp​,D分支同樣是encoder-decoder結構,值得留意的該分支的loss,由於D分支只關注細節特徵,所以通過ground truth alpha生成trimap,只在trimap的unknown區域計算d_pdp​與\alpha_gαg​的L_1L1​損失。F分支對語義資訊與細節預測進行融合,得到最終的alpha matte與ground truth計算L_1L1​損失,網路訓練的總損失為:L=\lambda_sL_s + \lambda_dL_d+\lambda_{\alpha}L_{\alpha}L=λsLs​+λdLd​+λα​。

最後,文章還提出了一種使視訊結果在時間上更平滑的後處理方式OFD,在前後兩幀較為相似而中間幀與前後兩幀距離較大時,使用前後幀的平均值平滑中間幀,但該方法會導致實際結果比輸入延遲一幀。

此外,U^2U2-Net、SIM等網路可以對影象進行顯著性摳圖,感興趣的話可以關注一下。

資料集

評價指標

常用的客觀評價指標來自於2009年CVPR一篇論文[^8],主要有:

此外,可以在paperwithcode上檢視Image Matting任務的相關文章,在Alpha Matting網站上檢視一些演算法的evaluation指標。

本專案的最終目的是在HiLens Kit硬體上落地實現實時視訊讀入與背景替換,開發環境為HiLens配套線上開發環境HiLens Studio,先上一下對比baseline的改進效果:

使用modnet預訓練模型modnet_photographic_portrait_matting.ckpt進行測試結果如下:

可以看到由於場景較為陌生、逆光等原因會導致摳圖結果有些閃爍,雖然modnet可以針對特定視訊進行自監督finetune,但我們的目的是在普遍意義上效果更好,因此沒有對本視訊進行自監督學習。

優化後的模型效果如下:

本視訊並沒有作為訓練資料。可以看到,摳圖的閃爍情況減少了很多,毛髮等細節也基本沒有損失。

工程落地

為了測試baseline效果,首先我們要在使用場景下對baseline進行工程落地。根據文件匯入/轉換本地開發模型可知

昇騰310 AI處理器支援模型格式為".om",對於Pytorch模型來說可以通過"Pytorch->Caffe->om"或"Pytorch->onnx->om"(新版本)的轉換方式得到,這裡我選擇的是第一種。Pytorch->Caffe模型轉換方法與注意事項在之前的部落格中有具體闡述過,這裡不贅述。轉換得到Caffe模型後,可以在HiLens Studio中直接轉為om模型,非常方便。

首先在HiLens Studio中新建一個技能,此處選擇了空模板,只需要修改一下技能名稱就可以。

將Caffe模型上傳到model資料夾下:

在控制檯中執行模型轉換命令即可得到可以執行的om模型:

/opt/ddk/bin/aarch64-linux-gcc7.3.0/omg --model=./modnet_portrait_320.prototxt --weight=./modnet_portrait_320.caffemodel --framework=0 --output=./modnet_portrait_320 --insert_op_conf=./aipp.cfg

接下來完善demo程式碼。在測試時HiLens Studio可以在工具欄選擇使用視訊模擬攝像頭輸入,或連線手機使用手機進行測試:

具體的demo程式碼如下:

# -*- coding: utf-8 -*-
# !/usr/bin/python3
# HiLens Framework 0.2.2 python demo
​
import cv2
import os
import hilens
import numpy as np
from utils import preprocess
import time
​
​
def run(work_path):
    hilens.init("hello")  # 與建立技能時的校驗值一致
​
    camera = hilens.VideoCapture('test/camera0_2.mp4')  # 模擬輸入的視訊路徑
    display = hilens.Display(hilens.HDMI)
​
    # 初始化模型
    model_path = os.path.join(work_path, 'model/modnet_portrait_320.om') # 模型路徑
    model = hilens.Model(model_path)
​
    while True:
        try:
            input_yuv = camera.read()
            input_rgb = cv2.cvtColor(input_yuv, cv2.COLOR_YUV2RGB_NV21)
            # 摳圖後替換的背景
            bg_img = cv2.cvtColor(cv2.imread('data/tiantan.jpg'), cv2.COLOR_BGR2RGB) 
            crop_img, input_img = preprocess(input_rgb)  # 預處理
            s = time.time()
            matte_tensor = model.infer([input_img.flatten()])[0]
            print('infer time:', time.time() - s)
            matte_tensor = matte_tensor.reshape(1, 1, 384, 384)
​
            alpha_t = matte_tensor[0].transpose(1, 2, 0)
            matte_np = cv2.resize(np.tile(alpha_t, (1, 1, 3)), (640, 640))
            fg_np = matte_np * crop_img + (1 - matte_np) * bg_img  # 替換背景
            view_np = np.uint8(np.concatenate((crop_img, fg_np), axis=1))
            print('all time:', time.time() - s)
​
            output_nv21 = hilens.cvt_color(view_np, hilens.RGB2YUV_NV21)
            display.show(output_nv21)
​
        except Exception as e:
            print(e)
            break
​
    hilens.terminate()

其中預處理部分的程式碼為:

import cv2
import numpy as np
​
​
TARGET_SIZE = 640
MODEL_SIZE = 384
​
​
def preprocess(ori_img):
    ori_img = cv2.flip(ori_img, 1)
    H, W, C = ori_img.shape
    x_start = max((W - min(H, W)) // 2, 0)
    y_start = max((H - min(H, W)) // 2, 0)
    crop_img = ori_img[y_start: y_start + min(H, W), x_start: x_start + min(H, W)]
    crop_img = cv2.resize(crop_img, (TARGET_SIZE, TARGET_SIZE))
    input_img = cv2.resize(crop_img, (MODEL_SIZE, MODEL_SIZE))
​
    return crop_img, input_img

demo部分的程式碼非常簡單,點選執行即可在模擬器中看到效果:

模型推理耗時44ms左右,端到端執行耗時60ms左右,達到了我們想要的實時的效果。

效果改進

預訓練模型在工程上存在著時序閃爍的問題,原論文中提出了一種使視訊結果在時間上更平滑的後處理方式OFD,即用前後兩幀平均誤差大的中間幀。但這種辦法只適合慢速運動,同時會導致一幀延遲,而我們希望可以對攝像頭輸入進行實時、普適的時序處理,因此OFD不適合我們的應用場景。

在Video Object Segmentation任務中有一些基於Memory Network的方法(如STM),摳圖領域也有新論文如DVM考慮引入時序記憶單元使摳圖結果在時序上更穩定,但這些方法普遍需要前後n幀資訊,在資源佔用、推理實時性、適用場景上都與我們希望的場景不符合。

考慮到資源消耗與效果的平衡,我們採用將前一幀的alpha結果cat到當前幀RGB影象後共同作為輸入的方法來使網路在時序上更穩定。

網路上的修改非常簡單,只需在模型初始化時指定in_channels = 4:

modnet = MODNet(in_channels=4, backbone_pretrained=False)

訓練資料方面,我們選擇一些VideoMatting的資料集:VideoMatte240K、ConferenceVideoSegmentationDataset。

最初,我們嘗試將前一幀alpha作為輸入、缺失前幀時補零這種簡單的策略對模型進行訓練:

if os.path.exists(os.path.join(self.alpha_path, alpha_pre_path)):
    alpha_pre = cv2.imread(os.path.join(self.alpha_path, alpha_pre_path))
else:
    alpha_pre = np.zeros_like(alpha)
 
net_input = torch.cat([image, alpha_pre], dim=0)

收斂部署後發現,在場景比較穩定時模型效果提升較大,而在人進、出畫面時模型適應較差,同時如果某一幀結果較差,將對後續幀產生很大影響。針對這些問題,考慮制定相應的資料增強的策略來解決問題。

  • 人進、出畫面時模型適應較差:資料集中空白幀較少,對人物入畫出畫學習不夠,因此在資料處理時增加空白幀概率:
if os.path.exists(os.path.join(self.alpha_path, alpha_pre_path)) and random.random() < 0.7:
    alpha_pre = cv2.imread(os.path.join(self.alpha_path, alpha_pre_path))
else:
    alpha_pre = np.zeros_like(alpha)
  • 某一幀結果較差,將對後續幀產生很大影響:目前的結果較為依賴前一幀alpha,沒有學會拋棄錯誤結果,因此在資料處理時對alpha_pre進行一定概率的仿射變換,使網路學會忽略偏差較大的結果;
  • 此外,光照問題仍然存在,在背光或光線較強處摳圖效果較差:對影象進行光照增強,具體的,一定概率情況下模擬點光源或線光源疊加到原圖中,使網路對光照更魯棒。光照資料增強有兩種比較常用的方式,一種是通過opencv進行簡單的模擬,具體可以參考augmentation.py,另外還有通過GAN生成資料,我們使用opencv進行模擬。

重新訓練後,我們的模型效果已經可以達到前文展示的效果,在16T算力的HiLens Kit上完全達到了實時、優雅的效果。進一步的,我還想要模型成為耗時更少、效果更好的優秀模型~目前在做的提升方向是:

  • 更換backbone:針對應用硬體選擇合適的backbone一向是提升模型價效比最高的方法,直接根據耗時與資源消耗針對硬體搜一個模型出來最不錯,目前搜出來的模型轉為onnx測試結果(輸入192x192):
GPU:
Average Performance excluding first iteration. Iterations 2 to 300. (Iterations greater than 1 only bind and evaluate)
  Average Bind: 0.124713 ms
  Average Evaluate: 16.0683 ms
​
  Average Working Set Memory usage (bind): 6.53219e-05 MB
  Average Working Set Memory usage (evaluate): 0.546117 MB
​
  Average Dedicated Memory usage (bind): 0 MB
  Average Dedicated Memory usage (evaluate): 0 MB
​
  Average Shared Memory usage (bind): 0 MB
  Average Shared Memory usage (evaluate): 0.000483382 MB
 
CPU:
Average Performance excluding first iteration. Iterations 2 to 300. (Iterations greater than 1 only bind and evaluate)
  Average Bind: 0.150212 ms
  Average Evaluate: 13.7656 ms
​
  Average Working Set Memory usage (bind): 9.14507e-05 MB
  Average Working Set Memory usage (evaluate): 0.566746 MB
​
  Average Dedicated Memory usage (bind): 0 MB
  Average Dedicated Memory usage (evaluate): 0 MB
​
  Average Shared Memory usage (bind): 0 MB
  Average Shared Memory usage (evaluate): 0 MB
  • 模型分支:在使用的觀察中發現,大部分較為穩定的場景可以使用較小的模型得到不錯的結果,所有考慮finetune LRBranch處理簡單場景,HRBranch與FusionBranch依舊用來處理複雜場景,這項工作還在進行中。

[^1]: Xu, Ning, et al. “Deep image matting.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2017
[^2]:Sengupta, Soumyadip, et al. “Background matting: The world is your green screen.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.
[^3]:Lin, Shanchuan, et al. “Real-Time High-Resolution Background Matting.” arXiv preprint arXiv:2012.07810 (2020).
[^4]:Li, Yaoyi, et al. “Inductive Guided Filter: Real-Time Deep Matting with Weakly Annotated Masks on Mobile Devices.” 2020 IEEE International Conference on Multimedia and Expo (ICME). IEEE, 2020.
[^5]: Yu, Qihang, et al. “Mask Guided Matting via Progressive Refinement Network.” arXiv e-prints (2020): arXiv-2012.
[^6]: Chen, Quan, et al. “Semantic human matting.” Proceedings of the 26th ACM international conference on Multimedia. 2018.
[^7]: Ke, Zhanghan, et al. “Is a Green Screen Really Necessary for Real-Time Human Matting?.” arXiv preprint arXiv:2011.11961 (2020).
[^8]:Rhemann, Christoph, et al. “A perceptually motivated online benchmark for image matting.” 2009 IEEE Conference on Computer Vision and Pattern Recognition. IEEE, 2009.

 

點選關注,第一時間瞭解華為雲新鮮技術~