masked Face Recognition Dataset available for download ? or any details of generating ? about facex-zoo HOT 3 OPEN

import math
import sys, os

import cv2
import yaml
import numpy as np
from tqdm import tqdm

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from core.image_cropper.arcface_cropper.FaceRecImageCropper import FaceRecImageCropper
from core.model_handler.face_alignment.FaceAlignModelHandler import FaceAlignModelHandler
from core.model_handler.face_detection.FaceDetModelHandler import FaceDetModelHandler
from core.model_loader.face_alignment.FaceAlignModelLoader import FaceAlignModelLoader
from core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader

from addition_module.face_mask_adding.FMA3D.face_masker import FaceMasker


def distance2center(x1, y1, x2, y2, image):
    im_cx = int(image.shape[1] / 2)
    im_cy = int(image.shape[0] / 2)
    cx = ((x2 + x1) / 2).astype(int)
    cy = ((y2 + y1) / 2).astype(int)
    return math.sqrt(math.pow(im_cx - cx, 2) + math.pow(im_cy - cy, 2))


def Filter2centerBox(boundingBoxes, frame):
    min_distance = 999999999
    min_idx = -1
    for i, det in enumerate(boundingBoxes):
        distance = distance2center(det[0], det[1], det[2], det[3], frame)
        # cv2.rectangle(frame, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (2, 255, 0), 1)
        # cv2.imshow("x", frame)
        # cv2.waitKey(0)
        if distance < min_distance:
            min_idx = i
            min_distance = distance
    return np.array([boundingBoxes[min_idx]])


def GetFaceLandMark(image_path):
    print("aligning {}".format(image_path))
    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
    dets = faceDetModelHandler.inference_on_image(image)
    dets = Filter2centerBox(dets, image)

    assert (len(dets) == 1)  # only one face in picture
    landmarks = faceAlignModelHandler.inference_on_image(image, dets[0])
    return landmarks


input_root = "/home/leo/workspace/data_set/megaface/facescrub_images"
out_root = "/home/leo/workspace/data_set/megaface/facescrub_images_addmasked"

if __name__ == '__main__':
    with open('config/model_conf.yaml') as f:
        model_conf = yaml.load(f)

    model_path = 'models'
    # detect init
    scene = 'non-mask'
    model_category = 'face_detection'
    model_name = model_conf[scene][model_category]
    faceDetModelLoader = FaceDetModelLoader(model_path, model_category, model_name)
    modelDet, cfgDet = faceDetModelLoader.load_model()
    faceDetModelHandler = FaceDetModelHandler(modelDet, 'cuda:0', cfgDet)

    # alignment init
    model_category = 'face_alignment'
    model_name = model_conf[scene][model_category]
    faceAlignModelLoader = FaceAlignModelLoader(model_path, model_category, model_name)
    modelAli, cfgAli = faceAlignModelLoader.load_model()
    faceAlignModelHandler = FaceAlignModelHandler(modelAli, 'cuda:0', cfgAli)

    # face croper
    face_cropper = FaceRecImageCropper()

    # face masker
    face_masker = FaceMasker(is_aug=True)

    template_name = ['0.png', '1.png', '2.png', '3.png', '4.png', '5.png', '6.png', '7.png']
    idx = 0
    all_pathes = []
    for root, dirs, files in os.walk(input_root):
        for file in files:
            image_path = os.path.join(root, file)
            all_pathes.append(image_path)

    for image_path in tqdm(all_pathes):
        masked_face_path = image_path.replace(input_root, out_root)
        if os.path.exists(os.path.dirname(masked_face_path)) is False:
            os.makedirs(os.path.dirname(masked_face_path))
        lms = GetFaceLandMark(image_path)
        face_lms = list(lms.reshape(-1))
        face_masker.add_mask_one(image_path, face_lms, template_name[idx % 8], masked_face_path)
        idx += 1