RVM最新AI人像摳圖技術，程式碼已經開源

在計算機視覺處理中，人像摳圖有著重要的意義，當然PS技術也能達到類似的效果，其實本質也是程式碼的實現，前2期的影片，我們使用了RVM人像摳圖技術來實現人像的摳圖，效果影片如下：

RVM是一款功能強大的實時影片人像摳圖技術，其由位元組跳動專案組開發。不同於現有神經網路將每一幀作為單獨圖片處理，RVM使用迴圈神經網路，在處理影片流時有時間記憶。RVM可在任意影片上做實時高畫質摳像。在Nvidia GTX 1080Ti上實現4K 76FPS和HD 104FPS。開源地址：

github.com/PeterL1n/RobustVideoMatting

我們可以從GitHub上下載其原始碼，下載完成後，在專案資料夾裡面可以建立一個test.py檔案，當然在開始本程式碼之前需要安裝如下第三方庫

1	av==8.0.3
2	torch==1.9.0
3	torchvision==0.10.0
4	tqdm==4.61.1
5	pims==0.5

安裝完成後，需要下載其專案預訓練好的模型，當然大家可以在GitHub上面找到其分享的模型地址進行下載，其模型包含2種預訓練模型：

rvm_mobilenetv3.pth
rvm_resnet50.pth

rvm_mobilenetv3.pth模型比較小，速度快，大家可以按照此模型進行嘗試，當然，若電腦配置較高，可以使用rvm_mobilenetv3.pth模型

前面的準備工作完成後，在test.py檔案中輸入如下程式碼：

import torch
from model import MattingNetwork
from inference import convert_video
model = MattingNetwork('mobilenetv3').eval().cuda() # or "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth')) # or rvm_resnet50.pth
convert_video(
model, # The model, can be on any device (cpu or cuda).
input_source='input.mp4', # A video file or an image sequence directory.
output_type='video', # Choose "video" or "png_sequence"
output_composition='output.mp4', # File path if video; directory path if png sequence.
output_video_mbps=4, # Output video mbps. Not needed for png sequence.
downsample_ratio=None, # A hyperparameter to adjust or use None for auto.
seq_chunk=12, # Process n frames at once for better parallelism.
)

model = MattingNetwork('mobilenetv3').eval().cuda()，這裡模型可以選擇mobilenetv3，當然也可以選擇resnet50，若電腦沒有GPU，刪除此行程式碼中的cuda()即可

torch.load('rvm_mobilenetv3.pth')載入模型檔案，此模型檔案需要放置在test.py資料夾的根目錄中

然後使用convert_video函式便可以成功生成背景為綠色的人像摳圖影片了

convert_video函式：

引數一：model為前定義的預訓練模型

引數二：input_source，輸入影片地址路徑

引數三：output_composition，輸出影片的路徑地址

引數四：output_video_mbps，輸出影片的mbps引數

引數五：downsample_ratio，一些超引數調整，可以配置成None,軟體自動配置

引數六：seq_chunk，由於此技術具有時間記憶功能，可以同時一次處理多個影片幀來加快影片處理的速度

當然若想輸出Pha通道與fgr通道

新增引數如下：

output_alpha=‘輸出路徑’

output_foreground=‘輸出路徑’

然後執行test.py，待程式碼執行完成後，便在設定的目錄下，自動生成摳圖後的效果影片

當然專案也可以按照如下方式進行摳圖

from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from inference_utils import VideoReader, VideoWriter
import cv2
reader = VideoReader('415809648-1-208.mp4', transform=ToTensor())
writer = VideoWriter('output1234.mp4', frame_rate=30)

bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1)  # Green background.
rec = [None] * 4                                       # Initial recurrent states.
downsample_ratio = 0.25                               # Adjust based on your video.

with torch.no_grad():
    for src in DataLoader(reader):                     # RGB tensor normalized to 0 ~ 1.
        fgr, pha, *rec = model(src, *rec, downsample_ratio)  # Cycle the recurrent states.
        com = fgr * pha + bgr * (1 - pha)              # Composite to green background.
        #cv2.imshow('com',com)
        writer.write(com)                              # Write frame.

當然也可以直接使用API 的方式

# Load the model.
model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # or "resnet50"
# Converter API.
convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")

分類： 健康

時間： 2021-10-27