20210520：人脸识别-口罩mask快速抠取

上篇文章讲述如何快速将口罩mask与人脸face贴合问题。但通常人脸数据较好收集，口罩数据较少，所以比较可行的方案就是：爬取带有口罩的图像，将口罩分割出来。

目前有两种方法：

1：使用PS工具，将口罩抠下来，前提是得会ps...对于不会ps的我，就得另想办法啦，ps无非就是后台使用代码帮你抠而已，自己实现不香嘛

2：使用标注工具【19年博文有介绍】，选择多边形进行标注，会将shape/landmark信息自动保存在xml文件中。然后配合下面的代码，就可以得到文中开头的口罩抠取效果图了。

代码过于简单，就不再进行注解了...

#!/usr/bin/python
# -*- coding:utf-8 -*-

import os
import cv2
import numpy as np
import xml.etree.ElementTree as ET

# define map between mask and bg
class_map = {
    '_background_': 0,
    'mask': 1,
}


# search file by suffix
def file_search(path, allfile):
    filelist = os.listdir(path)

    for filename in filelist:
        filepath = os.path.join(path, filename)
        if os.path.isdir(filepath):
            file_search(filepath, allfile)
        else:
            if filepath.split(".")[-1] in img_ext:
                # if filepath.endswith('_label_landmark.png'):
                allfile.append(filepath)
    return allfile


# parse box info
def parse_box_info(xml):
    tree = ET.parse(xml)
    # root = tree.getroot()
    nodes = tree.findall('object')
    bndboxes = []
    for child in nodes:
        bndbox = []
        tmp = child.find('bndbox')
        xmin = tmp.find('xmin').text
        ymin = tmp.find('ymin').text
        xmax = tmp.find('xmax').text
        ymax = tmp.find('ymax').text
        bndbox.append(xmin)
        bndbox.append(ymin)
        bndbox.append(xmax)
        bndbox.append(ymax)
        bndbox = np.array(bndbox, np.int32)
        bndboxes.append(bndbox)
    return bndboxes


# parse shape info
def parse_shape_info(xml):
    tree = ET.parse(xml)
    # root = tree.getroot()
    nodes = tree.findall('object')
    shapes = []
    for child in nodes:
        shape = []
        tmp = child.find('shape')
        tmp_shape = tmp.find('points')
        count  = 0
        coor_x = 0
        coor_y = 0
        for n in range(0, len(tmp_shape)):
            count += 1
            if n%2 == 0:
                value_x = int(tmp_shape[n].text)
            else:
                value_y = int(tmp_shape[n].text)
            if count == 2:
                count = 0      
                shape.append([value_x, value_y])
        shapes.append([shape])
    return shapes


# get mask area
def get_mask(img, mask_shape):
    b_channel, g_channel, r_channel = cv2.split(img)  
    alpha_channel = np.zeros(b_channel.shape, dtype=b_channel.dtype)
    cv2.fillPoly(alpha_channel, [np.array(mask_shape[0])] ,(255))    
    img_BGRA = cv2.merge((b_channel, g_channel, r_channel, alpha_channel))
    
    img_fill = img.copy()
    cv2.fillPoly(img_fill, [np.array(mask_shape[0])] ,(0,0,255))    
    return img_BGRA, img_fill

def labelcolormap(N):
    cmap = np.zeros((N, 3), dtype=np.uint8)
    for i in range(N):
        r = 0
        g = 0
        b = 0
        id = i
        for j in range(7):
            str_id = uint82bin(id)
            r = r ^ (np.uint8(str_id[-1]) << (7 - j))
            g = g ^ (np.uint8(str_id[-2]) << (7 - j))
            b = b ^ (np.uint8(str_id[-3]) << (7 - j))
            id = id >> 3
        cmap[i, 0] = r
        cmap[i, 1] = g
        cmap[i, 2] = b
    return cmap


if __name__ == '__main__':
    flag_show = True
    img_ext = ['jpg','png','bmp']
    folder_face = '/home/lilasi/sshfs/Code/face_mask_seg/imgs'
    
    # search image
    imgs_list = []
    file_search(folder_face, imgs_list)
    
    # loop
    for idx, img_file in enumerate(imgs_list):
        xml_file = img_file[:-4] + ".xml"
        
        # load image
        image = cv2.imread(img_file)
        im_h = image.shape[0]
        im_w = image.shape[1]
        
        # load mask shape point
        mask_shape = parse_shape_info(xml_file)
        
        # get mask area
        img_mask, img_fill = get_mask(image, mask_shape)
        name_img_mask = folder_face + '/' + os.path.basename(img_file)[:-4] + "_mask.png"
        name_img_fill = folder_face + '/' + os.path.basename(img_file)[:-4] + "_fill.png"
        cv2.imwrite(name_img_mask, img_mask)
        cv2.imwrite(name_img_fill, img_fill)
        if flag_show:
            cv2.imshow("mask", img_mask)
            cv2.imshow("fill", img_fill)
            cv2.waitKey(0)