first commit
This commit is contained in:
Your Name
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
Copyright StrangeAI authors @2019
|
||||
|
||||
assume you have to directly which you want
|
||||
convert A to B, just put all faces of A person to A,
|
||||
faces of B person to B
|
||||
|
||||
"""
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
import glob
|
||||
import os
|
||||
from alfred.dl.torch.common import device
|
||||
import cv2
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
import numpy as np
|
||||
from utils.umeyama import umeyama
|
||||
import cv2
|
||||
|
||||
random_transform_args = {
|
||||
'rotation_range': 10,
|
||||
'zoom_range': 0.05,
|
||||
'shift_range': 0.05,
|
||||
'random_flip': 0.4,
|
||||
}
|
||||
|
||||
|
||||
def random_transform(image, rotation_range, zoom_range, shift_range, random_flip):
|
||||
h, w = image.shape[0:2]
|
||||
rotation = np.random.uniform(-rotation_range, rotation_range)
|
||||
scale = np.random.uniform(1 - zoom_range, 1 + zoom_range)
|
||||
tx = np.random.uniform(-shift_range, shift_range) * w
|
||||
ty = np.random.uniform(-shift_range, shift_range) * h
|
||||
mat = cv2.getRotationMatrix2D((w // 2, h // 2), rotation, scale)
|
||||
mat[:, 2] += (tx, ty)
|
||||
result = cv2.warpAffine(image, mat, (w, h), borderMode=cv2.BORDER_REPLICATE)
|
||||
if np.random.random() < random_flip:
|
||||
result = result[:, ::-1]
|
||||
return result
|
||||
|
||||
|
||||
def random_warp_128(image):
|
||||
assert image.shape == (256, 256, 3), 'resize image to 256 256 first'
|
||||
range_ = np.linspace(128 - 120, 128 + 120, 9)
|
||||
mapx = np.broadcast_to(range_, (9, 9))
|
||||
mapy = mapx.T
|
||||
mapx = mapx + np.random.normal(size=(9, 9), scale=5)
|
||||
mapy = mapy + np.random.normal(size=(9, 9), scale=5)
|
||||
interp_mapx = cv2.resize(mapx, (144, 144))[8:136, 8:136].astype('float32')
|
||||
interp_mapy = cv2.resize(mapy, (144, 144))[8:136, 8:136].astype('float32')
|
||||
warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR)
|
||||
src_points = np.stack([mapx.ravel(), mapy.ravel()], axis=-1)
|
||||
dst_points = np.mgrid[0:129:16, 0:129:16].T.reshape(-1, 2)
|
||||
mat = umeyama(src_points, dst_points, True)[0:2]
|
||||
target_image = cv2.warpAffine(image, mat, (128, 128))
|
||||
return warped_image, target_image
|
||||
|
||||
|
||||
def random_warp_64(image):
|
||||
assert image.shape == (256, 256, 3)
|
||||
range_ = np.linspace(128 - 120, 128 + 120, 5)
|
||||
mapx = np.broadcast_to(range_, (5, 5))
|
||||
mapy = mapx.T
|
||||
mapx = mapx + np.random.normal(size=(5, 5), scale=5)
|
||||
mapy = mapy + np.random.normal(size=(5, 5), scale=5)
|
||||
interp_mapx = cv2.resize(mapx, (80, 80))[8:72, 8:72].astype('float32')
|
||||
interp_mapy = cv2.resize(mapy, (80, 80))[8:72, 8:72].astype('float32')
|
||||
warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR)
|
||||
src_points = np.stack([mapx.ravel(), mapy.ravel()], axis=-1)
|
||||
dst_points = np.mgrid[0:65:16, 0:65:16].T.reshape(-1, 2)
|
||||
mat = umeyama(src_points, dst_points, True)[0:2]
|
||||
target_image = cv2.warpAffine(image, mat, (64, 64))
|
||||
return warped_image, target_image
|
||||
|
||||
|
||||
class FacePairDataset(Dataset):
|
||||
|
||||
def __init__(self, a_dir, b_dir, target_size, transform):
|
||||
super(FacePairDataset, self).__init__
|
||||
self.a_dir = a_dir
|
||||
self.b_dir = b_dir
|
||||
self.target_size = target_size
|
||||
|
||||
self.transform = transform
|
||||
# extension can be changed here to png or others
|
||||
self.a_images_list = glob.glob(os.path.join(a_dir, '*.png'))
|
||||
self.b_images_list = glob.glob(os.path.join(b_dir, '*.png'))
|
||||
|
||||
def __getitem__(self, index):
|
||||
# return 2 image pair, A and B
|
||||
img_a = Image.open(self.a_images_list[index])
|
||||
img_b = Image.open(self.b_images_list[index])
|
||||
|
||||
# align the face first
|
||||
img_a = img_a.resize((self.target_size, self.target_size), Image.ANTIALIAS)
|
||||
img_b = img_b.resize((self.target_size, self.target_size), Image.ANTIALIAS)
|
||||
|
||||
# transform
|
||||
if self.transform:
|
||||
img_a = self.transform(img_a)
|
||||
img_b = self.transform(img_b)
|
||||
|
||||
# already resized, warp it
|
||||
img_a = random_transform(np.array(img_a), **random_transform_args)
|
||||
img_b = random_transform(np.array(img_b), **random_transform_args)
|
||||
img_a_input, img_a = random_warp(np.array(img_a), 256)
|
||||
img_b_input, img_b = random_warp(np.array(img_b), 256)
|
||||
img_a_tensor = torch.Tensor(img_a.transpose(2, 0, 1)/255.).float()
|
||||
img_a_input_tensor = torch.Tensor(img_a_input.transpose(2, 0, 1)/255.).float()
|
||||
img_b_tensor = torch.Tensor(img_b.transpose(2, 0, 1)/255.).float()
|
||||
img_b_input_tensor = torch.Tensor(img_b_input.transpose(2, 0, 1)/255.).float()
|
||||
return img_a_tensor, img_a_input_tensor, img_b_tensor, img_b_input_tensor
|
||||
|
||||
def __len__(self):
|
||||
return min(len(self.a_images_list), len(self.b_images_list))
|
||||
|
||||
|
||||
|
||||
class FacePairDataset64x64(Dataset):
|
||||
|
||||
def __init__(self, a_dir, b_dir, target_size, transform):
|
||||
super(FacePairDataset64x64, self).__init__
|
||||
self.a_dir = a_dir
|
||||
self.b_dir = b_dir
|
||||
self.target_size = target_size
|
||||
|
||||
self.transform = transform
|
||||
# extension can be changed here to png or others
|
||||
self.a_images_list = glob.glob(os.path.join(a_dir, '*.png'))
|
||||
self.b_images_list = glob.glob(os.path.join(b_dir, '*.png'))
|
||||
|
||||
def __getitem__(self, index):
|
||||
# return 2 image pair, A and B
|
||||
img_a = Image.open(self.a_images_list[index])
|
||||
img_b = Image.open(self.b_images_list[index])
|
||||
|
||||
# align the face first
|
||||
img_a = img_a.resize((256, 256), Image.ANTIALIAS)
|
||||
img_b = img_b.resize((256, 256), Image.ANTIALIAS)
|
||||
|
||||
# transform
|
||||
if self.transform:
|
||||
img_a = self.transform(img_a)
|
||||
img_b = self.transform(img_b)
|
||||
|
||||
# # already resized, warp it
|
||||
img_a = random_transform(np.array(img_a), **random_transform_args)
|
||||
img_b = random_transform(np.array(img_b), **random_transform_args)
|
||||
img_a_input, img_a = random_warp_64(np.array(img_a))
|
||||
img_b_input, img_b = random_warp_64(np.array(img_b))
|
||||
|
||||
img_a = np.array(img_a)
|
||||
img_b = np.array(img_b)
|
||||
|
||||
img_a_tensor = torch.Tensor(img_a.transpose(2, 0, 1)/255.).float()
|
||||
img_a_input_tensor = torch.Tensor(img_a_input.transpose(2, 0, 1)/255.).float()
|
||||
img_b_tensor = torch.Tensor(img_b.transpose(2, 0, 1)/255.).float()
|
||||
img_b_input_tensor = torch.Tensor(img_b_input.transpose(2, 0, 1)/255.).float()
|
||||
return img_a_tensor, img_a_input_tensor, img_b_tensor, img_b_input_tensor
|
||||
|
||||
def __len__(self):
|
||||
return min(len(self.a_images_list), len(self.b_images_list))
|
||||
|
||||
|
||||
|
||||
class FacePairDataset128x128(Dataset):
|
||||
|
||||
def __init__(self, a_dir, b_dir, target_size, transform):
|
||||
super(FacePairDataset128x128, self).__init__
|
||||
self.a_dir = a_dir
|
||||
self.b_dir = b_dir
|
||||
self.target_size = target_size
|
||||
|
||||
self.transform = transform
|
||||
self.a_images_list = glob.glob(os.path.join(a_dir, '*.png'))
|
||||
self.b_images_list = glob.glob(os.path.join(b_dir, '*.png'))
|
||||
|
||||
def __getitem__(self, index):
|
||||
# return 2 image pair, A and B
|
||||
img_a = Image.open(self.a_images_list[index])
|
||||
img_b = Image.open(self.b_images_list[index])
|
||||
|
||||
# align the face first
|
||||
img_a = img_a.resize((256, 256), Image.ANTIALIAS)
|
||||
img_b = img_b.resize((256, 256), Image.ANTIALIAS)
|
||||
|
||||
# transform
|
||||
if self.transform:
|
||||
img_a = self.transform(img_a)
|
||||
img_b = self.transform(img_b)
|
||||
|
||||
img_a = random_transform(np.array(img_a), **random_transform_args)
|
||||
img_b = random_transform(np.array(img_b), **random_transform_args)
|
||||
img_a_input, img_a = random_warp_128(np.array(img_a))
|
||||
img_b_input, img_b = random_warp_128(np.array(img_b))
|
||||
img_a_tensor = torch.Tensor(img_a.transpose(2, 0, 1)/255.).float()
|
||||
img_a_input_tensor = torch.Tensor(img_a_input.transpose(2, 0, 1)/255.).float()
|
||||
img_b_tensor = torch.Tensor(img_b.transpose(2, 0, 1)/255.).float()
|
||||
img_b_input_tensor = torch.Tensor(img_b_input.transpose(2, 0, 1)/255.).float()
|
||||
return img_a_tensor, img_a_input_tensor, img_b_tensor, img_b_input_tensor
|
||||
|
||||
def __len__(self):
|
||||
return min(len(self.a_images_list), len(self.b_images_list))
|
||||
@@ -0,0 +1,61 @@
|
||||
import numpy
|
||||
from utils.umeyama import umeyama
|
||||
import cv2
|
||||
|
||||
random_transform_args = {
|
||||
'rotation_range': 10,
|
||||
'zoom_range': 0.05,
|
||||
'shift_range': 0.05,
|
||||
'random_flip': 0.4,
|
||||
}
|
||||
|
||||
|
||||
def random_transform(image, rotation_range, zoom_range, shift_range, random_flip):
|
||||
h, w = image.shape[0:2]
|
||||
rotation = numpy.random.uniform(-rotation_range, rotation_range)
|
||||
scale = numpy.random.uniform(1 - zoom_range, 1 + zoom_range)
|
||||
tx = numpy.random.uniform(-shift_range, shift_range) * w
|
||||
ty = numpy.random.uniform(-shift_range, shift_range) * h
|
||||
mat = cv2.getRotationMatrix2D((w // 2, h // 2), rotation, scale)
|
||||
mat[:, 2] += (tx, ty)
|
||||
result = cv2.warpAffine(image, mat, (w, h), borderMode=cv2.BORDER_REPLICATE)
|
||||
if numpy.random.random() < random_flip:
|
||||
result = result[:, ::-1]
|
||||
return result
|
||||
|
||||
|
||||
# get pair of random warped images from aligened face image
|
||||
def random_warp(image):
|
||||
assert image.shape == (256, 256, 3)
|
||||
range_ = numpy.linspace(128 - 80, 128 + 80, 5)
|
||||
mapx = numpy.broadcast_to(range_, (5, 5))
|
||||
mapy = mapx.T
|
||||
|
||||
mapx = mapx + numpy.random.normal(size=(5, 5), scale=5)
|
||||
mapy = mapy + numpy.random.normal(size=(5, 5), scale=5)
|
||||
|
||||
interp_mapx = cv2.resize(mapx, (80, 80))[8:72, 8:72].astype('float32')
|
||||
interp_mapy = cv2.resize(mapy, (80, 80))[8:72, 8:72].astype('float32')
|
||||
|
||||
# just crop the image, remove the top left bottom right 8 pixels (in order to get the pure face)
|
||||
warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR)
|
||||
|
||||
src_points = numpy.stack([mapx.ravel(), mapy.ravel()], axis=-1)
|
||||
dst_points = numpy.mgrid[0:65:16, 0:65:16].T.reshape(-1, 2)
|
||||
mat = umeyama(src_points, dst_points, True)[0:2]
|
||||
target_image = cv2.warpAffine(image, mat, (64, 64))
|
||||
return warped_image, target_image
|
||||
|
||||
|
||||
def get_training_data(images, batch_size):
|
||||
indices = numpy.random.randint(len(images), size=batch_size)
|
||||
for i, index in enumerate(indices):
|
||||
image = images[index]
|
||||
image = random_transform(image, **random_transform_args)
|
||||
warped_img, target_img = random_warp(image)
|
||||
if i == 0:
|
||||
warped_images = numpy.empty((batch_size,) + warped_img.shape, warped_img.dtype)
|
||||
target_images = numpy.empty((batch_size,) + target_img.shape, warped_img.dtype)
|
||||
warped_images[i] = warped_img
|
||||
target_images[i] = target_img
|
||||
return warped_images, target_images
|
||||
Reference in New Issue
Block a user