使用图像分割技术实现视频特效-灵析社区

小乔学算法

一、前言

在许多视频对话软件中，都可以选择视频的背景。其原理就是将人像抠出来，把非人像部分替换。而大多数软件是切换图片背景，或者是动图背景。利用图像分割技术，可以实现更复杂的背景替换，最终结果就像电影特效。

许多特效拍摄会使用绿幕背景，而图像分割技术可以在不使用绿幕的情况下达到类似的效果，不过相较绿幕要差一些。

今天我们要做的就是使用开源的人像分割项目来完成视频背景替换的操作。我们会在原有项目的基础上继续开发。

二、项目介绍

2.1 项目运行

今天我们要使用的是一个视频转绿幕的项目。项目地址：github.com/PeterL1n/Ro…

首先需要下载项目：

git clone https://github.com/PeterL1n/RobustVideoMatting.git

下载完成后，安装对应的依赖：

cd RobustVideoMatting
pip install -r requirements_inference.txt

然后下载对应的模型：github.com/PeterL1n/Ro…

可以选择rvm_mobilenetv3.pth、rvm_resnet50.pth其中一个，放在项目目录下。

然后在项目下创建一个Python文件，写入下面的代码：

import torch
from model import MattingNetwork
from inference import convert_video


# 选择mobilenetv3或resnet50，对应前面下载的两个模型
model = MattingNetwork('mobilenetv3').eval().cuda()  # or "resnet50"
model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))

convert_video(
    model,                           # 模型
    input_source='input.mp4',        # 视频或者图片目录
    output_type='video',             # 选"video"或"png_sequence"
    output_composition='com.mp4',    # 输出路径
    output_alpha="pha.mp4"           # 预测的alpha通道
    output_video_mbps=4,             
    downsample_ratio=None,          
    seq_chunk=12,                 
)

运行上面代码就可以完成视频抠人像的操作，输出的com.mp4是一个绿幕视频，pha.mp4是一个黑白视频（人像为白，背景为黑）。下面是一些例子：

2.2 定义图片抠图和视频抠图函数

在项目中，没有提供自己抠图的代码，所以我们自己编写对应的代码。项目中model.MattingNetwork就是一个人像分割网络，我们调用它的前向传播方法，输入图片，它就会返回抠取的alpha通道，代码如下：

import cv2  
import torch  
from PIL import Image  
from torchvision.transforms import transforms  
  
from model import MattingNetwork  
  
device = "cuda" if torch.cuda.is_available() else "cpu"  
# 加载模型  
segmentor = MattingNetwork('resnet50').eval().cuda()  
segmentor.load_state_dict(torch.load('rvm_resnet50.pth'))  
  
  
def human_segment(model, image):  
    src = (transforms.PILToTensor()(image) / 255.)[None]  
    src = src.to(device)  
    # 抠图  
    with torch.no_grad():  
        fgr, pha, *rec = model(src)  
        segmented = torch.cat([src.cpu(), pha.cpu()], dim=1).squeeze(0).permute(1, 2, 0).numpy()  
        segmented = cv2.normalize(segmented, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)  
        return Image.fromarray(segmented)  

  
human_segment(segmentor, Image.open('xscn.jpg')).show()

现在只需要调用human_segment函数就能实现抠图操作了。而抠视频的操作也非常简单，在项目下提供了下面的代码：

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()  # 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.cuda(), *rec, downsample_ratio)  # Cycle the recurrent states.
        com = fgr * pha + bgr * (1 - pha)              # Composite to green background. 
        writer.write(com)                              # Write frame.

我们也可以自己编写代码用于视频转换，这里的操作就是逐帧读取，然后调用human_segment函数。这里写一个示例代码：

capture = cv2.VideoCapture("input.mp4")  
while True:  
    ret, frame = capture.read()  
    if not ret:  
        break  
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  
    result = human_segment(segmentor, Image.fromarray(image))  
    result = cv2.cvtColor(np.array(result), cv2.COLOR_RGB2BGR)  
    cv2.imshow("result", result)  
    cv2.waitKey(10)  
    cv2.destroyAllWindows()

不过上述代码显示时无法显示透明效果。

三、视频背景切换

视频背景切换的思路可以分为下面几个步骤：

1. 读取人像视频帧和背景视频帧
2. 对每一帧进行抠图
3. 把抠出来的人像与背景混合
4. 把混合结果写入视频

下面我们一步步来实现。其中1、2步我们已经实现了，下面要做的就是3、4两个步骤。

3.1 png图片切换背景

第三步可以看做是给png图片换背景的操作，可以把这一步实现为一个函数，代码如下：

from PIL import Image  
  
  
def change_background(image, background):  
    w, h = image.size  
    background = background.resize((w, h))  
    background.paste(image, (0, 0), image)  
    return background

其中image和background都是Pillow图片。

3.2 写入视频

最后就是写入视频的操作了，这个可以用OpenCV实现，代码如下：

# 读取人像视频和背景视频
capture = cv2.VideoCapture("input.mp4")  
capture_background = cv2.VideoCapture('background.mp4')  
# 获取画面大小  
fps = capture.get(cv2.CAP_PROP_FPS)  
width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))  
height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))  
size = (width, height)  
# 写入视频  
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  
out = cv2.VideoWriter('output.mp4', fourcc, fps, size)  

frames = min(capture.get(cv2.CAP_PROP_FRAME_COUNT), capture_background.get(cv2.CAP_PROP_FRAME_COUNT))  
bar = tqdm(total=frames)
while True:  
    ret1, frame1 = capture.read()  
    ret2, frame2 = capture_background.read()  
    # 如果有一个视频结束了，则结束
    if not ret1 or not ret2:  
        break  
    image = cv2.cvtColor(frame1, cv2.COLOR_BGR2RGB)  
    segmented = human_segment(segmentor, Image.fromarray(image))  
    background = Image.fromarray(cv2.cvtColor(frame2, cv2.COLOR_BGR2RGB))  
    changed = change_background(segmented, background)  
    changed = cv2.cvtColor(np.array(changed), cv2.COLOR_RGB2BGR)  
    out.write(changed)  
    bar.update(1)  
  
out.release()

上面的代码就是完成了本章开始提到的四个步骤，最后会输出一个视频。

在代码中，读取了人像和背景视频，而且不要求两个视频长度一样。程序会自动选取较短的视频作为输出视频的时长。

在代码中有一部分BGR转RBG、cv转Pillow的操作，这部分代码还能做一些改进。

点赞 评论 收藏 目录

阅读量：1742

点赞量：0

收藏量:0