使用BackgroundMattingV2模型进行抠图

作为五一前最后一篇文章,说点有用的东西。关于抠图,实际上这玩意都是老生常谈的,没什么创新的地方,目的自然是将其中的人物或事物从背景中分离出来。
关于BackgroundMattingV2抠图,实际上要求蛮高的,要求2张图片的尺寸要一致,而且需要均包含相同的事物。相关的项目可以参考,相关的模型可以点击如下链接进行下载。
对应的原图如下:

另一张图为:

下面提供相关的代码:

import torch
import cv2
import numpy as np
from model import MattingRefine

device = 'cuda' if torch.cuda.is_available() else 'cpu'
precision = torch.float32

model = MattingRefine(backbone='mobilenetv2',
                      backbone_scale=0.25,
                      refine_mode='sampling',
                      refine_sample_pixels=80_000)

model.load_state_dict(torch.load('pytorch_mobilenetv2.pth'))
model = model.eval().to(precision).to(device)

ori_src = cv2.imread("ziling.jpg") / 255.0
src = ori_src.transpose(2,1,0)
src = torch.Tensor(src).unsqueeze(0)
bgr = cv2.imread("out_img.png") / 255.0
bgr = bgr.transpose(2,1,0)
bgr = torch.Tensor(bgr).unsqueeze(0)

with torch.no_grad():
    pha, fgr = model(src, bgr)[:2]
    ...

整体来说代码没什么难度,只是送入模型的图片需要进行归一化处理。而输出需要对其乘以255并转换为整型。
最后是抠出后的效果:

效果可以说比较一般,从下面的蒙版图可以即可看出:

可以看到白色中出现一些裂缝,因此还是需要对其进行后处理修补。否则这么一个漂亮的妹子手臂的皮肤竟然腐烂了,影响心情。