face_masker.py

"""
@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 model.prnet import PRNet
from utils.cython.render import render_cy
import time

class PRN:
    """Process of PRNet.
    based on:
    https://github.com/YadiraF/PRNet/blob/master/api.py 
    """
    def __init__(self, model_path, device):
        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)
        self.device = device
        state_dict = torch.load(model_path, map_location=self.device)
        self.net.load_state_dict(state_dict)
        self.net.eval()
        # if torch.cuda.is_available():
        self.net.to(self.device)

    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_image = cropped_image.to(self.device)
            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, device):
        """init for FaceMasker
        
        Args:
            is_aug(bool): whether or not to add some augmentaion operation on the mask.
        """
        self.device = device
        self.uv_face_path = 'Data/uv-data/uv_face_mask.png'
        self.mask_template_folder = 'Data/mask-data'
        self.prn = PRN('model/prnet.pth', device = self.device)
        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:
            # print('Create UV map for template: ', mask_template)
            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, face_lms, template_name, masked_face_path, padded = None, write_image = True, pos_vertices = None):
        """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)
        # t1 = time.time()
        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]
        if pos_vertices is None:
            pos, vertices = self.get_vertices(face_lms, image) #3d reconstruction -> get texture. 
        else:
            print('Found exists vertices, use this params')
            pos, vertices = pos_vertices
        image = image/255. #!!
        texture = cv2.remap(image, pos[:,:,:2].astype(np.float32), None, 
                            interpolation=cv2.INTER_NEAREST, 
                            borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
        # print(texture.shape)
        # imsave('texture.jpg', texture) 
        # t2 = time.time()
        new_texture = self.get_new_texture(ref_texture_src, uv_mask_src, texture)
        new_colors = self.prn.get_colors_from_texture(new_texture)
        # print('Render cy')
        # t3 = time.time()
        # render
        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)
        # t4 = time.time()
        # imsave('face_mask.jpg', face_mask) 
        # imsave('new_image.jpg', new_image) 
        # print('Render done')
        face_mask = np.squeeze(np.floor(face_mask) > 0)
        
        tmp = new_image * face_mask[:, :, np.newaxis]
        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)!
        t5 = time.time()
        # print('[FaceMasker] Time preprocess: ', t2 - t1)
        # print('[FaceMasker] Time feed: ', t3 - t2)
        # print('[FaceMasker] Time render: ', t4 - t3)
        # print('[FaceMasker] Time post-process: ', t5 - t4)
        if padded is not None:
            if write_image:
                imsave(masked_face_path, new_image[padded:-padded, padded:-padded, :]) 
            return new_image[padded:-padded, padded:-padded, :]
        else:
            if write_image:
                imsave(masked_face_path, new_image) 
            return new_image
         

    def create_mask_one(self, image, image_segment, face_lms, output):
        """ Create mask for single input image
        """
        # image = imread(image_path)
        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_segment, pos[:,:,:2].astype(np.float32), None, 
                            interpolation=cv2.INTER_NEAREST, 
                            borderMode=cv2.BORDER_CONSTANT,borderValue=(0))

        imsave(output, texture) 

    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

class FaceMaskerMP:
    """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, device, n_processes = 4, max_queue_len = 64):
        """init for FaceMasker
        
        Args:
            is_aug(bool): whether or not to add some augmentaion operation on the mask.
        """
        self.device = device
        self.uv_face_path = 'Data/uv-data/uv_face_mask.png'
        self.mask_template_folder = 'Data/mask-data'
        self.prn = PRN('model/prnet.pth', device = self.device)
        self.template_name2ref_texture_src, self.template_name2uv_mask_src = self.get_ref_texture_src()
        self.is_aug = is_aug
        # self.n_processes = n_processes
        # self.q_in = [multiprocessing.Queue(max_queue_len) for i in range(self.n_processes)]
        # q_out = multiprocessing.Queue(max_queue_len)
        

    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:
            # print('Create UV map for template: ', mask_template)
            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, face_lms, template_name, masked_face_path, padded = None, write_image = True):
        """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)
        t1 = time.time()
        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))
        # print(texture.shape)
        # imsave('texture.jpg', texture) 
        t2 = time.time()
        new_texture = self.get_new_texture(ref_texture_src, uv_mask_src, texture)
        new_colors = self.prn.get_colors_from_texture(new_texture)
        # print('Render cy')
        t3 = time.time()
        # render
        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)
        t4 = time.time()
        # imsave('face_mask.jpg', face_mask) 
        # imsave('new_image.jpg', new_image) 
        # print('Render done')
        face_mask = np.squeeze(np.floor(face_mask) > 0)
        
        tmp = new_image * face_mask[:, :, np.newaxis]
        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)!
        t5 = time.time()
        # print('[FaceMasker] Time preprocess: ', t2 - t1)
        # print('[FaceMasker] Time feed: ', t3 - t2)
        # print('[FaceMasker] Time render: ', t4 - t3)
        # print('[FaceMasker] Time post-process: ', t5 - t4)
        if padded is not None:
            if write_image:
                imsave(masked_face_path, new_image[padded:-padded, padded:-padded, :]) 
            return new_image[padded:-padded, padded:-padded, :]
        else:
            if write_image:
                imsave(masked_face_path, new_image) 
            return new_image
    

    def mask_precompute(self, image, face_lms, template_name, masked_face_path, padded = None, write_image = True):
        """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]
        t1 = time.time()
        pos, vertices = self.get_vertices(face_lms, image) #3d reconstruction -> get texture. 
        t2 = time.time()
        print('get vertices: ', t2 - t1)
        image = image/255. #!!
        texture = cv2.remap(image, pos[:,:,:2].astype(np.float32), None, 
                            interpolation=cv2.INTER_NEAREST, 
                            borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
        # print(texture.shape)
        # imsave('texture.jpg', texture) 
        new_texture = self.get_new_texture(ref_texture_src, uv_mask_src, texture)
        new_colors = self.prn.get_colors_from_texture(new_texture)
        # print('Render cy')
        return (image, vertices, new_colors, self.prn.triangles, h, w)

    def mask_render(self, image, vertices, new_colors, triangles, h, w):
        face_mask, new_image = render_cy(np.ascontiguousarray(vertices.T), np.ascontiguousarray(new_colors.T), np.ascontiguousarray(triangles.T.astype(np.int64)), h, w)
        face_mask = np.squeeze(np.floor(face_mask) > 0)
        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)!
        return new_image


    def create_mask_one(self, image, image_segment, face_lms, output):
        """ Create mask for single input image
        """
        # image = imread(image_path)
        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_segment, pos[:,:,:2].astype(np.float32), None, 
                            interpolation=cv2.INTER_NEAREST, 
                            borderMode=cv2.BORDER_CONSTANT,borderValue=(0))

        imsave(output, texture) 

    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