| | import os
|
| | import cv2
|
| | import math
|
| | import random
|
| | import numpy as np
|
| | from skimage import transform
|
| |
|
| |
|
| | class Augmentation:
|
| | def __init__(self,
|
| | is_train=True,
|
| | aug_prob=1.0,
|
| | image_size=256,
|
| | crop_op=True,
|
| | std_lmk_5pts=None,
|
| | target_face_scale=1.0,
|
| | flip_rate=0.5,
|
| | flip_mapping=None,
|
| | random_shift_sigma=0.05,
|
| | random_rot_sigma=math.pi/180*18,
|
| | random_scale_sigma=0.1,
|
| | random_gray_rate=0.2,
|
| | random_occ_rate=0.4,
|
| | random_blur_rate=0.3,
|
| | random_gamma_rate=0.2,
|
| | random_nose_fusion_rate=0.2):
|
| | self.is_train = is_train
|
| | self.aug_prob = aug_prob
|
| | self.crop_op = crop_op
|
| | self._flip = Flip(flip_mapping, flip_rate)
|
| | if self.crop_op:
|
| | self._cropMatrix = GetCropMatrix(
|
| | image_size=image_size,
|
| | target_face_scale=target_face_scale,
|
| | align_corners=True)
|
| | else:
|
| | self._alignMatrix = GetAlignMatrix(
|
| | image_size=image_size,
|
| | target_face_scale=target_face_scale,
|
| | std_lmk_5pts=std_lmk_5pts)
|
| | self._randomGeometryMatrix = GetRandomGeometryMatrix(
|
| | target_shape=(image_size, image_size),
|
| | from_shape=(image_size, image_size),
|
| | shift_sigma=random_shift_sigma,
|
| | rot_sigma=random_rot_sigma,
|
| | scale_sigma=random_scale_sigma,
|
| | align_corners=True)
|
| | self._transform = Transform(image_size=image_size)
|
| | self._randomTexture = RandomTexture(
|
| | random_gray_rate=random_gray_rate,
|
| | random_occ_rate=random_occ_rate,
|
| | random_blur_rate=random_blur_rate,
|
| | random_gamma_rate=random_gamma_rate,
|
| | random_nose_fusion_rate=random_nose_fusion_rate)
|
| |
|
| | def process(self, img, lmk, lmk_5pts=None, scale=1.0, center_w=0, center_h=0, is_train=True):
|
| | if self.is_train and random.random() < self.aug_prob:
|
| | img, lmk, lmk_5pts, center_w, center_h = self._flip.process(img, lmk, lmk_5pts, center_w, center_h)
|
| | matrix_geoaug = self._randomGeometryMatrix.process()
|
| | if self.crop_op:
|
| | matrix_pre = self._cropMatrix.process(scale, center_w, center_h)
|
| | else:
|
| | matrix_pre = self._alignMatrix.process(lmk_5pts)
|
| | matrix = matrix_geoaug @ matrix_pre
|
| | aug_img, aug_lmk = self._transform.process(img, lmk, matrix)
|
| | aug_img = self._randomTexture.process(aug_img)
|
| | else:
|
| | if self.crop_op:
|
| | matrix = self._cropMatrix.process(scale, center_w, center_h)
|
| | else:
|
| | matrix = self._alignMatrix.process(lmk_5pts)
|
| | aug_img, aug_lmk = self._transform.process(img, lmk, matrix)
|
| | return aug_img, aug_lmk, matrix
|
| |
|
| |
|
| | class GetCropMatrix:
|
| | def __init__(self, image_size, target_face_scale, align_corners=False):
|
| | self.image_size = image_size
|
| | self.target_face_scale = target_face_scale
|
| | self.align_corners = align_corners
|
| |
|
| | def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center):
|
| | cosv = math.cos(angle)
|
| | sinv = math.sin(angle)
|
| |
|
| | fx, fy = from_center
|
| | tx, ty = to_center
|
| |
|
| | acos = scale * cosv
|
| | asin = scale * sinv
|
| |
|
| | a0 = acos
|
| | a1 = -asin
|
| | a2 = tx - acos * fx + asin * fy + shift_xy[0]
|
| |
|
| | b0 = asin
|
| | b1 = acos
|
| | b2 = ty - asin * fx - acos * fy + shift_xy[1]
|
| |
|
| | rot_scale_m = np.array([
|
| | [a0, a1, a2],
|
| | [b0, b1, b2],
|
| | [0.0, 0.0, 1.0]
|
| | ], np.float32)
|
| | return rot_scale_m
|
| |
|
| | def process(self, scale, center_w, center_h):
|
| | if self.align_corners:
|
| | to_w, to_h = self.image_size-1, self.image_size-1
|
| | else:
|
| | to_w, to_h = self.image_size, self.image_size
|
| |
|
| | rot_mu = 0
|
| | scale_mu = self.image_size / (scale * self.target_face_scale * 200.0)
|
| | shift_xy_mu = (0, 0)
|
| | matrix = self._compose_rotate_and_scale(
|
| | rot_mu, scale_mu, shift_xy_mu,
|
| | from_center=[center_w, center_h],
|
| | to_center=[to_w/2.0, to_h/2.0])
|
| | return matrix
|
| |
|
| |
|
| | class GetAlignMatrix:
|
| | def __init__(self, image_size, target_face_scale, std_lmk_5pts):
|
| | """
|
| | points in std_lmk_5pts range from -1 to 1.
|
| | """
|
| | self.std_lmk_5pts = (std_lmk_5pts * target_face_scale + 1) * \
|
| | np.array([image_size, image_size], np.float32) / 2.0
|
| |
|
| | def process(self, lmk_5pts):
|
| | assert lmk_5pts.shape[-2:] == (5, 2)
|
| | tform = transform.SimilarityTransform()
|
| | tform.estimate(lmk_5pts, self.std_lmk_5pts)
|
| | return tform.params
|
| |
|
| |
|
| | class GetRandomGeometryMatrix:
|
| | def __init__(self, target_shape, from_shape,
|
| | shift_sigma=0.1, rot_sigma=18*math.pi/180, scale_sigma=0.1,
|
| | shift_mu=0.0, rot_mu=0.0, scale_mu=1.0,
|
| | shift_normal=True, rot_normal=True, scale_normal=True,
|
| | align_corners=False):
|
| | self.target_shape = target_shape
|
| | self.from_shape = from_shape
|
| | self.shift_config = (shift_mu, shift_sigma, shift_normal)
|
| | self.rot_config = (rot_mu, rot_sigma, rot_normal)
|
| | self.scale_config = (scale_mu, scale_sigma, scale_normal)
|
| | self.align_corners = align_corners
|
| |
|
| | def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center):
|
| | cosv = math.cos(angle)
|
| | sinv = math.sin(angle)
|
| |
|
| | fx, fy = from_center
|
| | tx, ty = to_center
|
| |
|
| | acos = scale * cosv
|
| | asin = scale * sinv
|
| |
|
| | a0 = acos
|
| | a1 = -asin
|
| | a2 = tx - acos * fx + asin * fy + shift_xy[0]
|
| |
|
| | b0 = asin
|
| | b1 = acos
|
| | b2 = ty - asin * fx - acos * fy + shift_xy[1]
|
| |
|
| | rot_scale_m = np.array([
|
| | [a0, a1, a2],
|
| | [b0, b1, b2],
|
| | [0.0, 0.0, 1.0]
|
| | ], np.float32)
|
| | return rot_scale_m
|
| |
|
| | def _random(self, mu_sigma_normal, size=None):
|
| | mu, sigma, is_normal = mu_sigma_normal
|
| | if is_normal:
|
| | return np.random.normal(mu, sigma, size=size)
|
| | else:
|
| | return np.random.uniform(low=mu-sigma, high=mu+sigma, size=size)
|
| |
|
| | def process(self):
|
| | if self.align_corners:
|
| | from_w, from_h = self.from_shape[1]-1, self.from_shape[0]-1
|
| | to_w, to_h = self.target_shape[1]-1, self.target_shape[0]-1
|
| | else:
|
| | from_w, from_h = self.from_shape[1], self.from_shape[0]
|
| | to_w, to_h = self.target_shape[1], self.target_shape[0]
|
| |
|
| | if self.shift_config[:2] != (0.0, 0.0) or \
|
| | self.rot_config[:2] != (0.0, 0.0) or \
|
| | self.scale_config[:2] != (1.0, 0.0):
|
| | shift_xy = self._random(self.shift_config, size=[2]) * \
|
| | min(to_h, to_w)
|
| | rot_angle = self._random(self.rot_config)
|
| | scale = self._random(self.scale_config)
|
| | matrix_geoaug = self._compose_rotate_and_scale(
|
| | rot_angle, scale, shift_xy,
|
| | from_center=[from_w/2.0, from_h/2.0],
|
| | to_center=[to_w/2.0, to_h/2.0])
|
| |
|
| | return matrix_geoaug
|
| |
|
| |
|
| | class Transform:
|
| | def __init__(self, image_size):
|
| | self.image_size = image_size
|
| |
|
| | def _transformPoints2D(self, points, matrix):
|
| | """
|
| | points (nx2), matrix (3x3) -> points (nx2)
|
| | """
|
| | dtype = points.dtype
|
| |
|
| |
|
| | points = np.concatenate([points, np.ones_like(points[:, [0]])], axis=1)
|
| | points = points @ np.transpose(matrix)
|
| | points = points[:, :2] / points[:, [2, 2]]
|
| | return points.astype(dtype)
|
| |
|
| | def _transformPerspective(self, image, matrix):
|
| | """
|
| | image, matrix3x3 -> transformed_image
|
| | """
|
| | return cv2.warpPerspective(
|
| | image, matrix,
|
| | dsize=(self.image_size, self.image_size),
|
| | flags=cv2.INTER_LINEAR, borderValue=0)
|
| |
|
| | def process(self, image, landmarks, matrix):
|
| | t_landmarks = self._transformPoints2D(landmarks, matrix)
|
| | t_image = self._transformPerspective(image, matrix)
|
| | return t_image, t_landmarks
|
| |
|
| |
|
| | class RandomTexture:
|
| | def __init__(self, random_gray_rate=0, random_occ_rate=0, random_blur_rate=0, random_gamma_rate=0, random_nose_fusion_rate=0):
|
| | self.random_gray_rate = random_gray_rate
|
| | self.random_occ_rate = random_occ_rate
|
| | self.random_blur_rate = random_blur_rate
|
| | self.random_gamma_rate = random_gamma_rate
|
| | self.random_nose_fusion_rate = random_nose_fusion_rate
|
| | self.texture_augs = (
|
| | (self.add_occ, self.random_occ_rate),
|
| | (self.add_blur, self.random_blur_rate),
|
| | (self.add_gamma, self.random_gamma_rate),
|
| | (self.add_nose_fusion, self.random_nose_fusion_rate)
|
| | )
|
| |
|
| | def add_gray(self, image):
|
| | assert image.ndim == 3 and image.shape[-1] == 3
|
| | image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
| | image = np.tile(np.expand_dims(image, -1), [1, 1, 3])
|
| | return image
|
| |
|
| | def add_occ(self, image):
|
| | h, w, c = image.shape
|
| | rh = 0.2 + 0.6 * random.random()
|
| | rw = rh - 0.2 + 0.4 * random.random()
|
| | cx = int((h - 1) * random.random())
|
| | cy = int((w - 1) * random.random())
|
| | dh = int(h / 2 * rh)
|
| | dw = int(w / 2 * rw)
|
| | x0 = max(0, cx - dw // 2)
|
| | y0 = max(0, cy - dh // 2)
|
| | x1 = min(w - 1, cx + dw // 2)
|
| | y1 = min(h - 1, cy + dh // 2)
|
| | image[y0:y1+1, x0:x1+1] = 0
|
| | return image
|
| |
|
| | def add_blur(self, image):
|
| | blur_kratio = 0.05 * random.random()
|
| | blur_ksize = int((image.shape[0] + image.shape[1]) / 2 * blur_kratio)
|
| | if blur_ksize > 1:
|
| | image = cv2.blur(image, (blur_ksize, blur_ksize))
|
| | return image
|
| |
|
| | def add_gamma(self, image):
|
| | if random.random() < 0.5:
|
| | gamma = 0.25 + 0.75 * random.random()
|
| | else:
|
| | gamma = 1.0 + 3.0 * random.random()
|
| | image = (((image / 255.0) ** gamma) * 255).astype("uint8")
|
| | return image
|
| |
|
| | def add_nose_fusion(self, image):
|
| | h, w, c = image.shape
|
| | nose = np.array(bytearray(os.urandom(h * w * c)), dtype=image.dtype).reshape(h, w, c)
|
| | alpha = 0.5 * random.random()
|
| | image = (1 - alpha) * image + alpha * nose
|
| | return image.astype(np.uint8)
|
| |
|
| | def process(self, image):
|
| | image = image.copy()
|
| | if random.random() < self.random_occ_rate:
|
| | image = self.add_occ(image)
|
| | if random.random() < self.random_blur_rate:
|
| | image = self.add_blur(image)
|
| | if random.random() < self.random_gamma_rate:
|
| | image = self.add_gamma(image)
|
| | if random.random() < self.random_nose_fusion_rate:
|
| | image = self.add_nose_fusion(image)
|
| | """
|
| | orders = list(range(len(self.texture_augs)))
|
| | random.shuffle(orders)
|
| | for order in orders:
|
| | if random.random() < self.texture_augs[order][1]:
|
| | image = self.texture_augs[order][0](image)
|
| | """
|
| |
|
| | if random.random() < self.random_gray_rate:
|
| | image = self.add_gray(image)
|
| |
|
| | return image
|
| |
|
| |
|
| | class Flip:
|
| | def __init__(self, flip_mapping, random_rate):
|
| | self.flip_mapping = flip_mapping
|
| | self.random_rate = random_rate
|
| |
|
| | def process(self, image, landmarks, landmarks_5pts, center_w, center_h):
|
| | if random.random() >= self.random_rate or self.flip_mapping is None:
|
| | return image, landmarks, landmarks_5pts, center_w, center_h
|
| |
|
| |
|
| | if landmarks.shape[0] == 29:
|
| | flip_offset = 0
|
| |
|
| | elif landmarks.shape[0] in (68, 98):
|
| | flip_offset = -1
|
| | else:
|
| | flip_offset = -1
|
| |
|
| | h, w, _ = image.shape
|
| |
|
| | image_flip = np.fliplr(image).copy()
|
| | landmarks_flip = landmarks.copy()
|
| | for i, j in self.flip_mapping:
|
| | landmarks_flip[i] = landmarks[j]
|
| | landmarks_flip[j] = landmarks[i]
|
| | landmarks_flip[:, 0] = w + flip_offset - landmarks_flip[:, 0]
|
| | if landmarks_5pts is not None:
|
| | flip_mapping = ([0, 1], [3, 4])
|
| | landmarks_5pts_flip = landmarks_5pts.copy()
|
| | for i, j in flip_mapping:
|
| | landmarks_5pts_flip[i] = landmarks_5pts[j]
|
| | landmarks_5pts_flip[j] = landmarks_5pts[i]
|
| | landmarks_5pts_flip[:, 0] = w + flip_offset - landmarks_5pts_flip[:, 0]
|
| | else:
|
| | landmarks_5pts_flip = None
|
| |
|
| | center_w = w + flip_offset - center_w
|
| | return image_flip, landmarks_flip, landmarks_5pts_flip, center_w, center_h
|
| |
|
