62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
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
|