RobustVideoMatting下载 - RobustVideoMatting源码下载

RobustVideoMatting

Ai源码

1.0.0

下载

强大的视频抠图 (RVM)

预告片

英语 | 中文

具有时间指导的鲁棒高分辨率视频抠图论文的官方存储库。 RVM 专为强大的人体视频抠图而设计。与将帧作为独立图像处理的现有神经模型不同，RVM 使用循环神经网络来处理具有时间记忆的视频。 RVM 可以对任何视频进行实时抠图，无需额外输入。它在 Nvidia GTX 1080 Ti GPU 上实现了4K 76FPS和HD 104FPS 。该项目由字节跳动公司开发。

消息

[2021 年 11 月 3 日] 修复了 train.py 中的错误。
[2021 年 9 月 16 日] 代码根据 GPL-3.0 许可证重新发布。
[2021 年 8 月 25 日] 源代码和预训练模型已发布。
[2021 年 7 月 27 日] 论文被 WACV 2022 接收。

演出影片

观看表演视频（YouTube、Bilibili），了解模特的表演。

视频中的所有片段均可在 Google 云端硬盘中找到。

演示

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

下载

对于大多数用例，我们推荐 MobileNetv3 模型。 ResNet50 模型是较大的变体，性能改进较小。我们的模型可用于各种推理框架。有关更多说明，请参阅推理文档。

框架	下载	笔记
火炬	rvm_mobilenetv3.pth rvm_resnet50.pth	PyTorch 的官方权重。医生
火炬中心	没有什么可下载的。	在 PyTorch 项目中使用我们的模型的最简单方法。医生
TorchScript	rvm_mobilenetv3_fp32.torchscript rvm_mobilenetv3_fp16.torchscript rvm_resnet50_fp32.torch脚本 rvm_resnet50_fp16.torchscript	如果在移动设备上进行推理，请考虑自行导出 int8 量化模型。医生
奥恩克斯	rvm_mobilenetv3_fp32.onnx rvm_mobilenetv3_fp16.onnx rvm_resnet50_fp32.onnx rvm_resnet50_fp16.onnx	在具有 CPU 和 CUDA 后端的 ONNX 运行时上进行了测试。提供的模型使用 opset 12。Doc，Exporter。
TensorFlow	rvm_mobilenetv3_tf.zip rvm_resnet50_tf.zip	TensorFlow 2 SavedModel。医生
TensorFlow.js	rvm_mobilenetv3_tfjs_int8.zip	在网络上运行模型。演示、入门代码
核心机器学习	rvm_mobilenetv3_1280x720_s0.375_fp16.ml模型 rvm_mobilenetv3_1280x720_s0.375_int8.mlmodel rvm_mobilenetv3_1920x1080_s0.25_fp16.ml模型 rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel	CoreML 不支持动态解析。其他分辨率可以自行导出。型号需要 iOS 13+。 `s`表示`downsample_ratio` 。文件，出口商

所有模型均可在 Google Drive 和百度盘中获取（代码：gym7）。

PyTorch 示例

安装依赖项：

 pip install -r requests_inference.txt

加载模型：

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

为了转换视频，我们提供了一个简单的转换 API：

 from inference import Convert_videoconvert_video(model, # 模型，可以在任何设备上（cpu 或 cuda）。input_source='input.mp4', # 视频文件或图像序列 directory.output_type='video', # 选择“video” " or "png_sequence"output_composition='com.mp4', # 如果是 video，则为文件路径；如果为 png，则为目录路径 sequence.output_alpha="pha.mp4", # [可选] 输出原始 alpha Prediction.output_foreground="fgr.mp4", # [可选] 输出原始前景 Prediction.output_video_mbps=4, # 输出视频 mbps。png 不需要。 seq_chunk=12, # 一次处理 n 帧以获得更好的并行性。）

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

 from torch.utils.data 导入 DataLoaderfrom torchvision.transforms 导入 ToTensorfrom inference_utils 导入 VideoReader, VideoWriterreader = VideoReader('input.mp4', transform=ToTensor())writer = VideoWriter('output.mp4', frame_rate=30)bgr = torch .tensor([.47, 1, .6]).view(3, 1, 1).cuda() # Green background.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) # 循环循环 states.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测量速度以供参考。

图形处理器	类型	高清 (1920x1080)	4K (3840x2160)
RTX 3090	FP16	172帧/秒	154帧/秒
RTX 2060 超级	FP16	134帧/秒	108帧/秒
GTX 1080 钛	FP32	104帧/秒	74帧/秒

注1：高清使用downsample_ratio=0.25 ，4K 使用downsample_ratio=0.125 。所有测试均使用批量大小 1 和帧块 1。
注2：图灵架构之前的GPU不支持FP16推理，因此GTX 1080 Ti使用FP32。
注3：我们只测量张量吞吐量。该存储库中提供的视频转换脚本预计会慢得多，因为它不使用硬件视频编码/解码，并且没有在并行线程上完成张量传输。如果您有兴趣用Python实现硬件视频编解码，请参考PyNvCodec。