鲁棒视频抠图 (RVM)

English | 中文

本仓库为论文 具有时间引导的鲁棒高分辨率视频抠图 的官方代码库。RVM 专门用于鲁棒的人体视频抠图。与现有将每一帧视为独立图像处理的神经网络模型不同，RVM 使用循环神经网络结合时间记忆来处理视频。RVM 可以在没有任何额外输入的情况下对任意视频进行实时抠图，在 Nvidia GTX 1080 Ti 显卡上分别达到 4K 76FPS 和 HD 104FPS 的速度。该项目由 字节跳动 开发。

最新消息

• [2021年11月3日] 修复了 train.py 中的一个 bug。
• [2021年9月16日] 代码以 GPL-3.0 许可证重新发布。
• [2021年8月25日] 源代码和预训练模型正式公开。
• [2021年7月27日] 论文被 WACV 2022 接收。

展示视频

观看展示视频（YouTube、Bilibili），了解模型的实际效果。

视频中使用的所有素材均可在 Google Drive 中找到。

示例演示

• 摄像头演示：在浏览器中实时运行模型，并可视化循环状态。
• Colab 演示：使用免费 GPU 在您自己的视频上测试我们的模型。

下载

我们推荐在大多数用例中使用 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 和 Baidu Pan（提取码：gym7）中找到。

PyTorch 示例

1. 安装依赖：

pip install -r requirements_inference.txt

2. 加载模型：

import torch
from model import MattingNetwork

model = MattingNetwork('mobilenetv3').eval().cuda()  # 或 "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))

3. 为了转换视频，我们提供了一个简单的转换 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,             # 输出视频的码率。PNG 序列不需要此参数。
    downsample_ratio=None,           # 超参数，用于调整；若设为 None，则自动调整。
    seq_chunk=12,                    # 每次处理 n 帧，以提高并行性。
)

4. 或者编写自己的推理代码：

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                                       # 初始循环状态。
downsample_ratio = 0.25                                # 根据你的视频调整。

with torch.no_grad():
    for src in DataLoader(reader):                     # RGB 张量归一化到 0 ~ 1。
        fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio)  # 更新循环状态。
        com = fgr * pha + bgr * (1 - pha)              # 合成到绿色背景上。
        writer.write(com)                              # 写入帧。

5. 模型和转换器 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。

项目成员

• Shanchuan Lin
• Linjie Yang
• Imran Saleemi
• Soumyadip Sengupta

第三方项目

• NCNN C++ Android (@FeiGeChuanShu)
• lite.ai.toolkit (@DefTruth)
• Gradio Web Demo (@AK391)
• Unity Engine demo with NatML (@natsuite)
• MNN C++ Demo (@DefTruth)
• TNN C++ Demo (@DefTruth)

RobustVideoMatting 快速上手指南

RobustVideoMatting (RVM) 是一个专为鲁棒人像视频抠图设计的开源项目。它利用循环神经网络（RNN）处理视频的时间记忆，无需额外输入即可实时处理任意视频。在 Nvidia GTX 1080 Ti 上可实现 4K 76FPS 和 HD 104FPS 的推理速度。

环境准备

• 操作系统: Linux, macOS, Windows
• Python: 3.7+
• 深度学习框架: PyTorch 1.7+ (推荐 CUDA 版本以获取最佳性能)
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (推荐显存 4GB 以上)。若无 GPU，可使用 CPU 运行但速度较慢。
• 前置依赖:
  • ffmpeg: 用于视频读写（系统层面安装）。
  • Python 依赖库：torch, torchvision, numpy, Pillow, tqdm 等。

安装步骤

1. 安装系统依赖 (FFmpeg)

确保系统中已安装 ffmpeg。

• Ubuntu/Debian: sudo apt-get update && sudo apt-get install ffmpeg
• macOS: brew install ffmpeg
• Windows: 下载并安装 ffmpeg-builds 或将 ffmpeg.exe 添加到环境变量 PATH 中。

2. 克隆代码库

git clone https://github.com/PeterL1n/RobustVideoMatting.git
cd RobustVideoMatting

3. 安装 Python 依赖

推荐使用国内镜像源加速安装（如清华源）：

pip install -r requirements_inference.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

4. 下载预训练模型

对于大多数应用场景，推荐使用 MobileNetv3 模型（速度快、体积小）。如需更高精度可选 ResNet50。

方式 A：直接下载 (推荐国内用户使用百度网盘)

• 百度网盘: 链接 (提取码：gym7)
• 下载 rvm_mobilenetv3.pth 并放置在项目根目录。

方式 B：使用 TorchHub (无需手动下载) 后续代码示例中将展示如何自动加载。

基本使用

以下提供两种最简单的使用方式：一键视频转换 和 自定义推理代码。

方式一：使用官方转换 API (最简单)

此方法适合直接对视频文件进行批量抠图处理。

创建文件 inference_demo.py：

import torch
from model import MattingNetwork
from inference import convert_video

# 1. 加载模型 (自动下载或使用本地文件)
# 使用 TorchHub 加载 (首次运行会自动下载权重)
model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3").eval().cuda()

# 或者手动加载本地权重:
# model = MattingNetwork('mobilenetv3').eval().cuda()
# model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))

# 2. 执行视频转换
convert_video(
    model,                           # 模型实例
    input_source='input.mp4',        # 输入视频路径
    output_type='video',             # 输出类型："video" 或 "png_sequence"
    output_composition='com.mp4',    # 输出合成视频路径 (前景 + 新背景)
    output_alpha="pha.mp4",          # [可选] 输出纯 Alpha 通道视频
    output_foreground="fgr.mp4",     # [可选] 输出纯前景视频
    output_video_mbps=4,             # 输出视频码率
    downsample_ratio=None,           # 下采样比例，None 表示根据分辨率自动调整
    seq_chunk=12,                    # 并行处理的帧数块大小，提升速度
)

运行命令:

python inference_demo.py

方式二：自定义推理循环 (进阶)

如果你需要将抠图结果实时合成到特定背景，或集成到自己的 pipeline 中：

import torch
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from model import MattingNetwork
from inference_utils import VideoReader, VideoWriter

# 1. 初始化
model = MattingNetwork('mobilenetv3').eval().cuda()
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))

reader = VideoReader('input.mp4', transform=ToTensor())
writer = VideoWriter('output.mp4', frame_rate=30)

# 设置绿色背景 (BGR 格式，归一化到 0-1)
bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1).cuda()
rec = [None] * 4               # 初始化循环状态 (Recurrent States)
downsample_ratio = 0.25        # 根据视频分辨率调整，高分辨率建议设小 (如 0.125)

# 2. 逐帧推理
with torch.no_grad():
    for src in DataLoader(reader):
        # src: [B, C, H, W], 归一化到 0~1
        fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio)
        
        # 合成公式：前景 * Alpha + 背景 * (1 - Alpha)
        com = fgr * pha + bgr * (1 - pha)
        
        writer.write(com)

关键参数说明

• downsample_ratio: 重要的超参数。
  • None: 自动根据分辨率选择（推荐）。
  • 0.25: 适用于 1080P (HD) 视频。
  • 0.125: 适用于 4K 视频。
  • 值越小速度越快，但边缘细节可能略有损失。
• seq_chunk: 每次并行处理的帧数。增大此值可利用 GPU 并行能力加速，但会增加显存占用。