强大的视频抠图 (RVM)

RVM 专为强大的人类视频抠图而设计。与将帧作为独立图像处理的现有神经模型不同，RVM 使用循环神经网络来处理具有时间记忆的视频。RVM 可以在任何视频上实时执行抠图，无需额外输入。它在 Nvidia GTX 1080 Ti GPU 上实现了4K 76FPS和HD 104FPS。

论文：Robust High-Resolution Video Matting with Temporal Guidance

github地址：https://github.com/PeterL1n/RobustVideoMatting

这里只做模型测试，训练可以参考github代码

使用方法：点击上方菜单，选择 Run（运行） - Run All Cells（运行所有）

第一步 准备环境

import os
mox.file.copy('obs://obs-aigallery-zc/clf/code/RVM/RVM.zip','RVM.zip')
os.system('unzip RVM.zip')
!pip install --quiet av pims
!pip uninstall -y torch
!pip uninstall -y torchvision
!pip install torch
!pip install torchvision
%cd RVM

第二步 读取模型

import torch
from IPython.display import HTML
from base64 import b64decode, b64encode

model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3",pretrained=False).cuda() # or "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))
convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")

第三步 推理

convert_video(
    model,                           # 读取模型
    input_source='asianboss2.mp4',        # 输入视频
    downsample_ratio=None,           # [可选] 如 None, 最大分辨率降为512px.
    output_type='video',             # 选择 "video" 或 "png_sequence"
    output_composition='com.mp4',    # 输出视频
    output_alpha="pha.mp4",          # [可选] 透明度预测结果
    output_foreground="fgr.mp4",     # [可选] 前景预测结果
    output_video_mbps=4,             # 输出视频 mbps. png序列不需要改.
    seq_chunk=12,                    # 并行
    num_workers=1,                   # 读取线程 （图片读取时需要）
    progress=True                    # 打印转换过程
)

第四步 展示结果

outpath = "com.mp4"
mp4 = open(outpath,'rb').read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML("""
<video width=400 controls>
      <source src="%s" type="video/mp4">
</video>
""" % data_url)