使用京东FaceX-Zoo sdk生成戴口罩数据集

最近要做戴口罩的人脸识别，事实上当前开源的模型都不那么尽如人意，所以要优化算法再重新训练。

首先要解决的就是数据集的问题。因为从网络中能获取到的真实戴口罩人脸图像比较少，因此需要想办法生成一些戴口罩的人脸数据。这里推荐京东开源的FaceX-Zoo框架(https://gitee.com/yangdashi/FaceX-Zoo/tree/main#/yangdashi/FaceX-Zoo/blob/main/addition_module/face_mask_adding/FMA-3D)，它支持人脸检测、人脸对齐和人脸口罩生成，这里就是使用了以上三种方法生成了戴口罩的人脸数据集。

算法流程：
生成效果：

这里在这位博主开源的https://github.com/zengwb-lx/Face_Mask_Add基础上进行开发，他成功搭建了环境并调用起FaceX-Zoo的sdk，最终实现了单张人脸图片的口罩生成。我在其基础上，增加了批处理、随机抽样，并支持多核处理。

下面贴上核心代码并注意讲解。

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

import logging.config

logging.config.fileConfig("./config/logging.conf")
logger = logging.getLogger('api')

import os
import random
import shutil
from tqdm import trange
from multiprocessing import Pool
import argparse
import cv2
import numpy as np
from skimage.io import imsave
from face_det import FaceDet
from face_align import FaceAlign
from face_masker import FaceMasker


def separate_database():
    List_imgs = os.listdir(args.input_folder)
    Len_imgs = len(List_imgs)  # 数据集长度

    if args.num_cores == 1:  # 单核
        subset = List_imgs[:Len_imgs // 1]

        List_subsets = [subset]

    if args.num_cores == 2:  # 双核，将所有数据集分成两个子数据集
        subset1 = List_imgs[:Len_imgs // 2]
        subset2 = List_imgs[Len_imgs // 2:]

        List_subsets = [subset1, subset2]

    elif args.num_cores == 4:  # 四核，将所有数据集分成四个子数据集
        subset1 = List_imgs[:Len_imgs // 4]
        subset2 = List_imgs[Len_imgs // 4: Len_imgs // 2]
        subset3 = List_imgs[Len_imgs // 2: (Len_imgs * 3) // 4]
        subset4 = List_imgs[(Len_imgs * 3) // 4:]

        List_subsets = [subset1, subset2, subset3, subset4]

    elif args.num_cores >= 8:  # 八核以上，将所有数据集分成八个子数据集
        subset1 = List_imgs[:Len_imgs // 8]
        subset2 = List_imgs[Len_imgs // 8: Len_imgs // 4]
        subset3 = List_imgs[Len_imgs // 4: (Len_imgs * 3) // 8]
        subset4 = List_imgs[(Len_imgs * 3) // 8: Len_imgs // 2]
        subset5 = List_imgs[Len_imgs // 2: (Len_imgs * 5) // 8]
        subset6 = List_imgs[(Len_imgs * 5) // 8: (Len_imgs * 6) // 8]
        subset7 = List_imgs[(Len_imgs * 6) // 8: (Len_imgs * 7) // 8]
        subset8 = List_imgs[(Len_imgs * 7) // 8:]

        List_subsets = [subset1, subset2, subset3, subset4,
                        subset5, subset6, subset7, subset8]

    return List_subsets


def generate_mask(thread_index, input_folder, List_imgs):
    # load model
    face_detector = FaceDet(thread_index)
    face_align = FaceAlign(thread_index)
    face_masker = FaceMasker(True)

    for i in trange(len(List_imgs)):
        subfolder = List_imgs[i]
        images_path = os.path.join(input_folder, subfolder)
        output_subfolder = args.output_folder + subfolder

        if not os.path.exists(output_subfolder):
            os.mkdir(output_subfolder)

        for file in os.listdir(images_path):
            file_path = os.path.join(images_path, file)
            ran = random.randint(0, 9)
            if ran % 10 < args.random_ratio * 10:
                # 添加口罩存储图片
                image = cv2.imread(file_path, cv2.IMREAD_COLOR)

                bboxs = face_detector(image)
                image_show = image.copy()
                face_lms = []

                for box in bboxs:
                    landmarks = face_align(image, box)
                    lms = np.reshape(landmarks.astype(np.int32), (-1))
                    face_lms.append(lms)
                    cv2.rectangle(image_show, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
                    for (x, y) in landmarks.astype(np.int32):
                        cv2.circle(image_show, (x, y), 2, (255, 0, 0), -1)

                # face masker
                save_dir = output_subfolder + "/" + "masked_" + file
                mask_template_name = str(random.randint(0, 7)) + ".png"
                # 未检测到人脸直接跳过
                if len(face_lms) < 1:
                    continue
                # masked one face
                new_image = face_masker.add_mask_one(image, face_lms[0], mask_template_name, mask_template_name)

                imsave(save_dir, new_image)

            else:
                save_dir = output_subfolder + "/" + file
                # 不添加口罩直接复制图片
                shutil.copyfile(file_path, save_dir)


def start_job():
    p = Pool(args.num_cores)
    List_subsets = separate_database()

    # 对每个进程分配工作
    for i in range(args.num_cores):
        p.apply_async(generate_mask, args=(i, args.input_folder, List_subsets[i]))
    p.close()
    p.join()


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('-r', '--random_ratio', type=float, help='ratio of masked face')
    p.add_argument('-i', '--input_folder', type=str, help='input path')
    p.add_argument('-o', '--output_folder', type=str, help='output path')
    p.add_argument('-c', '--num_cores', type=int, default=1, help='number of GPU cores')
    args = p.parse_args()

    if args.random_ratio > 1:
        args.random_ratio = 1
    elif args.random_ratio < 0:
        args.random_ratio = 0

    if not os.path.exists(args.input_folder):
        print("Input path not exist!")
        exit()

    if not os.path.exists(args.output_folder):
        os.mkdir(args.output_folder)

    start_job()

运行指令：

python add_all_demo.py -r 0.1 -i ~/ -o ~/  -c 8  # 说明：-r 抽样率 -i 输入路径 -o 输出路径 -c 调用gpu核心数

• 随机抽样

因为我只需要数据集中存在10%比例的戴口罩人脸图像，因此这里使用了随机数：

ran = random.randint(0, 9)
if ran % 10 < args.random_ratio * 10:

随机到范围内的就戴上口罩，否则就直接转存到指定目录下。这样做是模拟真实使用情况，具体数值自行调整。

• 批处理

批处理前，通常要先把下载下来的数据集转换格式。我所下载的faces_webface_112x112和glint360k的数据格式都是.rec .idx .lst .bin格式的混合文件，需要先通过mxnet转换成图片格式。

随后读取数据，并按照标签存放到对应指定位置。

• 多核处理

这里多核处理的方式是按照gpu核心数分割数据集，然后使用multiprocess的线程池Pool和apply_async 模块进行多进程操作。对应的要修改指定gpu设备的方法，将进程索引分配给对应函数，使其调用起指定的gpu。

综上，我们就完成了戴口罩人脸数据集的制作。