基于阿里Semantatic Human Matting算法，实现精细化人物抠图

Posted 2022-06-10 AI浩

人像抠图

基于深度学习技术研发的人像抠图技术。可识别视频图像中的人像区域，包括头部、半身、全身位置，抠出人像部分后，配以不同背景图片、效果，实现娱乐化需求，支持用户玩转更多个性化操作，常用于直播、视频场景中。

2018年阿里的论文《Semantatic Human Matting》给出了抠图领域的一个新方法，论文、代码和使用的数据集如下：

代码：https://github.com/lizhengwei1992/Semantic_Human_Matting

论文链接：https://arxiv.org/pdf/1809.01354v2.pdf

数据集：https://pan.baidu.com/share/init?surl=R9PJJRT-KjSxh-2-3wCGxQ 密码是：dzsn

数据集介绍

北京玩星汇聚科技有限公司所高质量标注并开源的一份人像抠图(matting) 数据集，是目前已知最大的人像matting数据集，包含 34427 张图像和对应的matting结果图. 数据集中的原始图片来源于Flickr、百度、淘宝. 基于该数据集所训练的人像软分割模型已商用。

数据集下载

数据集分享在百度网盘(国内)：

链接：https://pan.baidu.com/s/1R9PJJRT-KjSxh-2-3wCGxQ

提取码：dzsn

MEGA(国外)：

https://mega.nz/#F!Gh8CFAyb!e2ppUh-copP76GbE8IWAEQ

数据集文件为：matting_human_half.zip，解压后其目录结构为：

- matting_human_half/
	- clip_img #人像图片(半身)，.jpg 格式
	- matting  #人像标注，.png 格式

clip_img文件夹：原图；
matting文件夹：原图对应的mask图，四个维度，第四个维度是图片对应的mask；

该数据集中，图片是经过人脸检测和区域裁剪后生成了600x800的半身人像. 标注的人像 matting 图片为 png 格式，可以从 png 图片中提取人像的 mask图，代码如下：

#!/usr/bin/python3
#!--*-- coding:utf-8 --*--
import cv2
import matplotlib.pyplot as plt

png = cv2.imread('./data/matting/1803151818/matting_00000000/1803151818-00000006.png', cv2.IMREAD_UNCHANGED)
print(png.shape)
#(800, 600, 4)

png_img = cv2.cvtColor(png[:,:,:3], cv2.COLOR_BGR2RGB)
alpha = png[:,:,3]

plt.figure(figsize=(10, 8))
plt.subplot(1, 2, 1)
plt.imshow(png_img)
plt.title("Matting PNG img")
plt.axis("off")
plt.subplot(1, 2, 2)
plt.imshow(alpha)
plt.title("Matting Alpha img")
plt.axis("off")
plt.show()

运行结果：

计算alpha图

下载代码，并解压，将clip_img文件夹和matting文件夹。如下图：

生成train.txt

在data文件夹下，新建get_train_txt.py生成train.txt

import os

pic_path = "matting/"

with open("train.txt", "w", encoding="UTF-8") as ff:
    for name_0 in os.listdir(pic_path):
        for name_1 in os.listdir(pic_path + "/" + name_0):
            for name_2 in os.listdir(pic_path + "/" + name_0 + "/" + name_1):
                pic_input_path = name_0 + "/" + name_1 + "/" + name_2
                ff.write(pic_input_path + "\\n")
    ff.close()
print("well done____________!")

遍历matting，写入train.txt，运行结果：

1803151818/matting_00000000/1803151818-00000003.png
1803151818/matting_00000000/1803151818-00000004.png
1803151818/matting_00000000/1803151818-00000006.png
1803151818/matting_00000000/1803151818-00000007.png
1803151818/matting_00000000/1803151818-00000008.png
1803151818/matting_00000000/1803151818-00000009.png
1803151818/matting_00000000/1803151818-00000010.png
1803151818/matting_00000000/1803151818-00000012.png
1803151818/matting_00000000/1803151818-00000015.png
1803151818/matting_00000000/1803151818-00000017.png
1803151818/matting_00000000/1803151818-00000020.png
1803151818/matting_00000000/1803151818-00000021.png

生成mask

在data新建get_mask.py生成mask。

import os
import cv2


matting_path = "matting/"
mask_path = "mask/"


for name_0 in os.listdir(matting_path):
    if not os.path.exists(mask_path + "/" + name_0):
        os.makedirs(mask_path + "/" + name_0)
    for name_1 in os.listdir(matting_path + "/" + name_0):
        if not os.path.exists(mask_path + name_0 + "/" + name_1):
            os.mkdir(mask_path + name_0 + "/" + name_1)
        for name_2 in os.listdir(matting_path + "/" + name_0 + "/" + name_1):

            pic_input_path = matting_path + "/" + name_0 + "/" + name_1 + "/" + name_2
            pic_output_path = mask_path + "/" + name_0 + "/" + name_1 + "/" + name_2
            print("pic_input_path=", pic_input_path)

            in_image = cv2.imread(pic_input_path, cv2.IMREAD_UNCHANGED)
            alpha = in_image[:, :, 3]
            cv2.imwrite(pic_output_path, alpha)

运行结果：

生成Trimap图

源代码有这个脚本文件gen_trimap.py，需要做一些修改。
首先对参数做修改：

def get_args():
    parser = argparse.ArgumentParser(description='Trimap')
    parser.add_argument('--mskDir', type=str, default='mask', help="masks directory")
    parser.add_argument('--saveDir', type=str, default='trimap', help="where trimap result save to")
    parser.add_argument('--list', type=str, default='train.txt', help="list of images id")
    parser.add_argument('--size', type=int, default=10, help="kernel size")
    args = parser.parse_args()
    print(args)
    return args

注释掉gen_trimap.py第37/43/49行的断言语句

# assert(cnt1 == cnt2 + cnt3)

导入os包

import os

在67行，插入：

        trimap_name_1 = trimap_name.split("/")[:-1]
        trimap_path = "/".join(trimap_name_1)
        if not os.path.exists(trimap_path):
            os.makedirs(trimap_path)

如果trimap路径不存在则创建路径。
修改完成后，运行gen_trimap.py，生成trimap图。运行结果：

单张图：

生成alpha图

修改./data/knn_matting.py脚本，主要修改该main方法。

def main():
    with open('train.txt') as fs:
        list_path = fs.read().splitlines()
        for image_path in list_path:
            img_name = os.path.join(rgbDir, image_path)
            trimap_name = os.path.join(trimapDir, image_path[:-4] + '.png')
            alpha_name = os.path.join(alphaDir, image_path[:-4] + '.png')
            img_name = img_name.replace('\\\\', '/')
            trimap_name = trimap_name.replace('\\\\', '/')
            alpha_name = alpha_name.replace('\\\\', '/')
            index_path = alpha_name.rindex('/')
            alpha_dir = alpha_name[:index_path]
            if not os.path.exists(alpha_dir):
                os.makedirs(alpha_dir)
            img_name = './' + img_name
            img = cv2.imread(img_name)
            trimap = cv2.imread(trimap_name)
            alpha = knn_matte(img, trimap)
            cv2.imwrite(alpha_name, alpha * 255)

运行结果：

可以看出，我们己经得到比较精细的Alpha图了。不过计算过程比较慢，几分钟一张图。如果时间太长可以直接使用mask图。

训练

修改dataset.py文件

将dataset.py文件第17/18/19行修改为：

image_name = os.path.join(data_dir, 'clip_img', file_name['image'].replace("matting", "clip").replace("png", "jpg"))
trimap_name = os.path.join(data_dir, 'trimap', file_name['trimap'].replace("clip", "matting"))
alpha_name = os.path.join(data_dir, 'alpha', file_name['alpha'].replace("clip", "matting"))

./data/clip_img/1803201916/clip_00000000/1803201916-00000117.png，这张图片的格式有问题，需要修改为 .jpg格式

将dataset.py文件第101/102/103行，修改为

trimap[trimap == 0] = 0
trimap[trimap >= 250] = 2
trimap[np.where(~((trimap == 0) | (trimap == 2)))] = 1

原因：trimap图是三色图，但是它的“三色”并不像上图中0/128/255只有这三色，它是在[0, 255]这个区间范围内。所以新改的代码，将这“三色”用区间区分，作为三种不同的label传入训练。

训练T-Net

python train.py --dataDir='./data' --saveDir='./ckpt' --trainData='human_matting_data' --trainList='./data/train.txt' --lrdecayType='keep' --nEpochs=200 --save_epoch=1 --load='human_matting' --patch_size=320 --lr=1e-5 --nThreads=4 --train_batch=4 --train_phase='pre_train_t_net'   

训练M-Net：

python train.py --dataDir='./data' --saveDir='./ckpt' --trainData='human_matting_data' --trainList='./data/train.txt' --lrdecayType='keep' --nEpochs=400 --save_epoch=1 --load='human_matting' --patch_size=320 --lr=5e-6 --gpus='0' --nThreads=4 --train_batch=4 --train_phase='end_to_end'

测试

打开test_camera.py，修改如下参数：

parser = argparse.ArgumentParser(description='human matting')
parser.add_argument('--model', default='./ckpt/human_matting/model/model_obj.pth', help='preTrained model')
parser.add_argument('--size', type=int, default=320,help='input size')
parser.add_argument('--without_gpu', action='store_true', default=False, help='no use gpu')

然后，运行test_camera.py就可以得到运行结果，由于目前没有训练好，只是调通了算法
等调通算法后再展示详细的结果。

代码

https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85564296?spm=1001.2014.3001.5503
参考文章：

【SHM】Semantic Human Matting抠图算法调试_zzZ_CMing的博客-CSDN博客_human matting