This project aims to add masks to the facial dataset, which is based on FMA-3D and constructs a effective, easy to operate, and efficient pipeline for facial detection, alignment, and mask wearing.
import os
import random
import shutil

import cv2
import argparse
import numpy as np
from tqdm import trange

from api_usage.face_detect import FaceDet
from api_usage.face_alignment import FaceAlign
from api_usage.face_masker import FaceMasker

def generate_mask(input_folder, output_folder, List_imgs, ratio, speed):
    # load model
    face_detector = FaceDet()
    face_align = FaceAlign()
    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 = os.path.join(output_folder, subfolder)

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

        images_path_list = os.listdir(images_path)  # os.listdir Having randomness
        pct = max(1, round(len(images_path_list) * ratio))  # Sampling by ratio, at least one, rounded off

        for file in images_path_list[:pct]:
            file_path = os.path.join(images_path, file)
            image = cv2.imread(file_path, cv2.IMREAD_COLOR)
            bboxs = face_detector(image)
            face_lms = [np.reshape(face_align(image, box).astype(np.int32), (-1)) for box in bboxs]

            # face masker
            file_name, file_lx = file.split(".")
            save_dir = os.path.join(output_subfolder, file_name + "_masked." + file_lx)
            mask_template_name = str(random.randint(0, 7)) + ".png"
            
            # no face continue
            if len(face_lms) < 1:
                continue
            
            face_masker.add_mask_one(file_path, face_lms[0], mask_template_name, save_dir, speed) # masked one face

        for file in images_path_list[pct:]:
            file_path = os.path.join(images_path, file)
            save_dir = os.path.join(output_subfolder, file)
            shutil.copyfile(file_path, save_dir)  # Copy the image without adding a mask


def main():
    List_imgs = os.listdir(args.input_folder)
    generate_mask(args.input_folder, args.output_folder, List_imgs, args.random_ratio, args.speed)


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('input_folder', type=str, help='input path')
    p.add_argument('-r', '--random_ratio', type=float, default=1, help="ratio of masked face")
    p.add_argument('-o', '--output_folder', type=str, default="./output", help="output path")
    p.add_argument('-s', '--speed', type=int, default=1, help="1: Using Cython to speed up, 0: No speed up")
    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)

    main()


add_mask.py
MetaInsight7/MaskFaceTool
import os
import random
import shutil
from multiprocessing import Pool

import cv2
import argparse
import numpy as np
from tqdm import trange

from api_usage.face_detect import FaceDet
from api_usage.face_alignment import FaceAlign
from api_usage.face_masker import FaceMasker

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

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

        List_subsets = [subset]

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

        List_subsets = [subset1, subset2]

    elif 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 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(input_folder, output_folder, List_imgs, ratio, speed):
    # load model
    face_detector = FaceDet()
    face_align = FaceAlign()
    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 = os.path.join(output_folder, subfolder)

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

        images_path_list = os.listdir(images_path)  # os.listdir Having randomness
        pct = max(1, round(len(images_path_list) * ratio))  # Sampling by ratio, at least one, rounded off

        for file in images_path_list[:pct]:
            file_path = os.path.join(images_path, file)
            image = cv2.imread(file_path, cv2.IMREAD_COLOR)
            bboxs = face_detector(image)
            face_lms = [np.reshape(face_align(image, box).astype(np.int32), (-1)) for box in bboxs]

            # face masker
            file_name, file_lx = file.split(".")
            save_dir = os.path.join(output_subfolder, file_name + "_masked." + file_lx)
            mask_template_name = str(random.randint(0, 7)) + ".png"
            
            # no face continue
            if len(face_lms) < 1:
                continue
            
            face_masker.add_mask_one(file_path, face_lms[0], mask_template_name, save_dir, speed) # masked one face

        for file in images_path_list[pct:]:
            file_path = os.path.join(images_path, file)
            save_dir = os.path.join(output_subfolder, file)
            shutil.copyfile(file_path, save_dir)  # Copy the image without adding a mask




def main():
    p = Pool(args.num_cores)
    List_subsets = separate_database(args.input_folder, args.num_cores)

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


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('input_folder', type=str, help='input path')
    p.add_argument('-r', '--random_ratio', type=float, default=1, help="ratio of masked face")
    p.add_argument('-o', '--output_folder', type=str, default="./output", help="output path")
    p.add_argument('-s', '--speed', type=int, default=1, help="1: Using Cython to speed up, 0: No speed up")
    p.add_argument('-c', '--num_cores', type=int, default=4, help='number of CPU 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)

    main()


add_mask_multiproc.py
MetaInsight7/MaskFaceTool
import sys
sys.path.append('.')

import torch
import yaml
import numpy as np
from core.model_loader.face_alignment.FaceAlignModelLoader import FaceAlignModelLoader
from core.model_handler.face_alignment.FaceAlignModelHandler import FaceAlignModelHandler

with open('config/model_conf.yaml') as f:
    model_conf = yaml.load(f, Loader=yaml.FullLoader)

class FaceAlign():
    def __init__(self):
         # common setting for all model, need not modify.
        model_path = 'models'

        # model setting, modified along with model
        scene = 'non-mask'
        model_category = 'face_alignment'
        model_name =  model_conf[scene][model_category]

        print('Start to load the face landmark model...')
        # load model
        try:
            faceAlignModelLoader = FaceAlignModelLoader(model_path, "cuda" if torch.cuda.is_available() else "cpu", model_category, model_name)
        except Exception as e:
            print('Failed to parse model configuration file!')
            print(e)
            sys.exit(-1)
        else:
            print('Successfully parsed the model configuration file model_meta.json!')

        try:
            model, cfg = faceAlignModelLoader.load_model()
        except Exception as e:
            print('Model loading failed!')
            print(e)
            sys.exit(-1)
        else:
            print('Successfully loaded the face landmark model!\n')

        self.faceAlignModelHandler = FaceAlignModelHandler(model, "cuda" if torch.cuda.is_available() else "cpu", cfg)
 

    def __call__(self, image, bbox):
        det = np.asarray(list(map(int, bbox[0:4])), dtype=np.int32)
        landmarks = self.faceAlignModelHandler.inference_on_image(image, det) 
        return landmarks
api_usage/face_alignment.py
MetaInsight7/MaskFaceTool
import sys
sys.path.append('.')

import torch
import yaml
from core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader
from core.model_handler.face_detection.FaceDetModelHandler import FaceDetModelHandler

with open('config/model_conf.yaml') as f:
    model_conf = yaml.load(f, Loader=yaml.FullLoader)

class FaceDet():
    def __init__(self):
        # common setting for all model, need not modify.
        model_path = 'models'

        # model setting, modified along with model
        scene = 'non-mask'
        model_category = 'face_detection'
        model_name =  model_conf[scene][model_category]

        print('Start to load the face detection model...')
        # load model
        try:
            faceDetModelLoader = FaceDetModelLoader(model_path, "cuda:0" if torch.cuda.is_available() else "cpu", model_category, model_name)
        except Exception as e:
            print('Failed to parse model configuration file!')
            print(e)
            sys.exit(-1)
        else:
            print('Successfully parsed the model configuration file model_meta.json!')

        try:
            model, cfg = faceDetModelLoader.load_model()
        except Exception as e:
            print('Model loading failed!')
            print(e)
            sys.exit(-1)
        else:
            print('Successfully loaded the face detection model!\n')

        self.faceDetModelHandler = FaceDetModelHandler(model, "cuda:0" if torch.cuda.is_available() else "cpu", cfg) 
    
    def __call__(self, image):
        try:
            dets = self.faceDetModelHandler.inference_on_image(image)
        except Exception as e:
            print('Face detection failed!')
            print(e)
            sys.exit(-1) 
        bboxs = dets 
        return bboxs
api_usage/face_detect.py
MetaInsight7/MaskFaceTool
"""
@author: Yinglu Liu, Jun Wang
@date: 20201012
@contact: jun21wangustc@gmail.com
"""

import os
from random import randint
import warnings
warnings.filterwarnings('ignore')
import cv2
import torch
import numpy as np
from skimage.io import imread, imsave
from skimage.transform import estimate_transform, warp
from utils import read_info
from models.mask_network.prnet import PRNet
from utils.cython.render import render_cy
from utils.render import render_texture

class PRN:
    """Process of PRNet.
    based on:
    https://github.com/YadiraF/PRNet/blob/master/api.py 
    """
    def __init__(self, model_path):
        self.resolution = 256
        self.MaxPos = self.resolution*1.1        
        self.face_ind = np.loadtxt('Data/uv-data/face_ind.txt').astype(np.int32)
        self.triangles = np.loadtxt('Data/uv-data/triangles.txt').astype(np.int32)
        self.net = PRNet(3, 3)
        state_dict = torch.load(model_path)
        self.net.load_state_dict(state_dict)
        self.net.eval()
        if torch.cuda.is_available():
            self.net = self.net.to('cuda')
    def process(self, image, image_info):
        if np.max(image_info.shape) > 4: # key points to get bounding box
            kpt = image_info
            if kpt.shape[0] > 3:
                kpt = kpt.T
            left = np.min(kpt[0, :]); right = np.max(kpt[0, :]); 
            top = np.min(kpt[1,:]); bottom = np.max(kpt[1,:])
        else:  # bounding box
            bbox = image_info
            left = bbox[0]; right = bbox[1]; top = bbox[2]; bottom = bbox[3]
        old_size = (right - left + bottom - top)/2
        center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0])
        size = int(old_size*1.6)
        # crop image
        src_pts = np.array([[center[0]-size/2, center[1]-size/2], 
                            [center[0] - size/2, center[1]+size/2], 
                            [center[0]+size/2, center[1]-size/2]])
        DST_PTS = np.array([[0,0], [0,self.resolution - 1], [self.resolution - 1, 0]])
        tform = estimate_transform('similarity', src_pts, DST_PTS)
        cropped_image = warp(image, tform.inverse, output_shape=(self.resolution, self.resolution))
        cropped_image = np.transpose(cropped_image[np.newaxis, :,:,:], (0, 3, 1, 2)).astype(np.float32)
        cropped_image = torch.from_numpy(cropped_image)
        if torch.cuda.is_available():
            cropped_image = cropped_image.cuda()
        with torch.no_grad():
            cropped_pos = self.net(cropped_image)
        cropped_pos = cropped_pos.cpu().detach().numpy()
        cropped_pos = np.transpose(cropped_pos, (0, 2, 3, 1)).squeeze() * self.MaxPos
        # restore 
        cropped_vertices = np.reshape(cropped_pos, [-1, 3]).T
        z = cropped_vertices[2,:].copy()/tform.params[0,0]
        cropped_vertices[2,:] = 1
        vertices = np.dot(np.linalg.inv(tform.params), cropped_vertices)
        vertices = np.vstack((vertices[:2,:], z))
        pos = np.reshape(vertices.T, [self.resolution, self.resolution, 3])
        return pos 
    def get_vertices(self, pos):
        all_vertices = np.reshape(pos, [self.resolution ** 2, -1])
        vertices = all_vertices[self.face_ind, :]
        return vertices
    def get_colors_from_texture(self, texture):
        all_colors = np.reshape(texture, [self.resolution**2, -1])
        colors = all_colors[self.face_ind, :]
        return colors

class FaceMasker:
    """Add a virtual mask in face.
    
    Attributes:
        uv_face_path(str): the path of uv_face. 
        mask_template_folder(str): the directory where all mask template in. 
        prn(object): PRN object, https://github.com/YadiraF/PRNet.
        template_name2ref_texture_src(dict): key is template name, value is the mask load by skimage.io.
        template_name2uv_mask_src(dict): key is template name, value is the uv_mask. 
        is_aug(bool): whether or not to add some augmentaion operation on the mask.
    """
    def __init__(self, is_aug):
        """init for FaceMasker
        
        Args:
            is_aug(bool): whether or not to add some augmentaion operation on the mask.
        """
        self.uv_face_path = 'Data/uv-data/uv_face_mask.png'
        self.mask_template_folder = 'Data/mask-data'
        self.prn = PRN('models/mask_network/prnet.pth')
        self.template_name2ref_texture_src, self.template_name2uv_mask_src = self.get_ref_texture_src()
        self.is_aug = is_aug

    def get_ref_texture_src(self):
        template_name2ref_texture_src = {}
        template_name2uv_mask_src = {}
        mask_template_list = os.listdir(self.mask_template_folder)
        uv_face = imread(self.uv_face_path, as_gray=True)/255.
        for mask_template in mask_template_list:
            mask_template_path = os.path.join(self.mask_template_folder, mask_template)
            ref_texture_src = imread(mask_template_path, as_gray=False)/255.
            if ref_texture_src.shape[2] == 4: # must 4 channel, how about 3 channel?
                uv_mask_src = ref_texture_src[:,:,3]
                ref_texture_src = ref_texture_src[:,:,:3]
            else:
                print('Fatal error!', mask_template_path)
            uv_mask_src[uv_face == 0] = 0
            template_name2ref_texture_src[mask_template] = ref_texture_src
            template_name2uv_mask_src[mask_template] = uv_mask_src
        return template_name2ref_texture_src, template_name2uv_mask_src

    def add_mask(self, face_root, image_name2lms, image_name2template_name, masked_face_root):
        for image_name, face_lms in image_name2lms.items():
            image_path = os.path.join(face_root, image_name)
            masked_face_path = os.path.join(masked_face_root, image_name)
            template_name = image_name2template_name[image_name]
            self.add_mask_one(image_path, face_lms, template_name, masked_face_path)

    # you can speed it up by a c++ version.
    def render(self, vertices, new_colors, h, w):
        vis_colors = np.ones((vertices.shape[0], 1))
        face_mask = render_texture(vertices.T, vis_colors.T, self.prn.triangles.T, h, w, c=1).astype(np.uint8)
        face_mask = np.squeeze(face_mask > 0)
        new_image = render_texture(vertices.T, new_colors.T, self.prn.triangles.T, h, w, c=3)
        return face_mask, new_image
        
    def add_mask_one(self, image_path, face_lms, template_name, masked_face_path, speed):
        """Add mask to one image.

        Args:
            image_path(str): the image to add mask.
            face_lms(str): face landmarks, [x1, y1, x2, y2, ..., x106, y106]
            template_name(str): the mask template to be added on the current image, 
                                got to '/Data/mask-data' for all template.
            masked_face_path(str): the path to save masked image.
        """
        image = imread(image_path)
        ref_texture_src = self.template_name2ref_texture_src[template_name] 
        uv_mask_src = self.template_name2uv_mask_src[template_name]
        if image.ndim == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        [h, w, c] = image.shape
        if c == 4:
            image = image[:,:,:3]
        pos, vertices = self.get_vertices(face_lms, image) #3d reconstruction -> get texture. 
        image = image/255. #!!
        texture = cv2.remap(image, pos[:,:,:2].astype(np.float32), None, 
                            interpolation=cv2.INTER_NEAREST, 
                            borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
        new_texture = self.get_new_texture(ref_texture_src, uv_mask_src, texture)
        new_colors = self.prn.get_colors_from_texture(new_texture)
        # render
        if speed:
            face_mask, new_image = render_cy(np.ascontiguousarray(vertices.T), np.ascontiguousarray(new_colors.T), np.ascontiguousarray(self.prn.triangles.T.astype(np.int64)), h, w)
            face_mask = np.squeeze(np.floor(face_mask) > 0)
        else:
            face_mask, new_image = self.render(vertices, new_colors, h, w)
        new_image = image * (1 - face_mask[:, :, np.newaxis]) + new_image * face_mask[:, :, np.newaxis]
        new_image = np.clip(new_image, -1, 1) #must clip to (-1, 1)!

        imsave(masked_face_path, np.uint8(new_image * 255)) 

    def get_vertices(self, face_lms, image):
        """Get vertices

        Args:
            face_lms: face landmarks.
            image:[0, 255]
        """
        lms_info = read_info.read_landmark_106_array(face_lms)
        pos = self.prn.process(image, lms_info) 
        vertices = self.prn.get_vertices(pos)
        return pos, vertices

    def get_new_texture(self, ref_texture_src, uv_mask_src, texture):
        """Get new texture
        Mainly for data augmentation.
        """
        x_offset = 5
        y_offset = 5
        alpha = '0.5,0.8'
        beta = 0
        erode_iter = 5
        
        # random augmentation
        ref_texture = ref_texture_src.copy()
        uv_mask = uv_mask_src.copy()
        if self.is_aug:
            # random flip
            if np.random.rand()>0.5:
                ref_texture = cv2.flip(ref_texture, 1, dst=None)
                uv_mask = cv2.flip(uv_mask, 1, dst=None)
            # random scale, 
            if np.random.rand()>0.5:
                x_offset = np.random.randint(x_offset)
                y_offset = np.random.randint(y_offset)
                ref_texture_temp = np.zeros_like(ref_texture)
                uv_mask_temp = np.zeros_like(uv_mask)
                target_size = (256-x_offset*2, 256-y_offset*2)
                ref_texture_temp[y_offset:256-y_offset, x_offset:256-x_offset,:] = cv2.resize(ref_texture, target_size)
                uv_mask_temp[y_offset:256-y_offset, x_offset:256-x_offset] = cv2.resize(uv_mask, target_size)
                ref_texture = ref_texture_temp
                uv_mask = uv_mask_temp
            # random erode
            if np.random.rand()>0.8:
                t = np.random.randint(erode_iter)
                kernel = np.ones((5,5),np.uint8)
                uv_mask = cv2.erode(uv_mask,kernel,iterations = t)
            # random contrast and brightness
            if np.random.rand()>0.5:
                alpha_r = [float(_) for _ in alpha.split(',')]
                alpha = (alpha_r[1] - alpha_r[0])*np.random.rand() + alpha_r[0]
                beta = beta
                img = ref_texture*255
                blank = np.zeros(img.shape, img.dtype)
                # dst = alpha * img + beta * blank
                dst = cv2.addWeighted(img, alpha, blank, 1-alpha, beta)
                ref_texture = dst.clip(0,255) / 255 
        new_texture = texture*(1 - uv_mask[:,:,np.newaxis]) + ref_texture[:,:,:3]*uv_mask[:,:,np.newaxis]
        return new_texture

api_usage/face_masker.py
MetaInsight7/MaskFaceTool
"""
@author: JiXuan Xu, Jun Wang
@date: 20201023
@contact: jun21wangustc@gmail.com 
"""

import cv2
import torch
import numpy as np
import torch.backends.cudnn as cudnn

from core.model_handler.BaseModelHandler import BaseModelHandler
from utils.BuzException import *
from torchvision import transforms

class FaceAlignModelHandler(BaseModelHandler):
    """Implementation of face landmark model handler

    Attributes:
        model: the face landmark model.
        device: use cpu or gpu to process.
        cfg(dict): testing config, inherit from the parent class.
    """
    def __init__(self, model, device, cfg):
        """
        Init FaceLmsModelHandler settings. 
        """
        super().__init__(model, device, cfg)
        self.img_size = self.cfg['img_size']
        
    def inference_on_image(self, image, dets):
        """Get the inference of the image and process the inference result.

        Returns:
            A numpy array, the landmarks prediction based on the shape of original image, shape: (106, 2), 
        """
        cudnn.benchmark = True
        try:
            image_pre = self._preprocess(image, dets)
        except Exception as e:
            raise e
        self.model = self.model.to(self.device)
        image_pre = image_pre.unsqueeze(0)
        with torch.no_grad():
            image_pre = image_pre.to(self.device)
            _, landmarks_normal = self.model(image_pre)
        landmarks = self._postprocess(landmarks_normal)
        return landmarks

    #Adapted from https://github.com/Hsintao/pfld_106_face_landmarks/blob/master/data/prepare.py
    def _preprocess(self, image, det):
        """Preprocess the input image, cutting the input image through the face detection information.
        Using the face detection result(dets) to get the face position in the input image.
        After determining the center of face position and the box size of face, crop the image
        and resize it into preset size.

        Returns:
           A torch tensor, the image after preprecess, shape: (3, 112, 112).
        """
        if not isinstance(image, np.ndarray):
            print('The input should be the ndarray read by cv2!')
            raise InputError()
        img = image.copy()
        self.image_org = image.copy()
        img = np.float32(img)

        xy = np.array([det[0], det[1]])
        zz = np.array([det[2], det[3]])
        wh = zz - xy + 1
        center = (xy + wh / 2).astype(np.int32)
        boxsize = int(np.max(wh) * 1.2)
        xy = center - boxsize // 2
        self.xy = xy
        self.boxsize = boxsize
        x1, y1 = xy
        x2, y2 = xy + boxsize
        height, width, _ = img.shape
        dx = max(0, -x1)
        dy = max(0, -y1)
        x1 = max(0, x1)
        y1 = max(0, y1)
        edx = max(0, x2 - width)
        edy = max(0, y2 - height)
        x2 = min(width, x2)
        y2 = min(height, y2)
        imageT = image[y1:y2, x1:x2]
        if dx > 0 or dy > 0 or edx > 0 or edy > 0:
            imageT = cv2.copyMakeBorder(
                imageT, dy, edy, dx, edx, cv2.BORDER_CONSTANT, 0)

        imageT = cv2.resize(imageT, (self.img_size, self.img_size))
        t = transforms.Compose([transforms.ToTensor()])
        img_after = t(imageT)
        return img_after

    def _postprocess(self, landmarks_normal):
        """Process the predicted landmarks into the form of the original image.

        Returns:
            A numpy array, the landmarks based on the shape of original image, shape: (106, 2), 
        """    
        landmarks_normal = landmarks_normal.cpu().numpy()
        landmarks_normal = landmarks_normal.reshape(landmarks_normal.shape[0], -1, 2)
        landmarks = landmarks_normal[0] * [self.boxsize, self.boxsize] + self.xy
        return landmarks

core/model_handler/face_alignment/FaceAlignModelHandler.py
MetaInsight7/MaskFaceTool
"""
@author: JiXuan Xu, Jun Wang
@date: 20201019
@contact: jun21wangustc@gmail.com 
"""

import torch
import numpy as np
from math import ceil
from itertools import product as product
import torch.backends.cudnn as cudnn

from core.model_handler.BaseModelHandler import BaseModelHandler
from utils.BuzException import *


class FaceDetModelHandler(BaseModelHandler):
    """Implementation of face detection model handler

    Attributes:
        model: the face detection model.
        device: use cpu or gpu to process.
        cfg(dict): testing config, inherit from the parent class.
    """
    def __init__(self, model, device, cfg):
        """
        Init FaceDetModelHandler settings. 
        """
        super().__init__(model, device, cfg)
        self.variance = self.cfg['variance']
        
    def inference_on_image(self, image):
        """Get the inference of the image and process the inference result.

        Returns:
            A numpy array, the shape is N * (x, y, w, h, confidence), 
            N is the number of detection box.
        """
        cudnn.benchmark = True
        input_height, input_width, _ = image.shape
        try:
            image, scale = self._preprocess(image)
        except Exception as e:
            raise e
        self.model = self.model.to(self.device)
        image = torch.from_numpy(image).unsqueeze(0)
        with torch.no_grad():
            image = image.to(self.device)
            scale = scale.to(self.device)
            loc, conf, landms = self.model(image)
        dets = self._postprocess(loc, conf, scale, input_height, input_width)
        return dets

    def _preprocess(self, image):
        """Preprocess the image, such as standardization and other operations.

        Returns:
            A numpy array list, the shape is channel * h * w.
            A tensor, the shape is 4.
        """
        if not isinstance(image, np.ndarray):
            print('The input should be the ndarray read by cv2!')
            raise InputError()
        img = np.float32(image)
        scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
        img -= (104, 117, 123)
        img = img.transpose(2, 0, 1)
        return img, scale

    def _postprocess(self, loc, conf, scale, input_height, input_width):
        """Postprecess the prediction result.
        Decode detection result, set the confidence threshold and do the NMS
        to keep the appropriate detection box. 

        Returns:
            A numpy array, the shape is N * (x, y, w, h, confidence), 
            N is the number of detection box.
        """
        priorbox = PriorBox(self.cfg, image_size=(input_height, input_width))
        priors = priorbox.forward()
        priors = priors.to(self.device)
        prior_data = priors.data
        boxes = self.decode(loc.data.squeeze(0), prior_data, self.cfg['variance'])
        boxes = boxes * scale
        boxes = boxes.cpu().numpy()
        scores = conf.squeeze(0).data.cpu().numpy()[:, 1]

        # ignore low scores
        inds = np.where(scores > self.cfg['confidence_threshold'])[0]
        boxes = boxes[inds]
        scores = scores[inds]

        # keep top-K before NMS
        order = scores.argsort()[::-1]
        boxes = boxes[order]
        scores = scores[order]

        # do NMS
        nms_threshold = 0.2
        dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
        keep = self.py_cpu_nms(dets, nms_threshold)
        dets = dets[keep, :]
        return dets

    # Adapted from https://github.com/chainer/chainercv
    def decode(self, loc, priors, variances):
        """Decode locations from predictions using priors to undo
        the encoding we did for offset regression at train time.
        Args:
            loc (tensor): location predictions for loc layers,
                Shape: [num_priors,4]
            priors (tensor): Prior boxes in center-offset form.
                Shape: [num_priors,4].
            variances: (list[float]) Variances of priorboxes

        Return:
            decoded bounding box predictions
        """
        boxes = torch.cat((priors[:, :2], priors[:, 2:]), 1)
        boxes[:, :2] = priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:]
        boxes[:, 2:] = priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])
        boxes[:, :2] -= boxes[:, 2:] / 2
        boxes[:, 2:] += boxes[:, :2]
        return boxes

    # Adapted from https://github.com/biubug6/Pytorch_Retinaface
    def py_cpu_nms(self, dets, thresh):
        """Python version NMS.

        Returns:
            The kept index after NMS.
        """
        x1 = dets[:, 0]
        y1 = dets[:, 1]
        x2 = dets[:, 2]
        y2 = dets[:, 3]
        scores = dets[:, 4]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]
        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            inds = np.where(ovr <= thresh)[0]
            order = order[inds + 1]
        return keep

# Adapted from https://github.com/biubug6/Pytorch_Retinafacey
class PriorBox(object):
    """Compute the suitable parameters of anchors for later decode operation

    Attributes:
        cfg(dict): testing config.
        image_size(tuple): the input image size.
    """
    def __init__(self, cfg, image_size=None):
        """
        Init priorBox settings related to the generation of anchors. 
        """
        super(PriorBox, self).__init__()
        self.min_sizes = cfg['min_sizes']
        self.steps = cfg['steps']
        self.image_size = image_size
        self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]
        self.name = "s"

    def forward(self):
        anchors = []
        for k, f in enumerate(self.feature_maps):
            min_sizes = self.min_sizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / self.image_size[1]
                    s_ky = min_size / self.image_size[0]
                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]
        # back to torch land
        output = torch.Tensor(anchors).view(-1, 4)
        return output

core/model_handler/face_detection/FaceDetModelHandler.py
MetaInsight7/MaskFaceTool
"""
@author: JiXuan Xu, Jun Wang
@date: 20201015
@contact: jun21wangustc@gmail.com 
"""
import os
import sys
sys.path.append('models/network_def')

from abc import ABCMeta, abstractmethod

import json

class BaseModelLoader(metaclass=ABCMeta):
    """Base class for all model loader.
    All the model loaders need to inherit this base class, 
    and each new model needs to implement the "load model" method
    """
    def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'):
        model_root_dir = os.path.join(model_path, model_category, model_name)
        meta_file_path = os.path.join(model_root_dir, meta_file)
        self.cfg = {}
        try:
            self.meta_conf = json.load(open(meta_file_path, 'r'))
        except IOError as e:
            raise e
        except Exception as e:
            raise e
        # common configs for all model
        self.cfg['model_path'] = model_path
        self.cfg['model_category'] = model_category
        self.cfg['model_name'] = model_name
        self.cfg['model_type'] = self.meta_conf['model_type']
        self.cfg['model_info'] = self.meta_conf['model_info']
        self.cfg['model_file_path'] = os.path.join(model_root_dir, self.meta_conf['model_file'])
        self.cfg['release_date'] = self.meta_conf['release_date']
        self.cfg['input_height'] = self.meta_conf['input_height']
        self.cfg['input_width'] = self.meta_conf['input_width']

    @abstractmethod
    def load_model(self):
        """Should be overridden by all subclasses.
        Different models may have different configuration information,
        such as mean, so each model implements its own loader
        """
        pass

core/model_loader/BaseModelLoader.py
MetaInsight7/MaskFaceTool
"""
@author: JiXuan Xu, Jun Wang
@date: 20201023
@contact: jun21wangustc@gmail.com 
"""


import torch

from core.model_loader.BaseModelLoader import BaseModelLoader

class FaceAlignModelLoader(BaseModelLoader):
    def __init__(self, model_path, device, model_category, model_name, meta_file='model_meta.json'):
        super().__init__(model_path, model_category, model_name, meta_file)
        self.cfg['img_size'] = self.meta_conf['input_width']
        self.device = device
        
    def load_model(self):
        try:
            model = torch.load(self.cfg['model_file_path'], map_location=self.device)
        except Exception as e:
            raise e
        else:
            return model, self.cfg

core/model_loader/face_alignment/FaceAlignModelLoader.py
MetaInsight7/MaskFaceTool
"""
@author: JiXuan Xu, Jun Wang
@date: 20201019
@contact: jun21wangustc@gmail.com 
"""


import torch

from core.model_loader.BaseModelLoader import BaseModelLoader

class FaceDetModelLoader(BaseModelLoader):
    def __init__(self, model_path, device, model_category, model_name, meta_file='model_meta.json'):
        super().__init__(model_path, model_category, model_name, meta_file)
        self.cfg['min_sizes'] = self.meta_conf['min_sizes']
        self.cfg['steps'] = self.meta_conf['steps']
        self.cfg['variance'] = self.meta_conf['variance']
        self.cfg['in_channel'] = self.meta_conf['in_channel']
        self.cfg['out_channel'] = self.meta_conf['out_channel']
        self.cfg['confidence_threshold'] = self.meta_conf['confidence_threshold']
        self.device = device
        
    def load_model(self):
        try:
            model = torch.load(self.cfg['model_file_path'], map_location=self.device)
        except Exception as e:
            raise e
        else:
            return model, self.cfg

core/model_loader/face_detection/FaceDetModelLoader.py
MetaInsight7/MaskFaceTool