training/dataset/library/bi_online_generation.py

import dlib
from skimage import io
from skimage import transform as sktransform
import numpy as np
from matplotlib import pyplot as plt
import json
import os
import random
from PIL import Image
from imgaug import augmenters as iaa
from .DeepFakeMask import dfl_full,facehull,components,extended
import cv2
import tqdm

def name_resolve(path):
    name = os.path.splitext(os.path.basename(path))[0]
    vid_id, frame_id = name.split('_')[0:2]
    return vid_id, frame_id 
    
def total_euclidean_distance(a,b):
    assert len(a.shape) == 2
    return np.sum(np.linalg.norm(a-b,axis=1))

def random_get_hull(landmark,img1,hull_type):
    if hull_type == 0:
        mask = dfl_full(landmarks=landmark.astype('int32'),face=img1, channels=3).mask
        return mask/255
    elif hull_type == 1:
        mask = extended(landmarks=landmark.astype('int32'),face=img1, channels=3).mask
        return mask/255
    elif hull_type == 2:
        mask = components(landmarks=landmark.astype('int32'),face=img1, channels=3).mask
        return mask/255
    elif hull_type == 3:
        mask = facehull(landmarks=landmark.astype('int32'),face=img1, channels=3).mask
        return mask/255

def random_erode_dilate(mask, ksize=None):
    if random.random()>0.5:
        if ksize is  None:
            ksize = random.randint(1,21)
        if ksize % 2 == 0:
            ksize += 1
        mask = np.array(mask).astype(np.uint8)*255
        kernel = np.ones((ksize,ksize),np.uint8)
        mask = cv2.erode(mask,kernel,1)/255
    else:
        if ksize is  None:
            ksize = random.randint(1,5)
        if ksize % 2 == 0:
            ksize += 1
        mask = np.array(mask).astype(np.uint8)*255
        kernel = np.ones((ksize,ksize),np.uint8)
        mask = cv2.dilate(mask,kernel,1)/255
    return mask


# borrow from https://github.com/MarekKowalski/FaceSwap
def blendImages(src, dst, mask, featherAmount=0.2):
   
    maskIndices = np.where(mask != 0)
    
    src_mask = np.ones_like(mask)
    dst_mask = np.zeros_like(mask)

    maskPts = np.hstack((maskIndices[1][:, np.newaxis], maskIndices[0][:, np.newaxis]))
    faceSize = np.max(maskPts, axis=0) - np.min(maskPts, axis=0)
    featherAmount = featherAmount * np.max(faceSize)

    hull = cv2.convexHull(maskPts)
    dists = np.zeros(maskPts.shape[0])
    for i in range(maskPts.shape[0]):
        dists[i] = cv2.pointPolygonTest(hull, (maskPts[i, 0], maskPts[i, 1]), True)

    weights = np.clip(dists / featherAmount, 0, 1)

    composedImg = np.copy(dst)
    composedImg[maskIndices[0], maskIndices[1]] = weights[:, np.newaxis] * src[maskIndices[0], maskIndices[1]] + (1 - weights[:, np.newaxis]) * dst[maskIndices[0], maskIndices[1]]

    composedMask = np.copy(dst_mask)
    composedMask[maskIndices[0], maskIndices[1]] = weights[:, np.newaxis] * src_mask[maskIndices[0], maskIndices[1]] + (
                1 - weights[:, np.newaxis]) * dst_mask[maskIndices[0], maskIndices[1]]

    return composedImg, composedMask


# borrow from https://github.com/MarekKowalski/FaceSwap
def colorTransfer(src, dst, mask):
    transferredDst = np.copy(dst)
    
    maskIndices = np.where(mask != 0)
    

    maskedSrc = src[maskIndices[0], maskIndices[1]].astype(np.int32)
    maskedDst = dst[maskIndices[0], maskIndices[1]].astype(np.int32)

    meanSrc = np.mean(maskedSrc, axis=0)
    meanDst = np.mean(maskedDst, axis=0)

    maskedDst = maskedDst - meanDst
    maskedDst = maskedDst + meanSrc
    maskedDst = np.clip(maskedDst, 0, 255)

    transferredDst[maskIndices[0], maskIndices[1]] = maskedDst

    return transferredDst

class BIOnlineGeneration():
    def __init__(self):
        with open('precomuted_landmarks.json', 'r') as f:
            self.landmarks_record =  json.load(f)
            for k,v in self.landmarks_record.items():
                self.landmarks_record[k] = np.array(v)
        # extract all frame from all video in the name of {videoid}_{frameid}
        self.data_list = [
                    '000_0000.png',
                    '001_0000.png'      
                    ] * 10000
        
        # predefine mask distortion
        self.distortion = iaa.Sequential([iaa.PiecewiseAffine(scale=(0.01, 0.15))])
        
    def gen_one_datapoint(self):
        background_face_path = random.choice(self.data_list)
        data_type = 'real' if random.randint(0,1) else 'fake'
        if data_type == 'fake' :
            face_img,mask =  self.get_blended_face(background_face_path)
            mask = ( 1 - mask ) * mask * 4
        else:
            face_img = io.imread(background_face_path)
            mask = np.zeros((317, 317, 1))
        
        # randomly downsample after BI pipeline
        if random.randint(0,1):
            aug_size = random.randint(64, 317)
            face_img = Image.fromarray(face_img)
            if random.randint(0,1):
                face_img = face_img.resize((aug_size, aug_size), Image.BILINEAR)
            else:
                face_img = face_img.resize((aug_size, aug_size), Image.NEAREST)
            face_img = face_img.resize((317, 317),Image.BILINEAR)
            face_img = np.array(face_img)
            
        # random jpeg compression after BI pipeline
        if random.randint(0,1):
            quality = random.randint(60, 100)
            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
            face_img_encode = cv2.imencode('.jpg', face_img, encode_param)[1]
            face_img = cv2.imdecode(face_img_encode, cv2.IMREAD_COLOR)
        
        face_img = face_img[60:317,30:287,:]
        mask = mask[60:317,30:287,:]
        
        # random flip
        if random.randint(0,1):
            face_img = np.flip(face_img,1)
            mask = np.flip(mask,1)
            
        return face_img,mask,data_type
        
    def get_blended_face(self,background_face_path):
        background_face = io.imread(background_face_path)
        background_landmark = self.landmarks_record[background_face_path]
        
        foreground_face_path = self.search_similar_face(background_landmark,background_face_path)
        foreground_face = io.imread(foreground_face_path)
        
        # down sample before blending
        aug_size = random.randint(128,317)
        background_landmark = background_landmark * (aug_size/317)
        foreground_face = sktransform.resize(foreground_face,(aug_size,aug_size),preserve_range=True).astype(np.uint8)
        background_face = sktransform.resize(background_face,(aug_size,aug_size),preserve_range=True).astype(np.uint8)
        
        # get random type of initial blending mask
        mask = random_get_hull(background_landmark, background_face)
       
        #  random deform mask
        mask = self.distortion.augment_image(mask)
        mask = random_erode_dilate(mask)
        
        # filte empty mask after deformation
        if np.sum(mask) == 0 :
            raise NotImplementedError

        # apply color transfer
        foreground_face = colorTransfer(background_face, foreground_face, mask*255)
        
        # blend two face
        blended_face, mask = blendImages(foreground_face, background_face, mask*255)
        blended_face = blended_face.astype(np.uint8)
       
        # resize back to default resolution
        blended_face = sktransform.resize(blended_face,(317,317),preserve_range=True).astype(np.uint8)
        mask = sktransform.resize(mask,(317,317),preserve_range=True)
        mask = mask[:,:,0:1]
        return blended_face,mask
    
    def search_similar_face(self,this_landmark,background_face_path):
        vid_id, frame_id = name_resolve(background_face_path)
        min_dist = 99999999
        
        # random sample 5000 frame from all frams:
        all_candidate_path = random.sample( self.data_list, k=5000) 
        
        # filter all frame that comes from the same video as background face
        all_candidate_path = filter(lambda k:name_resolve(k)[0] != vid_id, all_candidate_path)
        all_candidate_path = list(all_candidate_path)
        
        # loop throungh all candidates frame to get best match
        for candidate_path in all_candidate_path:
            candidate_landmark = self.landmarks_record[candidate_path].astype(np.float32)
            candidate_distance = total_euclidean_distance(candidate_landmark, this_landmark)
            if candidate_distance < min_dist:
                min_dist = candidate_distance
                min_path = candidate_path

        return min_path
    
if __name__ == '__main__':
    ds = BIOnlineGeneration()
    from tqdm import tqdm
    all_imgs = []
    for _ in tqdm(range(50)):
        img,mask,label = ds.gen_one_datapoint()
        mask = np.repeat(mask,3,2)
        mask = (mask*255).astype(np.uint8)
        img_cat = np.concatenate([img,mask],1)
        all_imgs.append(img_cat)
    all_in_one = Image.new('RGB', (2570,2570))

    for x in range(5):
        for y in range(10):
            idx = x*10+y   
            im = Image.fromarray(all_imgs[idx])
            
            dx = x*514
            dy = y*257
            
            all_in_one.paste(im, (dx,dy))

    all_in_one.save("all_in_one.jpg")