ModNet抠图算法及摄像头实时抠图示例

数据猴赛雷

已于 2023-06-09 23:58:13 修改

阅读量2.2k

点赞数 1

分类专栏：神经网络文章标签：算法

于 2023-06-09 23:54:28 首次发布

本文链接：https://blog.csdn.net/benben044/article/details/131136506

版权

神经网络专栏收录该内容

61 篇文章 9 订阅

订阅专栏

三、Deep Image Matting算法

四、ModNet算法：Trimap-Free Portrait Matting in Real Time

1、网络结构图

2、算法解读

五、ModNet抠图实践

一、视频抠图采用绿幕的原因

1、摄像机成色原因

主流摄像机传感器为RGB三通道，所以为了抠图最精准最好采用三原色中原始颜色。此外，相机的CMOS传感器矩阵多数都是采用拜耳阵列，该阵列中绿色感光点是2个，高于红色和蓝色，所以信息更丰富更容易抠除。

2、抠图效果原因

视频中的人物和皮肤，多数都是绿色的补色，反差大，这样电脑在渲染处理时就更容易区分边缘和纹理毛发，从而减少抠图的工作量。

3、经济成本

绿背景亮度高，拍摄时光可以亮度调小点从而省电。

二、抠图背景知识

人像抠图：算法概述及工程实现（一）-云社区-华为云

1、Trimap

最常用的先验知识，它是一个三元图，每个像素取值为{0, 128, 255}其中之一，分别代表前景、未知与背景。

2、什么是抠图

对于一张图I，我们感兴趣的人像部分称为前景F，其余部分为背景B，则图像I可以视为F与B的加权融合:
I=alpha∗F+(1−alpha)∗B，alpha的shape与I一致。

而抠图任务就是找到合适的权重alpha矩阵。

将按照上述公式前景图和背景图融合的过程举例如下：

假如一张图的中间圆圈部分为前景，其余部分为背景。则上述两张图按照公式结合后，中间圆圈都是前景相关的像素，而圆圈之外都是背景相关的像素。Alpha对应的是前景图的概率矩阵。

假如alpha训练完成后，若要完成一张图的抠图，只要alpha*原图 + (1-alpha)*白底图即可。

Alpha是介于[0, 1]之间的连续值，可以理解为像素属于前景的概率，这与人像分割是不同的。在人像分割任务中，alpha只能取0或1，本质上是分类任务，而抠图是回归任务。

抠图任务的ground truth，可以看到值分布在0~1之间。

语义分割的ground truth，可以看到值非0即1。

3、抠图算法分类

目前流行的抠图算法大致可以分为两类。

一种是需要先验信息的Trimap-based的方法，宽泛的先验信息包括Trimap、粗糙mask、无人的背景图像、pose信息等，网络使用先验信息与图片信息共同预测alpha

另一种则是Trimap-free的方法，仅根据图片信息预测alpha，对实际应用更友好，但效果普遍不如Trimap-based的方法。

目前主流是trimap-free算法。

三、Deep Image Matting算法

1、网络结构图

2、算法解读

网络包括Encoder-Decoder阶段和Refinement阶段

（1）Encoder-Decoder阶段

输入为RGB图像的patch和对应trimap的concat，所以包含4通道，经过编码和解码后输出单通道的raw alpha pred。该阶段的loss由两部分组成：

第一部分是预测的alpha和真实的alpha之间的绝对误差，考虑到L1 loss在0处不可微，使用Charbonnier Loss去近似：

第二部分是由预测的alpha、真实的前景和真实的背景组成的RGB图像与真实的RGB图像之间的绝对误差，其作用是对网络施加约束，同样使用Charbonnier Loss去近似：

最终的Loss是两部分的加权求和：

（2）Refinement阶段

它的输入为Encoder-Decoder阶段输出的raw alpha pred与原始RGB图像的concat，同样为4通道，原始RGB能够为refine提供边界细节信息。重点是使用了一个skip connection，将Encoder-Decoder阶段输出的raw alpha pred与Refinement阶段输出的refined alpha pred做一个add操作，然后输出最终的预测结果。其实Refinement阶段就是一个residual block，通过残差学习对边界信息进行建模，与去噪模型对噪声建模如出一辙。

Refinement阶段只有一个loss：refined alpha pred与GT alpha matte计算Charbonnier Loss。

四、ModNet算法：Trimap-Free Portrait Matting in Real Time

1、网络结构图

2、算法解读

网络结构由：语义估计分支、细节预测分支、语义-细节融合分支组成。

五、ModNet抠图实践

参考文章：

【Matting】MODNet：实时人像抠图模型-onnx python部署_onnx模型下载_嘟嘟太菜了的博客-CSDN博客

原作者的onnix模型链接：https://download.csdn.net/download/qq_40035462/85046509

代码示例：

import cv2
import time
from tqdm import tqdm
import numpy as np
import onnxruntime as rt


class Matting:
    def __init__(self, model_path='onnx_model\modnet.onnx', input_size=(512, 512)):
        self.model_path = model_path
        self.sess = rt.InferenceSession(self.model_path, providers=['CUDAExecutionProvider'])
        # self.sess = rt.InferenceSession(self.model_path)  # 默认使用cpu
        self.input_name = self.sess.get_inputs()[0].name
        self.label_name = self.sess.get_outputs()[0].name
        self.input_size = input_size
        self.txt_font = cv2.FONT_HERSHEY_PLAIN

    def normalize(self, im, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
        im = im.astype(np.float32, copy=False) / 255.0
        im -= mean
        im /= std
        return im

    def resize(self, im, target_size=608, interp=cv2.INTER_LINEAR):
        if isinstance(target_size, list) or isinstance(target_size, tuple):
            w = target_size[0]
            h = target_size[1]
        else:
            w = target_size
            h = target_size
        im = cv2.resize(im, (w, h), interpolation=interp)
        return im

    def preprocess(self, image, target_size=(512, 512), interp=cv2.INTER_LINEAR):
        image = self.normalize(image)
        image = self.resize(image, target_size=target_size, interp=interp)
        image = np.transpose(image, [2, 0, 1])
        image = image[None, :, :, :]
        return image

    def predict_frame(self, bgr_image):
        assert len(bgr_image.shape) == 3, "Please input RGB image."
        raw_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
        h, w, c = raw_image.shape
        image = self.preprocess(raw_image, target_size=self.input_size)

        pred = self.sess.run(
            [self.label_name],
            {self.input_name: image.astype(np.float32)}
        )[0]
        pred = pred[0, 0]
        matte_np = self.resize(pred, target_size=(w, h), interp=cv2.INTER_NEAREST)
        matte_np = np.expand_dims(matte_np, axis=-1)
        return matte_np

    def predict_image(self, source_image_path, save_image_path):
        bgr_image = cv2.imread(source_image_path)
        assert len(bgr_image.shape) == 3, "Please input RGB image."
        matte_np = self.predict_frame(bgr_image)
        matting_frame = matte_np * bgr_image + (1 - matte_np) * np.full(bgr_image.shape, 255.0)
        matting_frame = matting_frame.astype('uint8')
        cv2.imwrite(save_image_path, matting_frame)

    def predict_camera(self):
        cap_video = cv2.VideoCapture(0)
        if not cap_video.isOpened():
            raise IOError("Error opening video stream or file.")
        beg = time.time()
        count = 0
        while cap_video.isOpened():
            ret, raw_frame = cap_video.read()
            if ret:
                count += 1
                matte_np = self.predict_frame(raw_frame)
                matting_frame = matte_np * raw_frame + (1 - matte_np) * np.full(raw_frame.shape, 255.0)
                matting_frame = matting_frame.astype('uint8')

                end = time.time()
                fps = round(count / (end - beg), 2)
                if count >= 50:
                    count = 0
                    beg = end

                cv2.putText(matting_frame, "fps: " + str(fps), (20, 20), self.txt_font, 2, (0, 0, 255), 1)

                cv2.imshow('Matting', matting_frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            else:
                break
        cap_video.release()
        cv2.destroyWindow()

    def check_video(self, src_path, dst_path):
        cap1 = cv2.VideoCapture(src_path)
        fps1 = int(cap1.get(cv2.CAP_PROP_FPS))
        number_frames1 = cap1.get(cv2.CAP_PROP_FRAME_COUNT)
        cap2 = cv2.VideoCapture(dst_path)
        fps2 = int(cap2.get(cv2.CAP_PROP_FPS))
        number_frames2 = cap2.get(cv2.CAP_PROP_FRAME_COUNT)
        assert fps1 == fps2 and number_frames1 == number_frames2, "fps or number of frames not equal."

    def predict_video(self, video_path, save_path, threshold=2e-7):
        # 使用odf策略
        time_beg = time.time()
        pre_t2 = None  # 前2步matte
        pre_t1 = None  # 前1步matte

        cap = cv2.VideoCapture(video_path)
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
        number_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        print("source video fps: {}, video resolution: {}, video frames: {}".format(fps, size, number_frames))
        videoWriter = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc('I', '4', '2', '0'), fps, size)

        ret, frame = cap.read()
        with tqdm(range(int(number_frames))) as t:
            for c in t:
                matte_np = self.predict_frame(frame)
                if pre_t2 is None:
                    pre_t2 = matte_np
                elif pre_t1 is None:
                    pre_t1 = matte_np
                    # 第一帧写入
                    matting_frame = pre_t2 * frame + (1 - pre_t2) * np.full(frame.shape, 255.0)
                    videoWriter.write(matting_frame.astype('uint8'))
                else:
                    # odf
                    error_interval = np.mean(np.abs(pre_t2 - matte_np))
                    error_neigh = np.mean(np.abs(pre_t1 - pre_t2))
                    if error_interval < threshold < error_neigh:
                        pre_t1 = pre_t2

                    matting_frame = pre_t1 * frame + (1 - pre_t1) * np.full(frame.shape, 255.0)
                    videoWriter.write(matting_frame.astype('uint8'))
                    pre_t2 = pre_t1
                    pre_t1 = matte_np

                ret, frame = cap.read()
            # 最后一帧写入
            matting_frame = pre_t1 * frame + (1 - pre_t1) * np.full(frame.shape, 255.0)
            videoWriter.write(matting_frame.astype('uint8'))
            cap.release()
        print("video matting over, time consume: {}, fps: {}".format(time.time() - time_beg, number_frames / (time.time() - time_beg)))


if __name__ == '__main__':
    model = Matting(model_path='onnx_model\modnet.onnx', input_size=(512, 512))
    model.predict_camera()
    # model.predict_image('images\\1.jpeg', 'output\\1.png')
    # model.predict_image('images\\2.jpeg', 'output\\2.png')
    # model.predict_image('images\\3.jpeg', 'output\\3.png')
    # model.predict_image('images\\4.jpeg', 'output\\4.png')
    # model.predict_video("video\dance.avi", "output\dance_matting.avi")

代码中涉及的modnet.onnx文件见最上面的附件。