RVM最新AI人像抠图技术，代码已经开源

在计算机视觉处理中，人像抠图有着重要的意义，当然PS技术也能达到类似的效果，其实本质也是代码的实现，前2期的视频，我们使用了RVM人像抠图技术来实现人像的抠图，效果视频如下：

人工智能人像抠图来进行清华虚拟AI人华智冰抠图

实时动态视频背景抠图技术，让没有绿布的你也可以制作绿幕大片

RVM是一款功能强大的实时视频人像抠图技术，其由字节跳动项目组开发。不同于现有神经网络将每一帧作为单独图片处理，RVM使用循环神经网络，在处理视频流时有时间记忆。RVM可在任意视频上做实时高清抠像。在Nvidia GTX 1080Ti上实现4K 76FPS和HD 104FPS。开源地址：

github.com/PeterL1n/RobustVideoMatting

我们可以从GitHub上下载其源代码，下载完成后，在项目文件夹里面可以建立一个test.py文件，当然在开始本代码之前需要安装如下第三方库

安装完成后，需要下载其项目预训练好的模型，当然大家可以在GitHub上面找到其分享的模型地址进行下载，其模型包含2种预训练模型：

1. rvm_mobilenetv3.pth
2. rvm_resnet50.pth

rvm_mobilenetv3.pth模型比较小，速度快，大家可以按照此模型进行尝试，当然，若电脑配置较高，可以使用rvm_mobilenetv3.pth模型

前面的准备工作完成后，在test.py文件中输入如下代码：

import torch
from model import MattingNetwork
from inference import convert_video
model = MattingNetwork('mobilenetv3').eval().cuda() # or "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth')) # or rvm_resnet50.pth
convert_video(
model, # The model, can be on any device (cpu or cuda).
input_source='input.mp4', # A video file or an image sequence directory.
output_type='video', # Choose "video" or "png_sequence"
output_composition='output.mp4', # File path if video; directory path if png sequence.
output_video_mbps=4, # Output video mbps. Not needed for png sequence.
downsample_ratio=None, # A hyperparameter to adjust or use None for auto.
seq_chunk=12, # Process n frames at once for better parallelism.
)

model = MattingNetwork('mobilenetv3').eval().cuda()，这里模型可以选择mobilenetv3，当然也可以选择resnet50，若电脑没有GPU，删除此行代码中的cuda()即可

torch.load('rvm_mobilenetv3.pth')加载模型文件，此模型文件需要放置在test.py文件夹的根目录中

然后使用convert_video函数便可以成功生成背景为绿色的人像抠图视频了

convert_video函数：

参数一：model为前定义的预训练模型

参数二：input_source，输入视频地址路径

参数三：output_composition，输出视频的路径地址

参数四：output_video_mbps，输出视频的mbps参数

参数五：downsample_ratio，一些超参数调整，可以配置成None,软件自动配置

参数六：seq_chunk，由于此技术具有时间记忆功能，可以同时一次处理多个视频帧来加快视频处理的速度

当然若想输出Pha通道与fgr通道

添加参数如下：

output_alpha=‘输出路径’

output_foreground=‘输出路径’

然后运行test.py，待代码运行完成后，便在设置的目录下，自动生成抠图后的效果视频

当然项目也可以按照如下方式进行抠图

from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from inference_utils import VideoReader, VideoWriter
import cv2
reader = VideoReader('415809648-1-208.mp4', transform=ToTensor())
writer = VideoWriter('output1234.mp4', frame_rate=30)

bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1)  # Green background.
rec = [None] * 4                                       # Initial recurrent states.
downsample_ratio = 0.25                               # Adjust based on your video.

with torch.no_grad():
    for src in DataLoader(reader):                     # RGB tensor normalized to 0 ~ 1.
        fgr, pha, *rec = model(src, *rec, downsample_ratio)  # Cycle the recurrent states.
        com = fgr * pha + bgr * (1 - pha)              # Composite to green background. 
        #cv2.imshow('com',com)
        writer.write(com)                              # Write frame.

当然也可以直接使用API 的方式

# Load the model.
model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # or "resnet50"
# Converter API.
convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")