本文分享自華爲雲社區《人像摳圖:算法概述及工程實現(一)》,原文做者:杜甫蓋房子 。python
本文將從算法概述、工程實現、優化改進三個方面闡述如何實現一個實時、優雅、精確的視頻人像摳圖項目。linux
對於一張圖I, 咱們感興趣的人像部分稱爲前景F,其他部分爲背景B,則圖像I能夠視爲F與B的加權融合:I = alpha F + (1 - alpha) BI=alpha∗F+(1−alpha)∗B,而摳圖任務就是找到合適的權重alpha。值得一提的是,如圖,查看摳圖ground truth能夠看到,alpha是[0, 1]之間的連續值,能夠理解爲像素屬於前景的機率,這與人像分割是不一樣的。如圖,在人像分割任務中,alpha只能取0或1,本質上是分類任務,而摳圖是迴歸任務。算法
摳圖ground truth:
segmentfault
分割ground truth:
網絡
咱們主要關注比較有表明性的基於深度學習的摳圖算法。目前流行的摳圖算法大體能夠分爲兩類,一種是須要先驗信息的Trimap-based的方法,寬泛的先驗信息包括Trimap、粗糙mask、無人的背景圖像、Pose信息等,網絡使用先驗信息與圖片信息共同預測alpha;另外一種則是Trimap-free的方法,僅根據圖片信息預測alpha,對實際應用更友好,但效果廣泛不如Trimap-based的方法。dom
Trimap是最經常使用的先驗知識,顧名思義Trimap是一個三元圖,每一個像素取值爲{0,128,255}其中之一,分別表明前景、未知與背景,如圖。
ide
多數摳圖算法採用了Trimap做爲先驗知識。Adobe在17年提出了Deep Image Matting1,這是首個端到端預測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 Matting2是華盛頓大學提出的摳圖算法,後續發佈了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獲得了不錯的效果,但該項目沒法實時運行,也沒法很好的處理高分辨率輸入。因此項目團隊又推出了Background Matting V23,該項目能夠以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 Filter4、MGMatting5等,使用粗略的mask做爲先驗信息預測alpha matte,在應用時也比trimap友好不少。MGMatting同時也提出了一個有636張精確標註人像的摳圖數據集RealWorldPortrait-636,能夠經過合成等數據增廣方法擴展使用。
實際應用中先驗信息獲取起來是很不方便的,一些文章將先驗信息獲取的部分也放在網絡中進行。
阿里巴巴提出的Semantic Human Matting6一樣分解了摳圖任務,網絡分爲三個部分,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=Lα+Lc+Lt。文章實現了端到端Trimap-free的摳圖算法,但較爲臃腫。此外文章提出Fashion Model數據集,從電商網站收集整理了35000+標註的圖片,但並無開放。
modnet7認爲神經網絡更擅長學習單一任務,因此將摳圖任務分爲三個子任務,分別進行顯式監督訓練和同步優化,最終能夠以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+λαLα。
最後,文章還提出了一種使視頻結果在時間上更平滑的後處理方式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)
從新訓練後,咱們的模型效果已經能夠達到前文展現的效果,在16T算力的HiLens Kit上徹底達到了實時、優雅的效果。進一步的,我還想要模型成爲耗時更少、效果更好的優秀模型~目前在作的提高方向是:
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