Excel的主要功能概述
986
2022-05-30
人像抠图:算法概述及工程实现(一)
Pytorch->Caffe模型转换
上一篇概述了人像抠图现有的算法,平衡速度与精度选择了MODNet作为baseline,本文将重点阐述工程实现、优化改进的细节。
本项目的最终目的是在HiLens Kit硬件上落地实现实时视频读入与背景替换,开发环境为HiLens配套在线开发环境HiLens Studio,先上一下对比baseline的改进效果:
使用modnet预训练模型modnet_photographic_portrait_matting.ckpt进行测试结果如下:
可以看到由于场景较为陌生、逆光等原因会导致抠图结果有些闪烁,虽然modnet可以针对特定视频进行自监督finetune,但我们的目的是在普遍意义上效果更好,因此没有对本视频进行自监督学习。
优化后的模型效果如下:
注:原视频来自Human-centric video matting发布的数据集
本视频并没有作为训练数据。可以看到,抠图的闪烁情况减少了很多,毛发等细节也基本没有损失。
工程落地
为了测试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依旧用来处理复杂场景,这项工作还在进行中。
后续还会进行一些量化蒸馏的优化尝试,期待更好更快的结果。
华为HiLens 视频
版权声明:本文内容由网络用户投稿,版权归原作者所有,本站不拥有其著作权,亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容,请联系我们jiasou666@gmail.com 处理,核实后本网站将在24小时内删除侵权内容。