Robust Video Matting (RVM)
English | 中文
这是论文Robust High-Resolution Video Matting with Temporal Guidance的官方仓库。RVM 专为稳健的人物视频抠像设计。与将帧作为独立图像处理的现有神经网络模型不同,RVM 使用递归神经网络处理具有时间记忆的视频。RVM 可以实时对任何视频进行抠像,无需额外输入。它在 Nvidia GTX 1080 Ti GPU 上可实现 4K 76FPS 和 HD 104FPS。该项目由字节跳动公司开发。
新闻
- [2021年11月03日] 修复了 train.py 中的一个错误。
- [2021年9月16日] 代码重新以 GPL-3.0 许可证发布。
- [2021年8月25日] 发布了源代码和预训练模型。
- [2021年7月27日] 论文被 WACV 2022 接收。
宣传片
观看宣传片视频 (YouTube, Bilibili) 了解模型的表现。
视频中的所有素材均可在 Google Drive 中获取。
演示
下载
我们建议在大多数使用场景中选择 MobileNetv3 模型。ResNet50 模型是更大版本,性能略有提升。我们的模型支持多种推理框架。请参阅推理文档了解更多说明。
框架 | 下载 | 备注 |
PyTorch |
rvm_mobilenetv3.pth rvm_resnet50.pth | PyTorch 官方权重。文档 |
TorchHub | 无需下载。 | 在您的 PyTorch 项目中使用我们的模型的最简单方法。文档 |
TorchScript |
rvm_mobilenetv3_fp32.torchscript rvm_mobilenetv3_fp16.torchscript rvm_resnet50_fp32.torchscript rvm_resnet50_fp16.torchscript | 如果在移动设备上进行推理,请考虑自行导出 int8 量化模型。文档 |
ONNX |
rvm_mobilenetv3_fp32.onnx rvm_mobilenetv3_fp16.onnx rvm_resnet50_fp32.onnx rvm_resnet50_fp16.onnx | 已在 ONNX Runtime 上使用 CPU 和 CUDA 后端进行测试。提供的模型使用 opset 12。文档,导出器。 |
TensorFlow |
rvm_mobilenetv3_tf.zip rvm_resnet50_tf.zip | TensorFlow 2 SavedModel。文档 |
TensorFlow.js |
rvm_mobilenetv3_tfjs_int8.zip | 在网页上运行模型。演示,入门代码 |
CoreML |
rvm_mobilenetv3_1280x720_s0.375_fp16.mlmodel rvm_mobilenetv3_1280x720_s0.375_int8.mlmodel rvm_mobilenetv3_1920x1080_s0.25_fp16.mlmodel rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel |
CoreML 不支持动态分辨率。其他分辨率可自行导出。模型需 iOS 13+。s 表示 downsample_ratio 。文档,导出器
|
所有模型都可以在 Google Drive 和 百度网盘 (提取码: gym7) 下载。
PyTorch 示例
- 安装依赖项:
pip install -r requirements_inference.txt
- 加载模型:
import torch
from model import MattingNetwork
model = MattingNetwork('mobilenetv3').eval().cuda() # 或 "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))
- 要转换视频,我们提供了一个简单的转换 API:
from inference import convert_video
convert_video(
model, # 模型,可以在任何设备上运行 (cpu 或 cuda)。
input_source='input.mp4', # 视频文件或图像序列目录。
output_type='video', # 选择 "video" 或 "png_sequence"
output_composition='com.mp4', # 如果是视频,填写文件路径;如果是 png 序列,填写目录路径。
output_alpha="pha.mp4", # [可选] 输出原始 alpha 预测。
output_foreground="fgr.mp4", # [可选] 输出原始前景预测。
output_video_mbps=4, # 输出视频的 mbps。不需要 png 序列。
downsample_ratio=None, # 调整的超参数,或者使用 None 自动调整。
seq_chunk=12, # 一次处理 n 帧以提高并行性。
)
在 torch.no_grad():
的情况下:
对于 DataLoader(reader)
中的 src
: # RGB 张量归一化到 0 ~ 1。
fgr
, pha
, 以及 rec
= model(src.cuda(), *rec, downsample_ratio)
# 循环递归状态。
com
= fgr
* pha
+ bgr
* (1 - pha
) # 合成到绿色背景。
writer.write(com)
# 写入帧。
模型和转换器 API 也可通过 TorchHub 获取。
# 加载模型。
model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # 或 "resnet50"
# 转换器 API。
convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")
有关 downsample_ratio
超参数、更多转换器参数和更高级用法的详细信息,请参阅 推理文档。
训练与评估
请参考 训练文档 来训练和评估你自己的模型。
速度
速度参考由 inference_speed_test.py
测量。
GPU | 数据类型 | HD (1920x1080) | 4K (3840x2160) |
---|---|---|---|
RTX 3090 | FP16 | 172 FPS | 154 FPS |
RTX 2060 Super | FP16 | 134 FPS | 108 FPS |
GTX 1080 Ti | FP32 | 104 FPS | 74 FPS |
- 注 1: HD 使用
downsample_ratio=0.25
,4K 使用downsample_ratio=0.125
。所有测试使用批量大小 1 和帧块 1。 - 注 2: 图灵架构之前的 GPU 不支持 FP16 推理,因此 GTX 1080 Ti 使用 FP32。
- 注 3: 我们仅测量张量吞吐量。本仓库提供的视频转换脚本预计会慢得多,因为它没有利用硬件视频编码/解码,并且没有在并行线程上完成张量传输。如果你有兴趣在 Python 中实现硬件视频编码/解码,请参考 PyNvCodec。