用上 RNN，这个视频抠像工具效果绝了
用上 RNN，这个视频抠像工具效果绝了
内容导读
本文为字节跳动团队发布的视频抠像工具 RVM 代码解析及论文《Robust High-Resolution Video Matting with Temporal Guidance》概要。
视频抠像，把人物从动态画面中抠除，用于进一步替换背景或其他用途。看似是专业拍摄才用得上的工具，但在有些场景里的确派得上大用处。
快捷、高效的视频抠像有多重要？这一点，今年《披荆斩棘的哥哥》的后期团队一定深有体会。
近期，字节跳动的论文 《Robust High-Resolution Video Matting with Temporal Guidance 》被 WACV 2022 收录，官方 GitHub 库于近期以 GPL-3.0 许可证开源。
RVM 项目开放线上公开测试
此次开源的 RVM 项目专为稳定人物视频抠像设计，可在任意视频上做实时高清抠像。不同于现有神经网络将每一帧作为单独图片处理，RVM 使用循环神经网络，在处理视频流时有时间记忆。
经测试，RVM 支持在 Nvidia GTX 1080Ti 上实现 4K 76FPS 和 HD 104FPS。
此研究项目的团队来自字节跳动。
项目 GitHub https://github.com/PeterL1n/RobustVideoMatting
项目论文 https://arxiv.org/abs/2108.11515
墙内 Colab
https://openbayes.com/console/open-tutorials/containers/oqv42tbd8ko
模型核心：RNN
RVM 的网络架构包括：
- 特征提取编码器：提取单帧特征；
- 循环解码器：综合时间信息；
- 深度引导滤波（DGF）模块：高分辨率采样。
值得注意的是，模型在内部将高分辨率输入缩小做初步的处理，然后再放大做细分处理。所以建议用户设置 downsample_ratio 使缩小后的分辨率维持在 256 到 512 像素之间。
另外，正因为此模型基于循环神经网络（Recurrent Neural Network），所以必须按顺序处理视频每帧，并提供网络循环记忆。
多框架支持
项目成员非常贴心，模型分别基于 MobileNetV3、ResNet50 实现，还支持了PyTorch、TorchHub、TorchScript、ONNX、TensorFlow、TensorFlow.js、CoreML 等主流框架。
PyTorch
- 载入模型
import torch
from model import MattingNetwork
model = MattingNetwork(variant='mobilenetv3').eval().cuda() # 或 variant="resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))- 推断循环（完整实例）
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from inference_utils import VideoReader, VideoWriter
reader = VideoReader('input.mp4', transform=ToTensor())
writer = VideoWriter('output.mp4', frame_rate=30)
bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1).cuda() # 绿背景
rec = [None] * 4 # 初始记忆
with torch.no_grad():
for src in DataLoader(reader):
fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio=0.25) # 将上一帧的记忆给下一帧
writer.write(fgr * pha + bgr * (1 - pha))
- 视频转换 API
from inference import convert_video
convert_video(
model, # 模型，可以加载到任何设备（cpu 或 cuda）
input_source='input.mp4', # 视频文件，或图片序列文件夹
input_resize=(1920, 1080), # [可选项] 缩放视频大小
downsample_ratio=0.25, # [可选项] 下采样比，若 None，自动下采样至 512px
output_type='video', # 可选 "video"（视频）或 "png_sequence"（PNG 序列）
output_composition='com.mp4', # 若导出视频，提供文件路径。若导出 PNG 序列，提供文件夹路径
output_alpha="pha.mp4", # [可选项] 输出透明度预测
output_foreground="fgr.mp4", # [可选项] 输出前景预测
output_video_mbps=4, # 若导出视频，提供视频码率
seq_chunk=12, # 设置多帧并行计算
num_workers=1, # 只适用于图片序列输入，读取线程
progress=True # 显示进度条
)
TensorFlow
- 范例（Channal Last 格式）
import tensorflow as tf
model = tf.keras.models.load_model('rvm_mobilenetv3_tf')
model = tf.function(model)
rec = [ tf.constant(0.) ] * 4 # 初始记忆
downsample_ratio = tf.constant(0.25) # 下采样率，根据视频调整
for src in YOUR_VIDEO: # src 张量是 [B, H, W, C] 的形状，而不是 [B, C, H, W]!
out = model([src, *rec, downsample_ratio])
fgr, pha, *rec = out['fgr'], out['pha'], out['r1o'], out['r2o'], out['r3o'], out['r4o']
项目 GitHub https://github.com/PeterL1n/RobustVideoMatting
项目论文 https://arxiv.org/abs/2108.11515
Colab
https://openbayes.com/console/open-tutorials/containers/oqv42tbd8ko