import math import torch from torch import autograd from torch.nn import functional as F import numpy as np from model.stylegan.distributed import reduce_sum from model.stylegan.op import upfirdn2d class AdaptiveAugment: def __init__(self, ada_aug_target, ada_aug_len, update_every, device): self.ada_aug_target = ada_aug_target self.ada_aug_len = ada_aug_len self.update_every = update_every self.ada_update = 0 self.ada_aug_buf = torch.tensor([0.0, 0.0], device=device) self.r_t_stat = 0 self.ada_aug_p = 0 @torch.no_grad() def tune(self, real_pred): self.ada_aug_buf += torch.tensor( (torch.sign(real_pred).sum().item(), real_pred.shape[0]), device=real_pred.device, ) self.ada_update += 1 if self.ada_update % self.update_every == 0: self.ada_aug_buf = reduce_sum(self.ada_aug_buf) pred_signs, n_pred = self.ada_aug_buf.tolist() self.r_t_stat = pred_signs / n_pred if self.r_t_stat > self.ada_aug_target: sign = 1 else: sign = -1 self.ada_aug_p += sign * n_pred / self.ada_aug_len self.ada_aug_p = min(1, max(0, self.ada_aug_p)) self.ada_aug_buf.mul_(0) self.ada_update = 0 return self.ada_aug_p SYM6 = ( 0.015404109327027373, 0.0034907120842174702, -0.11799011114819057, -0.048311742585633, 0.4910559419267466, 0.787641141030194, 0.3379294217276218, -0.07263752278646252, -0.021060292512300564, 0.04472490177066578, 0.0017677118642428036, -0.007800708325034148, ) def translate_mat(t_x, t_y, device="cpu"): batch = t_x.shape[0] mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1) translate = torch.stack((t_x, t_y), 1) mat[:, :2, 2] = translate return mat def rotate_mat(theta, device="cpu"): batch = theta.shape[0] mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1) sin_t = torch.sin(theta) cos_t = torch.cos(theta) rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2) mat[:, :2, :2] = rot return mat def scale_mat(s_x, s_y, device="cpu"): batch = s_x.shape[0] mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1) mat[:, 0, 0] = s_x mat[:, 1, 1] = s_y return mat def translate3d_mat(t_x, t_y, t_z): batch = t_x.shape[0] mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1) translate = torch.stack((t_x, t_y, t_z), 1) mat[:, :3, 3] = translate return mat def rotate3d_mat(axis, theta): batch = theta.shape[0] u_x, u_y, u_z = axis eye = torch.eye(3).unsqueeze(0) cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0) outer = torch.tensor(axis) outer = (outer.unsqueeze(1) * outer).unsqueeze(0) sin_t = torch.sin(theta).view(-1, 1, 1) cos_t = torch.cos(theta).view(-1, 1, 1) rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1) eye_4[:, :3, :3] = rot return eye_4 def scale3d_mat(s_x, s_y, s_z): batch = s_x.shape[0] mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1) mat[:, 0, 0] = s_x mat[:, 1, 1] = s_y mat[:, 2, 2] = s_z return mat def luma_flip_mat(axis, i): batch = i.shape[0] eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1) axis = torch.tensor(axis + (0,)) flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1) return eye - flip def saturation_mat(axis, i): batch = i.shape[0] eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1) axis = torch.tensor(axis + (0,)) axis = torch.ger(axis, axis) saturate = axis + (eye - axis) * i.view(-1, 1, 1) return saturate def lognormal_sample(size, mean=0, std=1, device="cpu"): return torch.empty(size, device=device).log_normal_(mean=mean, std=std) def category_sample(size, categories, device="cpu"): category = torch.tensor(categories, device=device) sample = torch.randint(high=len(categories), size=(size,), device=device) return category[sample] def uniform_sample(size, low, high, device="cpu"): return torch.empty(size, device=device).uniform_(low, high) def normal_sample(size, mean=0, std=1, device="cpu"): return torch.empty(size, device=device).normal_(mean, std) def bernoulli_sample(size, p, device="cpu"): return torch.empty(size, device=device).bernoulli_(p) def random_mat_apply(p, transform, prev, eye, device="cpu"): size = transform.shape[0] select = bernoulli_sample(size, p, device=device).view(size, 1, 1) select_transform = select * transform + (1 - select) * eye return select_transform @ prev def sample_affine(p, size, height, width, device="cpu"): G = torch.eye(3, device=device).unsqueeze(0).repeat(size, 1, 1) eye = G # flip param = category_sample(size, (0, 1)) Gc = scale_mat(1 - 2.0 * param, torch.ones(size), device=device) G = random_mat_apply(p, Gc, G, eye, device=device) # print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\n') # 90 rotate #param = category_sample(size, (0, 3)) #Gc = rotate_mat(-math.pi / 2 * param, device=device) #G = random_mat_apply(p, Gc, G, eye, device=device) # print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\n') # integer translate param = uniform_sample(size, -0.125, 0.125) param_height = torch.round(param * height) / height param_width = torch.round(param * width) / width Gc = translate_mat(param_width, param_height, device=device) G = random_mat_apply(p, Gc, G, eye, device=device) # print('integer translate', G, translate_mat(param_width, param_height), sep='\n') # isotropic scale param = lognormal_sample(size, std=0.2 * math.log(2)) Gc = scale_mat(param, param, device=device) G = random_mat_apply(p, Gc, G, eye, device=device) # print('isotropic scale', G, scale_mat(param, param), sep='\n') p_rot = 1 - math.sqrt(1 - p) # pre-rotate param = uniform_sample(size, -math.pi, math.pi) Gc = rotate_mat(-param, device=device) G = random_mat_apply(p_rot, Gc, G, eye, device=device) # print('pre-rotate', G, rotate_mat(-param), sep='\n') # anisotropic scale param = lognormal_sample(size, std=0.2 * math.log(2)) Gc = scale_mat(param, 1 / param, device=device) G = random_mat_apply(p, Gc, G, eye, device=device) # print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\n') # post-rotate param = uniform_sample(size, -math.pi, math.pi) Gc = rotate_mat(-param, device=device) G = random_mat_apply(p_rot, Gc, G, eye, device=device) # print('post-rotate', G, rotate_mat(-param), sep='\n') # fractional translate param = normal_sample(size, std=0.125) Gc = translate_mat(param, param, device=device) G = random_mat_apply(p, Gc, G, eye, device=device) # print('fractional translate', G, translate_mat(param, param), sep='\n') return G def sample_color(p, size): C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1) eye = C axis_val = 1 / math.sqrt(3) axis = (axis_val, axis_val, axis_val) # brightness param = normal_sample(size, std=0.2) Cc = translate3d_mat(param, param, param) C = random_mat_apply(p, Cc, C, eye) # contrast param = lognormal_sample(size, std=0.5 * math.log(2)) Cc = scale3d_mat(param, param, param) C = random_mat_apply(p, Cc, C, eye) # luma flip param = category_sample(size, (0, 1)) Cc = luma_flip_mat(axis, param) C = random_mat_apply(p, Cc, C, eye) # hue rotation param = uniform_sample(size, -math.pi, math.pi) Cc = rotate3d_mat(axis, param) C = random_mat_apply(p, Cc, C, eye) # saturation param = lognormal_sample(size, std=1 * math.log(2)) Cc = saturation_mat(axis, param) C = random_mat_apply(p, Cc, C, eye) return C def make_grid(shape, x0, x1, y0, y1, device): n, c, h, w = shape grid = torch.empty(n, h, w, 3, device=device) grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device) grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1) grid[:, :, :, 2] = 1 return grid def affine_grid(grid, mat): n, h, w, _ = grid.shape return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2) def get_padding(G, height, width, kernel_size): device = G.device cx = (width - 1) / 2 cy = (height - 1) / 2 cp = torch.tensor( [(-cx, -cy, 1), (cx, -cy, 1), (cx, cy, 1), (-cx, cy, 1)], device=device ) cp = G @ cp.T pad_k = kernel_size // 4 pad = cp[:, :2, :].permute(1, 0, 2).flatten(1) pad = torch.cat((-pad, pad)).max(1).values pad = pad + torch.tensor([pad_k * 2 - cx, pad_k * 2 - cy] * 2, device=device) pad = pad.max(torch.tensor([0, 0] * 2, device=device)) pad = pad.min(torch.tensor([width - 1, height - 1] * 2, device=device)) pad_x1, pad_y1, pad_x2, pad_y2 = pad.ceil().to(torch.int32) return pad_x1, pad_x2, pad_y1, pad_y2 def try_sample_affine_and_pad(img, p, kernel_size, G=None): batch, _, height, width = img.shape G_try = G if G is None: G_try = torch.inverse(sample_affine(p, batch, height, width)) pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(G_try, height, width, kernel_size) img_pad = F.pad(img, (pad_x1, pad_x2, pad_y1, pad_y2), mode="reflect") return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2) class GridSampleForward(autograd.Function): @staticmethod def forward(ctx, input, grid): out = F.grid_sample( input, grid, mode="bilinear", padding_mode="zeros", align_corners=False ) ctx.save_for_backward(input, grid) return out @staticmethod def backward(ctx, grad_output): input, grid = ctx.saved_tensors grad_input, grad_grid = GridSampleBackward.apply(grad_output, input, grid) return grad_input, grad_grid class GridSampleBackward(autograd.Function): @staticmethod def forward(ctx, grad_output, input, grid): op = torch._C._jit_get_operation("aten::grid_sampler_2d_backward") grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False) ctx.save_for_backward(grid) return grad_input, grad_grid @staticmethod def backward(ctx, grad_grad_input, grad_grad_grid): grid, = ctx.saved_tensors grad_grad_output = None if ctx.needs_input_grad[0]: grad_grad_output = GridSampleForward.apply(grad_grad_input, grid) return grad_grad_output, None, None grid_sample = GridSampleForward.apply def scale_mat_single(s_x, s_y): return torch.tensor(((s_x, 0, 0), (0, s_y, 0), (0, 0, 1)), dtype=torch.float32) def translate_mat_single(t_x, t_y): return torch.tensor(((1, 0, t_x), (0, 1, t_y), (0, 0, 1)), dtype=torch.float32) def random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6): kernel = antialiasing_kernel len_k = len(kernel) kernel = torch.as_tensor(kernel).to(img) # kernel = torch.ger(kernel, kernel).to(img) kernel_flip = torch.flip(kernel, (0,)) img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad( img, p, len_k, G ) G_inv = ( translate_mat_single((pad_x1 - pad_x2).item() / 2, (pad_y1 - pad_y2).item() / 2) @ G ) up_pad = ( (len_k + 2 - 1) // 2, (len_k - 2) // 2, (len_k + 2 - 1) // 2, (len_k - 2) // 2, ) img_2x = upfirdn2d(img_pad, kernel.unsqueeze(0), up=(2, 1), pad=(*up_pad[:2], 0, 0)) img_2x = upfirdn2d(img_2x, kernel.unsqueeze(1), up=(1, 2), pad=(0, 0, *up_pad[2:])) G_inv = scale_mat_single(2, 2) @ G_inv @ scale_mat_single(1 / 2, 1 / 2) G_inv = translate_mat_single(-0.5, -0.5) @ G_inv @ translate_mat_single(0.5, 0.5) batch_size, channel, height, width = img.shape pad_k = len_k // 4 shape = (batch_size, channel, (height + pad_k * 2) * 2, (width + pad_k * 2) * 2) G_inv = ( scale_mat_single(2 / img_2x.shape[3], 2 / img_2x.shape[2]) @ G_inv @ scale_mat_single(1 / (2 / shape[3]), 1 / (2 / shape[2])) ) grid = F.affine_grid(G_inv[:, :2, :].to(img_2x), shape, align_corners=False) img_affine = grid_sample(img_2x, grid) d_p = -pad_k * 2 down_pad = ( d_p + (len_k - 2 + 1) // 2, d_p + (len_k - 2) // 2, d_p + (len_k - 2 + 1) // 2, d_p + (len_k - 2) // 2, ) img_down = upfirdn2d( img_affine, kernel_flip.unsqueeze(0), down=(2, 1), pad=(*down_pad[:2], 0, 0) ) img_down = upfirdn2d( img_down, kernel_flip.unsqueeze(1), down=(1, 2), pad=(0, 0, *down_pad[2:]) ) return img_down, G def apply_color(img, mat): batch = img.shape[0] img = img.permute(0, 2, 3, 1) mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3) mat_add = mat[:, :3, 3].view(batch, 1, 1, 3) img = img @ mat_mul + mat_add img = img.permute(0, 3, 1, 2) return img def random_apply_color(img, p, C=None): if C is None: C = sample_color(p, img.shape[0]) img = apply_color(img, C.to(img)) return img, C def augment(img, p, transform_matrix=(None, None)): img, G = random_apply_affine(img, p, transform_matrix[0]) if img.shape[1] == 3: img, C = random_apply_color(img, p, transform_matrix[1]) else: tmp, C = random_apply_color(img[:,0:3], p, transform_matrix[1]) img = torch.cat((tmp, img[:,3:]), dim=1) return img, (G, C)