|
|
|
|
|
|
|
|
|
|
|
|
|
import torch
|
|
from torchvision import datasets,transforms,utils
|
|
from torch.utils.data import Dataset,IterableDataset
|
|
from glob import glob
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
import random
|
|
import cv2
|
|
from torch import nn
|
|
import sys
|
|
import scipy as sp
|
|
from skimage.measure import label, regionprops
|
|
from training.dataset.library.bi_online_generation import random_get_hull
|
|
import albumentations as alb
|
|
|
|
import warnings
|
|
warnings.filterwarnings('ignore')
|
|
|
|
|
|
def alpha_blend(source,target,mask):
|
|
mask_blured = get_blend_mask(mask)
|
|
img_blended=(mask_blured * source + (1 - mask_blured) * target)
|
|
return img_blended,mask_blured
|
|
|
|
|
|
def dynamic_blend(source,target,mask):
|
|
mask_blured = get_blend_mask(mask)
|
|
blend_list=[0.25,0.5,0.75,1,1,1]
|
|
blend_ratio = blend_list[np.random.randint(len(blend_list))]
|
|
mask_blured*=blend_ratio
|
|
img_blended=(mask_blured * source + (1 - mask_blured) * target)
|
|
return img_blended,mask_blured
|
|
|
|
|
|
def get_blend_mask(mask):
|
|
H,W=mask.shape
|
|
size_h=np.random.randint(192,257)
|
|
size_w=np.random.randint(192,257)
|
|
mask=cv2.resize(mask,(size_w,size_h))
|
|
kernel_1=random.randrange(5,26,2)
|
|
kernel_1=(kernel_1,kernel_1)
|
|
kernel_2=random.randrange(5,26,2)
|
|
kernel_2=(kernel_2,kernel_2)
|
|
|
|
mask_blured = cv2.GaussianBlur(mask, kernel_1, 0)
|
|
mask_blured = mask_blured/(mask_blured.max())
|
|
mask_blured[mask_blured<1]=0
|
|
|
|
mask_blured = cv2.GaussianBlur(mask_blured, kernel_2, np.random.randint(5,46))
|
|
mask_blured = mask_blured/(mask_blured.max())
|
|
mask_blured = cv2.resize(mask_blured,(W,H))
|
|
return mask_blured.reshape((mask_blured.shape+(1,)))
|
|
|
|
|
|
def get_alpha_blend_mask(mask):
|
|
kernel_list=[(11,11),(9,9),(7,7),(5,5),(3,3)]
|
|
blend_list=[0.25,0.5,0.75]
|
|
kernel_idxs=random.choices(range(len(kernel_list)), k=2)
|
|
blend_ratio = blend_list[random.sample(range(len(blend_list)), 1)[0]]
|
|
mask_blured = cv2.GaussianBlur(mask, kernel_list[0], 0)
|
|
|
|
mask_blured[mask_blured<mask_blured.max()]=0
|
|
mask_blured[mask_blured>0]=1
|
|
|
|
mask_blured = cv2.GaussianBlur(mask_blured, kernel_list[kernel_idxs[1]], 0)
|
|
mask_blured = mask_blured/(mask_blured.max())
|
|
return mask_blured.reshape((mask_blured.shape+(1,)))
|
|
|
|
|
|
class RandomDownScale(alb.core.transforms_interface.ImageOnlyTransform):
|
|
def apply(self,img,**params):
|
|
return self.randomdownscale(img)
|
|
|
|
def randomdownscale(self,img):
|
|
keep_ratio=True
|
|
keep_input_shape=True
|
|
H,W,C=img.shape
|
|
ratio_list=[2,4]
|
|
r=ratio_list[np.random.randint(len(ratio_list))]
|
|
img_ds=cv2.resize(img,(int(W/r),int(H/r)),interpolation=cv2.INTER_NEAREST)
|
|
if keep_input_shape:
|
|
img_ds=cv2.resize(img_ds,(W,H),interpolation=cv2.INTER_LINEAR)
|
|
|
|
return img_ds
|
|
|
|
|
|
|
|
def get_boundary(mask, apply_dilation=True, apply_motion_blur=True):
|
|
if len(mask.shape) == 3:
|
|
mask = mask[:, :, 0]
|
|
|
|
mask = cv2.GaussianBlur(mask, (3, 3), 0)
|
|
if mask.max() > 1:
|
|
boundary = mask / 255.
|
|
else:
|
|
boundary = mask
|
|
boundary = 4 * boundary * (1. - boundary)
|
|
|
|
boundary = boundary * 255
|
|
boundary = random_dilate(boundary)
|
|
|
|
if apply_motion_blur:
|
|
boundary = random_motion_blur(boundary)
|
|
boundary = boundary / 255.
|
|
return boundary
|
|
|
|
def random_dilate(mask, max_kernel_size=5):
|
|
kernel_size = random.randint(1, max_kernel_size)
|
|
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
|
dilated_mask = cv2.dilate(mask, kernel, iterations=1)
|
|
return dilated_mask
|
|
|
|
def random_motion_blur(mask, max_kernel_size=5):
|
|
kernel_size = random.randint(1, max_kernel_size)
|
|
kernel = np.zeros((kernel_size, kernel_size))
|
|
anchor = random.randint(0, kernel_size - 1)
|
|
kernel[:, anchor] = 1 / kernel_size
|
|
motion_blurred_mask = cv2.filter2D(mask, -1, kernel)
|
|
return motion_blurred_mask
|
|
|
|
|
|
|
|
class SBI_API:
|
|
def __init__(self,phase='train',image_size=256):
|
|
|
|
assert phase == 'train', f"Current SBI API only support train phase, but got {phase}"
|
|
|
|
self.image_size=(image_size,image_size)
|
|
self.phase=phase
|
|
|
|
self.transforms=self.get_transforms()
|
|
self.source_transforms = self.get_source_transforms()
|
|
self.bob_transforms = self.get_source_transforms_for_bob()
|
|
|
|
|
|
def __call__(self,img,landmark=None):
|
|
try:
|
|
assert landmark is not None, "landmark of the facial image should not be None."
|
|
|
|
|
|
if random.random() < 1.0:
|
|
|
|
img_r,img_f,mask_f=self.self_blending(img.copy(),landmark.copy())
|
|
else:
|
|
|
|
img_r,img_f,mask_f=self.bob(img.copy(),landmark.copy())
|
|
|
|
if self.phase=='train':
|
|
transformed=self.transforms(image=img_f.astype('uint8'),image1=img_r.astype('uint8'))
|
|
img_f=transformed['image']
|
|
img_r=transformed['image1']
|
|
return img_f,img_r
|
|
except Exception as e:
|
|
print(e)
|
|
return None,None
|
|
|
|
|
|
def get_source_transforms(self):
|
|
return alb.Compose([
|
|
alb.Compose([
|
|
alb.RGBShift((-20,20),(-20,20),(-20,20),p=0.3),
|
|
alb.HueSaturationValue(hue_shift_limit=(-0.3,0.3), sat_shift_limit=(-0.3,0.3), val_shift_limit=(-0.3,0.3), p=1),
|
|
alb.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1,0.1), p=1),
|
|
],p=1),
|
|
|
|
alb.OneOf([
|
|
RandomDownScale(p=1),
|
|
alb.Sharpen(alpha=(0.2, 0.5), lightness=(0.5, 1.0), p=1),
|
|
],p=1),
|
|
|
|
], p=1.)
|
|
|
|
|
|
def get_transforms(self):
|
|
return alb.Compose([
|
|
|
|
alb.RGBShift((-20,20),(-20,20),(-20,20),p=0.3),
|
|
alb.HueSaturationValue(hue_shift_limit=(-0.3,0.3), sat_shift_limit=(-0.3,0.3), val_shift_limit=(-0.3,0.3), p=0.3),
|
|
alb.RandomBrightnessContrast(brightness_limit=(-0.3,0.3), contrast_limit=(-0.3,0.3), p=0.3),
|
|
alb.ImageCompression(quality_lower=40,quality_upper=100,p=0.5),
|
|
|
|
],
|
|
additional_targets={f'image1': 'image'},
|
|
p=1.)
|
|
|
|
|
|
def randaffine(self,img,mask):
|
|
f=alb.Affine(
|
|
translate_percent={'x':(-0.03,0.03),'y':(-0.015,0.015)},
|
|
scale=[0.95,1/0.95],
|
|
fit_output=False,
|
|
p=1)
|
|
|
|
g=alb.ElasticTransform(
|
|
alpha=50,
|
|
sigma=7,
|
|
alpha_affine=0,
|
|
p=1,
|
|
)
|
|
|
|
transformed=f(image=img,mask=mask)
|
|
img=transformed['image']
|
|
|
|
mask=transformed['mask']
|
|
transformed=g(image=img,mask=mask)
|
|
mask=transformed['mask']
|
|
return img,mask
|
|
|
|
|
|
def get_source_transforms_for_bob(self):
|
|
return alb.Compose([
|
|
alb.Compose([
|
|
alb.ImageCompression(quality_lower=40,quality_upper=100,p=1),
|
|
],p=1),
|
|
|
|
alb.OneOf([
|
|
RandomDownScale(p=1),
|
|
alb.Sharpen(alpha=(0.2, 0.5), lightness=(0.5, 1.0), p=1),
|
|
],p=1),
|
|
|
|
], p=1.)
|
|
|
|
def bob(self,img,landmark):
|
|
H,W=len(img),len(img[0])
|
|
if np.random.rand()<0.25:
|
|
landmark=landmark[:68]
|
|
|
|
|
|
hull_type = random.choice([0, 1, 2, 3])
|
|
mask=random_get_hull(landmark,img,hull_type)[:,:,0]
|
|
|
|
source = img.copy()
|
|
source = self.bob_transforms(image=source.astype(np.uint8))['image']
|
|
source, mask = self.randaffine(source,mask)
|
|
mask = get_blend_mask(mask)
|
|
|
|
|
|
boundary = get_boundary(mask)
|
|
|
|
blend_list = [0.25,0.5,0.75,1,1,1]
|
|
blend_ratio = blend_list[np.random.randint(len(blend_list))]
|
|
boundary *= blend_ratio
|
|
boundary = np.repeat(boundary[:, :, np.newaxis], 3, axis=2)
|
|
img_blended = (boundary * source + (1 - boundary) * img)
|
|
|
|
img_blended = img_blended.astype(np.uint8)
|
|
img = img.astype(np.uint8)
|
|
|
|
return img,img_blended,boundary.squeeze()
|
|
|
|
|
|
def self_blending(self,img,landmark):
|
|
H,W=len(img),len(img[0])
|
|
if np.random.rand()<0.25:
|
|
landmark=landmark[:68]
|
|
|
|
|
|
hull_type = random.choice([0, 1, 2, 3])
|
|
mask=random_get_hull(landmark,img,hull_type)[:,:,0]
|
|
|
|
source = img.copy()
|
|
if np.random.rand()<0.5:
|
|
source = self.source_transforms(image=source.astype(np.uint8))['image']
|
|
else:
|
|
img = self.source_transforms(image=img.astype(np.uint8))['image']
|
|
|
|
source, mask = self.randaffine(source,mask)
|
|
|
|
img_blended,mask=dynamic_blend(source,img,mask)
|
|
img_blended = img_blended.astype(np.uint8)
|
|
img = img.astype(np.uint8)
|
|
|
|
return img,img_blended,mask
|
|
|
|
|
|
def reorder_landmark(self,landmark):
|
|
landmark_add=np.zeros((13,2))
|
|
for idx,idx_l in enumerate([77,75,76,68,69,70,71,80,72,73,79,74,78]):
|
|
landmark_add[idx]=landmark[idx_l]
|
|
landmark[68:]=landmark_add
|
|
return landmark
|
|
|
|
|
|
def hflip(self,img,mask=None,landmark=None,bbox=None):
|
|
H,W=img.shape[:2]
|
|
landmark=landmark.copy()
|
|
if bbox is not None:
|
|
bbox=bbox.copy()
|
|
|
|
if landmark is not None:
|
|
landmark_new=np.zeros_like(landmark)
|
|
|
|
|
|
landmark_new[:17]=landmark[:17][::-1]
|
|
landmark_new[17:27]=landmark[17:27][::-1]
|
|
|
|
landmark_new[27:31]=landmark[27:31]
|
|
landmark_new[31:36]=landmark[31:36][::-1]
|
|
|
|
landmark_new[36:40]=landmark[42:46][::-1]
|
|
landmark_new[40:42]=landmark[46:48][::-1]
|
|
|
|
landmark_new[42:46]=landmark[36:40][::-1]
|
|
landmark_new[46:48]=landmark[40:42][::-1]
|
|
|
|
landmark_new[48:55]=landmark[48:55][::-1]
|
|
landmark_new[55:60]=landmark[55:60][::-1]
|
|
|
|
landmark_new[60:65]=landmark[60:65][::-1]
|
|
landmark_new[65:68]=landmark[65:68][::-1]
|
|
if len(landmark)==68:
|
|
pass
|
|
elif len(landmark)==81:
|
|
landmark_new[68:81]=landmark[68:81][::-1]
|
|
else:
|
|
raise NotImplementedError
|
|
landmark_new[:,0]=W-landmark_new[:,0]
|
|
|
|
else:
|
|
landmark_new=None
|
|
|
|
if bbox is not None:
|
|
bbox_new=np.zeros_like(bbox)
|
|
bbox_new[0,0]=bbox[1,0]
|
|
bbox_new[1,0]=bbox[0,0]
|
|
bbox_new[:,0]=W-bbox_new[:,0]
|
|
bbox_new[:,1]=bbox[:,1].copy()
|
|
if len(bbox)>2:
|
|
bbox_new[2,0]=W-bbox[3,0]
|
|
bbox_new[2,1]=bbox[3,1]
|
|
bbox_new[3,0]=W-bbox[2,0]
|
|
bbox_new[3,1]=bbox[2,1]
|
|
bbox_new[4,0]=W-bbox[4,0]
|
|
bbox_new[4,1]=bbox[4,1]
|
|
bbox_new[5,0]=W-bbox[6,0]
|
|
bbox_new[5,1]=bbox[6,1]
|
|
bbox_new[6,0]=W-bbox[5,0]
|
|
bbox_new[6,1]=bbox[5,1]
|
|
else:
|
|
bbox_new=None
|
|
|
|
if mask is not None:
|
|
mask=mask[:,::-1]
|
|
else:
|
|
mask=None
|
|
img=img[:,::-1].copy()
|
|
return img,mask,landmark_new,bbox_new
|
|
|
|
|
|
if __name__=='__main__':
|
|
seed=10
|
|
random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
torch.cuda.manual_seed(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
api=SBI_API(phase='train',image_size=256)
|
|
|
|
img_path = 'FaceForensics++/original_sequences/youtube/c23/frames/000/000.png'
|
|
img = cv2.imread(img_path)
|
|
landmark_path = img_path.replace('frames', 'landmarks').replace('png', 'npy')
|
|
landmark = np.load(landmark_path)
|
|
sbi_img, ori_img = api(img, landmark)
|
|
|
