提供了编程的基础技术教程

zl程序教程

您现在的位置是：首页 >  其他

当前栏目

用上 RNN，这个视频抠像工具效果绝了

项目 工具
2023-03-14 22:34:05 时间

内容导读

本文为字节跳动团队发布的视频抠像工具 RVM 代码解析及论文《Robust High-Resolution Video Matting with Temporal Guidance》概要。

视频抠像，把人物从动态画面中抠除，用于进一步替换背景或其他用途。看似是专业拍摄才用得上的工具，但在有些场景里的确派得上大用处。

快捷、高效的视频抠像有多重要？这一点，今年《披荆斩棘的哥哥》的后期团队一定深有体会。

近期，字节跳动的论文 《Robust High-Resolution Video Matting with Temporal Guidance 》被 WACV 2022 收录，官方 GitHub 库于近期以 GPL-3.0 许可证开源。

RVM 项目开放线上公开测试

此次开源的 RVM 项目专为稳定人物视频抠像设计，可在任意视频上做实时高清抠像。不同于现有神经网络将每一帧作为单独图片处理，RVM 使用循环神经网络，在处理视频流时有时间记忆。

经测试，RVM 支持在 Nvidia GTX 1080Ti 上实现 4K 76FPS 和 HD 104FPS。

此研究项目的团队来自字节跳动。

项目 GitHub https://github.com/PeterL1n/RobustVideoMatting

项目论文 https://arxiv.org/abs/2108.11515

墙内 Colab

https://openbayes.com/console/open-tutorials/containers/oqv42tbd8ko

模型核心：RNN

RVM 的网络架构包括：

- 特征提取编码器：提取单帧特征；

- 循环解码器：综合时间信息；

- 深度引导滤波（DGF）模块：高分辨率采样。

值得注意的是，模型在内部将高分辨率输入缩小做初步的处理，然后再放大做细分处理。所以建议用户设置 downsample_ratio 使缩小后的分辨率维持在 256 到 512 像素之间。

另外，正因为此模型基于循环神经网络（Recurrent Neural Network），所以必须按顺序处理视频每帧，并提供网络循环记忆。

多框架支持

项目成员非常贴心，模型分别基于 MobileNetV3、ResNet50 实现，还支持了PyTorch、TorchHub、TorchScript、ONNX、TensorFlow、TensorFlow.js、CoreML 等主流框架。

PyTorch

- 载入模型

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

- 推断循环（完整实例）

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                                       # 初始记忆

with torch.no_grad():
    for src in DataLoader(reader):
        fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio=0.25)  # 将上一帧的记忆给下一帧
        writer.write(fgr * pha + bgr * (1 - pha))

- 视频转换 API

from inference import convert_video

convert_video(
    model,                           # 模型，可以加载到任何设备（cpu 或 cuda）
    input_source='input.mp4',        # 视频文件，或图片序列文件夹
    input_resize=(1920, 1080),       # [可选项] 缩放视频大小
    downsample_ratio=0.25,           # [可选项] 下采样比，若 None，自动下采样至 512px
    output_type='video',             # 可选 "video"（视频）或 "png_sequence"（PNG 序列）
    output_composition='com.mp4',    # 若导出视频，提供文件路径。若导出 PNG 序列，提供文件夹路径
    output_alpha="pha.mp4",          # [可选项] 输出透明度预测
    output_foreground="fgr.mp4",     # [可选项] 输出前景预测
    output_video_mbps=4,             # 若导出视频，提供视频码率
    seq_chunk=12,                    # 设置多帧并行计算
    num_workers=1,                   # 只适用于图片序列输入，读取线程
    progress=True                    # 显示进度条
)

TensorFlow

- 范例（Channal Last 格式）

import tensorflow as tf

model = tf.keras.models.load_model('rvm_mobilenetv3_tf')
model = tf.function(model)

rec = [ tf.constant(0.) ] * 4         # 初始记忆
downsample_ratio = tf.constant(0.25)  # 下采样率，根据视频调整

for src in YOUR_VIDEO:  # src 张量是 [B, H, W, C] 的形状，而不是 [B, C, H, W]!
    out = model([src, *rec, downsample_ratio])
    fgr, pha, *rec = out['fgr'], out['pha'], out['r1o'], out['r2o'], out['r3o'], out['r4o']

项目 GitHub https://github.com/PeterL1n/RobustVideoMatting

项目论文 https://arxiv.org/abs/2108.11515

Colab

https://openbayes.com/console/open-tutorials/containers/oqv42tbd8ko

相关文章