vlambda博客
学习文章列表

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

from torch.utils.data import DataLoaderfrom torchvision.transforms import ToTensorfrom 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))

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                    # 显示进度条)

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']