# 稳定视频抠像 (RVM)
< p align = "center" > < a href = "README.md" > English< / a > | 中文< / p >
论文 [Robust High-Resolution Video Matting with Temporal Guidance ](https://peterl1n.github.io/RobustVideoMatting/ ) 的官方 GitHub 库。RVM 专为稳定人物视频抠像设计。不同于现有神经网络将每一帧作为单独图片处理， ， **4K 76FPS** 和 **HD 104FPS** 。此研究项目来自[字节跳动](https://www.bytedance.com/)。
< br >
## 更新
* [2021年9月16日] 代码重新以 GPL-3.0 许可发布。
* [2021年8月25日] 公开代码和模型。
* [2021年7月27日] 论文被 WACV 2022 收录。
< br >
## 展示视频
观看展示视频 ([YouTube](https://youtu.be/Jvzltozpbpk), [Bilibili ](https://www.bilibili.com/video/BV1Z3411B7g7/ ))，了解模型能力。
< p align = "center" >
< a href = "https://youtu.be/Jvzltozpbpk" >
< img src = "documentation/image/showreel.gif" >
< / a >
< / p >
视频中的所有素材都提供下载，可用于测试模型：[Google Drive](https://drive.google.com/drive/folders/1VFnWwuu-YXDKG-N6vcjK_nL7YZMFapMU?usp=sharing) 或[百度网盘](https://pan.baidu.com/s/1igMteDwN5rO1Sn7YIhBlvQ)（密码: tb3w）
< br >
## Demo
* [网页 ](https://peterl1n.github.io/RobustVideoMatting/#/demo ): 在浏览器里看摄像头抠像效果，展示模型内部循环记忆值。
* [Colab ](https://colab.research.google.com/drive/10z-pNKRnVNsp0Lq9tH1J_XPZ7CBC_uHm?usp=sharing ): 用我们的模型转换你的视频。
< br >
## 下载
推荐在通常情况下使用 MobileNetV3 的模型。ResNet50 的模型大很多，效果稍有提高。我们的模型支持很多框架。详情请阅读[推断文档](documentation/inference_zh_Hans.md)。
< table >
< thead >
< tr >
< td > 框架< / td >
< td > 下载< / td >
< td > 备注< / td >
< / tr >
< / thead >
< tbody >
< tr >
< td > PyTorch< / td >
< td >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3.pth" > rvm_mobilenetv3.pth< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50.pth" > rvm_resnet50.pth< / a >
< / td >
< td >
官方 PyTorch 模型权值。< a href = "documentation/inference_zh_Hans.md#pytorch" > 文档< / a >
< / td >
< / tr >
< tr >
< td > TorchHub< / td >
< td >
无需手动下载。
< / td >
< td >
更方便地在你的 PyTorch 项目里使用此模型。< a href = "documentation/inference_zh_Hans.md#torchhub" > 文档< / a >
< / td >
< / tr >
< tr >
< td > TorchScript< / td >
< td >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.torchscript" > rvm_mobilenetv3_fp32.torchscript< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp16.torchscript" > rvm_mobilenetv3_fp16.torchscript< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50_fp32.torchscript" > rvm_resnet50_fp32.torchscript< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50_fp16.torchscript" > rvm_resnet50_fp16.torchscript< / a >
< / td >
< td >
若需在移动端推断，可以考虑自行导出 int8 量化的模型。< a href = "documentation/inference_zh_Hans.md#torchscript" > 文档< / a >
< / td >
< / tr >
< tr >
< td > ONNX< / td >
< td >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx" > rvm_mobilenetv3_fp32.onnx< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp16.onnx" > rvm_mobilenetv3_fp16.onnx< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50_fp32.onnx" > rvm_resnet50_fp32.onnx< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50_fp16.onnx" > rvm_resnet50_fp16.onnx< / a >
< / td >
< td >
在 ONNX Runtime 的 CPU 和 CUDA backend 上测试过。提供的模型用 opset 12。< a href = "documentation/inference_zh_Hans.md#onnx" > 文档< / a > ， < a href = "https://github.com/PeterL1n/RobustVideoMatting/tree/onnx" > 导出< / a >
< / td >
< / tr >
< tr >
< td > TensorFlow< / td >
< td >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_tf.zip" > rvm_mobilenetv3_tf.zip< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50_tf.zip" > rvm_resnet50_tf.zip< / a >
< / td >
< td >
TensorFlow 2 SavedModel 格式。< a href = "documentation/inference_zh_Hans.md#tensorflow" > 文档< / a >
< / td >
< / tr >
< tr >
< td > TensorFlow.js< / td >
< td >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_tfjs_int8.zip" > rvm_mobilenetv3_tfjs_int8.zip< / a > < br >
< / td >
< td >
在网页上跑模型。< a href = "https://peterl1n.github.io/RobustVideoMatting/#/demo" > 展示< / a > ， < a href = "https://github.com/PeterL1n/RobustVideoMatting/tree/tfjs" > 示范代码< / a >
< / td >
< / tr >
< tr >
< td > CoreML< / td >
< td >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_1280x720_s0.375_fp16.mlmodel" > rvm_mobilenetv3_1280x720_s0.375_fp16.mlmodel< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_1280x720_s0.375_int8.mlmodel" > rvm_mobilenetv3_1280x720_s0.375_int8.mlmodel< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_1920x1080_s0.25_fp16.mlmodel" > rvm_mobilenetv3_1920x1080_s0.25_fp16.mlmodel< / a > < br >
< a href = "https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel" > rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel< / a > < br >
< / td >
< td >
CoreML 只能导出固定分辨率，其他分辨率可自行导出。支持 iOS 13+。< code > s< / code > 代表下采样比。< a href = "documentation/inference_zh_Hans.md#coreml" > 文档< / a > ， < a href = "https://github.com/PeterL1n/RobustVideoMatting/tree/coreml" > 导出< / a >
< / td >
< / tr >
< / tbody >
< / table >
所有模型可在 [Google Drive ](https://drive.google.com/drive/folders/1pBsG-SCTatv-95SnEuxmnvvlRx208VKj?usp=sharing ) 或[百度网盘](https://pan.baidu.com/s/1puPSxQqgBFOVpW4W7AolkA)（密码: gym7）
< br >
## PyTorch 范例
1. 安装 Python 库:
```sh
pip install -r requirements_inference.txt
```
2. 加载模型:
```python
import torch
from model import MattingNetwork
model = MattingNetwork('mobilenetv3').eval().cuda() # 或 "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))
```
3. 若只需要做视频抠像处理，我们提供简单的 API:
```python
from inference import convert_video
convert_video(
model, # 模型， （ ）
input_source='input.mp4', # 视频文件，或图片序列文件夹
output_type='video', # 可选 "video"（视频）或 "png_sequence"（
output_composition='output.mp4', # 若导出视频，提供文件路径。若导出 PNG 序列，提供文件夹路径
output_video_mbps=4, # 若导出视频，提供视频码率
downsample_ratio=None, # 下采样比，可根据具体视频调节，或 None 选择自动
seq_chunk=12, # 设置多帧并行计算
)
```
4. 或自己写推断逻辑:
```python
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): # 输入张量， ， ～
fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio) # 将上一帧的记忆给下一帧
com = fgr * pha + bgr * (1 - pha) # 将前景合成到绿色背景
writer.write(com) # 输出帧
```
5. 模型和 API 也可通过 TorchHub 快速载入。
```python
# 加载模型
model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # 或 "resnet50"
# 转换 API
convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")
```
[推断文档 ](documentation/inference_zh_Hans.md )里有对 `downsample_ratio` 参数，
< br >
## 训练和评估
请参照[训练文档（英文）](documentation/training.md)。
< br >
## 速度
速度用 `inference_speed_test.py` 测量以供参考。
| GPU | dType | 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： `downsample_ratio=0.25` ， `downsample_ratio=0.125` 。 所有测试都使用 batch size 1 和 frame chunk 1。
* 注释2：
* 注释3： （ ） PyNvCodec ](https://github.com/NVIDIA/VideoProcessingFramework )。
< br >
## 第三方资源
* NCNN C++ and Android Demo: [ncnn_Android_RobustVideoMatting ](https://github.com/FeiGeChuanShu/ncnn_Android_RobustVideoMatting ) from [FeiGeChuanShu ](https://github.com/FeiGeChuanShu )
* ONNXRuntime C++ Demo: [lite.ai.toolkit ](https://github.com/DefTruth/lite.ai.toolkit/blob/main/ort/cv/rvm.cpp ) and [demo ](https://github.com/DefTruth/RobustVideoMatting.lite.ai.toolkit ) from [DefTruth ](https://github.com/DefTruth )
## 项目成员
* [Shanchuan Lin ](https://www.linkedin.com/in/shanchuanlin/ )
* [Linjie Yang ](https://sites.google.com/site/linjieyang89/ )
* [Imran Saleemi ](https://www.linkedin.com/in/imran-saleemi/ )
* [Soumyadip Sengupta ](https://homes.cs.washington.edu/~soumya91/ )
