Arnaudding001 commited on
Commit
355bf1a
1 Parent(s): 15d7fed

Create simple_augment.py

Browse files
Files changed (1) hide show
  1. simple_augment.py +468 -0
simple_augment.py ADDED
@@ -0,0 +1,468 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # almost the same as model.stylegan.non_leaking
2
+ # we only modify the parameters in sample_affine() to make the transformations mild
3
+
4
+ import math
5
+
6
+ import torch
7
+ from torch import autograd
8
+ from torch.nn import functional as F
9
+ import numpy as np
10
+
11
+ from model.stylegan.distributed import reduce_sum
12
+ from model.stylegan.op import upfirdn2d
13
+
14
+
15
+ class AdaptiveAugment:
16
+ def __init__(self, ada_aug_target, ada_aug_len, update_every, device):
17
+ self.ada_aug_target = ada_aug_target
18
+ self.ada_aug_len = ada_aug_len
19
+ self.update_every = update_every
20
+
21
+ self.ada_update = 0
22
+ self.ada_aug_buf = torch.tensor([0.0, 0.0], device=device)
23
+ self.r_t_stat = 0
24
+ self.ada_aug_p = 0
25
+
26
+ @torch.no_grad()
27
+ def tune(self, real_pred):
28
+ self.ada_aug_buf += torch.tensor(
29
+ (torch.sign(real_pred).sum().item(), real_pred.shape[0]),
30
+ device=real_pred.device,
31
+ )
32
+ self.ada_update += 1
33
+
34
+ if self.ada_update % self.update_every == 0:
35
+ self.ada_aug_buf = reduce_sum(self.ada_aug_buf)
36
+ pred_signs, n_pred = self.ada_aug_buf.tolist()
37
+
38
+ self.r_t_stat = pred_signs / n_pred
39
+
40
+ if self.r_t_stat > self.ada_aug_target:
41
+ sign = 1
42
+
43
+ else:
44
+ sign = -1
45
+
46
+ self.ada_aug_p += sign * n_pred / self.ada_aug_len
47
+ self.ada_aug_p = min(1, max(0, self.ada_aug_p))
48
+ self.ada_aug_buf.mul_(0)
49
+ self.ada_update = 0
50
+
51
+ return self.ada_aug_p
52
+
53
+
54
+ SYM6 = (
55
+ 0.015404109327027373,
56
+ 0.0034907120842174702,
57
+ -0.11799011114819057,
58
+ -0.048311742585633,
59
+ 0.4910559419267466,
60
+ 0.787641141030194,
61
+ 0.3379294217276218,
62
+ -0.07263752278646252,
63
+ -0.021060292512300564,
64
+ 0.04472490177066578,
65
+ 0.0017677118642428036,
66
+ -0.007800708325034148,
67
+ )
68
+
69
+
70
+ def translate_mat(t_x, t_y, device="cpu"):
71
+ batch = t_x.shape[0]
72
+
73
+ mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
74
+ translate = torch.stack((t_x, t_y), 1)
75
+ mat[:, :2, 2] = translate
76
+
77
+ return mat
78
+
79
+
80
+ def rotate_mat(theta, device="cpu"):
81
+ batch = theta.shape[0]
82
+
83
+ mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
84
+ sin_t = torch.sin(theta)
85
+ cos_t = torch.cos(theta)
86
+ rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2)
87
+ mat[:, :2, :2] = rot
88
+
89
+ return mat
90
+
91
+
92
+ def scale_mat(s_x, s_y, device="cpu"):
93
+ batch = s_x.shape[0]
94
+
95
+ mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
96
+ mat[:, 0, 0] = s_x
97
+ mat[:, 1, 1] = s_y
98
+
99
+ return mat
100
+
101
+
102
+ def translate3d_mat(t_x, t_y, t_z):
103
+ batch = t_x.shape[0]
104
+
105
+ mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
106
+ translate = torch.stack((t_x, t_y, t_z), 1)
107
+ mat[:, :3, 3] = translate
108
+
109
+ return mat
110
+
111
+
112
+ def rotate3d_mat(axis, theta):
113
+ batch = theta.shape[0]
114
+
115
+ u_x, u_y, u_z = axis
116
+
117
+ eye = torch.eye(3).unsqueeze(0)
118
+ cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0)
119
+ outer = torch.tensor(axis)
120
+ outer = (outer.unsqueeze(1) * outer).unsqueeze(0)
121
+
122
+ sin_t = torch.sin(theta).view(-1, 1, 1)
123
+ cos_t = torch.cos(theta).view(-1, 1, 1)
124
+
125
+ rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer
126
+
127
+ eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
128
+ eye_4[:, :3, :3] = rot
129
+
130
+ return eye_4
131
+
132
+
133
+ def scale3d_mat(s_x, s_y, s_z):
134
+ batch = s_x.shape[0]
135
+
136
+ mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
137
+ mat[:, 0, 0] = s_x
138
+ mat[:, 1, 1] = s_y
139
+ mat[:, 2, 2] = s_z
140
+
141
+ return mat
142
+
143
+
144
+ def luma_flip_mat(axis, i):
145
+ batch = i.shape[0]
146
+
147
+ eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
148
+ axis = torch.tensor(axis + (0,))
149
+ flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1)
150
+
151
+ return eye - flip
152
+
153
+
154
+ def saturation_mat(axis, i):
155
+ batch = i.shape[0]
156
+
157
+ eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
158
+ axis = torch.tensor(axis + (0,))
159
+ axis = torch.ger(axis, axis)
160
+ saturate = axis + (eye - axis) * i.view(-1, 1, 1)
161
+
162
+ return saturate
163
+
164
+
165
+ def lognormal_sample(size, mean=0, std=1, device="cpu"):
166
+ return torch.empty(size, device=device).log_normal_(mean=mean, std=std)
167
+
168
+
169
+ def category_sample(size, categories, device="cpu"):
170
+ category = torch.tensor(categories, device=device)
171
+ sample = torch.randint(high=len(categories), size=(size,), device=device)
172
+
173
+ return category[sample]
174
+
175
+
176
+ def uniform_sample(size, low, high, device="cpu"):
177
+ return torch.empty(size, device=device).uniform_(low, high)
178
+
179
+
180
+ def normal_sample(size, mean=0, std=1, device="cpu"):
181
+ return torch.empty(size, device=device).normal_(mean, std)
182
+
183
+
184
+ def bernoulli_sample(size, p, device="cpu"):
185
+ return torch.empty(size, device=device).bernoulli_(p)
186
+
187
+
188
+ def random_mat_apply(p, transform, prev, eye, device="cpu"):
189
+ size = transform.shape[0]
190
+ select = bernoulli_sample(size, p, device=device).view(size, 1, 1)
191
+ select_transform = select * transform + (1 - select) * eye
192
+
193
+ return select_transform @ prev
194
+
195
+
196
+ def sample_affine(p, size, height, width, device="cpu"):
197
+ G = torch.eye(3, device=device).unsqueeze(0).repeat(size, 1, 1)
198
+ eye = G
199
+
200
+ # flip
201
+ param = category_sample(size, (0, 1))
202
+ Gc = scale_mat(1 - 2.0 * param, torch.ones(size), device=device)
203
+ G = random_mat_apply(p, Gc, G, eye, device=device)
204
+ # print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\n')
205
+
206
+ # 90 rotate
207
+ #param = category_sample(size, (0, 3))
208
+ #Gc = rotate_mat(-math.pi / 2 * param, device=device)
209
+ #G = random_mat_apply(p, Gc, G, eye, device=device)
210
+ # print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\n')
211
+
212
+ # integer translate
213
+ param = uniform_sample(size, -0.125, 0.125)
214
+ param_height = torch.round(param * height) / height
215
+ param_width = torch.round(param * width) / width
216
+ Gc = translate_mat(param_width, param_height, device=device)
217
+ G = random_mat_apply(p, Gc, G, eye, device=device)
218
+ # print('integer translate', G, translate_mat(param_width, param_height), sep='\n')
219
+
220
+ # isotropic scale
221
+ param = lognormal_sample(size, std=0.1 * math.log(2))
222
+ Gc = scale_mat(param, param, device=device)
223
+ G = random_mat_apply(p, Gc, G, eye, device=device)
224
+ # print('isotropic scale', G, scale_mat(param, param), sep='\n')
225
+
226
+ p_rot = 1 - math.sqrt(1 - p)
227
+
228
+ # pre-rotate
229
+ param = uniform_sample(size, -math.pi * 0.25, math.pi * 0.25)
230
+ Gc = rotate_mat(-param, device=device)
231
+ G = random_mat_apply(p_rot, Gc, G, eye, device=device)
232
+ # print('pre-rotate', G, rotate_mat(-param), sep='\n')
233
+
234
+ # anisotropic scale
235
+ param = lognormal_sample(size, std=0.1 * math.log(2))
236
+ Gc = scale_mat(param, 1 / param, device=device)
237
+ G = random_mat_apply(p, Gc, G, eye, device=device)
238
+ # print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\n')
239
+
240
+ # post-rotate
241
+ param = uniform_sample(size, -math.pi * 0.25, math.pi * 0.25)
242
+ Gc = rotate_mat(-param, device=device)
243
+ G = random_mat_apply(p_rot, Gc, G, eye, device=device)
244
+ # print('post-rotate', G, rotate_mat(-param), sep='\n')
245
+
246
+ # fractional translate
247
+ param = normal_sample(size, std=0.125)
248
+ Gc = translate_mat(param, param, device=device)
249
+ G = random_mat_apply(p, Gc, G, eye, device=device)
250
+ # print('fractional translate', G, translate_mat(param, param), sep='\n')
251
+
252
+ return G
253
+
254
+
255
+ def sample_color(p, size):
256
+ C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1)
257
+ eye = C
258
+ axis_val = 1 / math.sqrt(3)
259
+ axis = (axis_val, axis_val, axis_val)
260
+
261
+ # brightness
262
+ param = normal_sample(size, std=0.2)
263
+ Cc = translate3d_mat(param, param, param)
264
+ C = random_mat_apply(p, Cc, C, eye)
265
+
266
+ # contrast
267
+ param = lognormal_sample(size, std=0.5 * math.log(2))
268
+ Cc = scale3d_mat(param, param, param)
269
+ C = random_mat_apply(p, Cc, C, eye)
270
+
271
+ # luma flip
272
+ param = category_sample(size, (0, 1))
273
+ Cc = luma_flip_mat(axis, param)
274
+ C = random_mat_apply(p, Cc, C, eye)
275
+
276
+ # hue rotation
277
+ param = uniform_sample(size, -math.pi, math.pi)
278
+ Cc = rotate3d_mat(axis, param)
279
+ C = random_mat_apply(p, Cc, C, eye)
280
+
281
+ # saturation
282
+ param = lognormal_sample(size, std=1 * math.log(2))
283
+ Cc = saturation_mat(axis, param)
284
+ C = random_mat_apply(p, Cc, C, eye)
285
+
286
+ return C
287
+
288
+
289
+ def make_grid(shape, x0, x1, y0, y1, device):
290
+ n, c, h, w = shape
291
+ grid = torch.empty(n, h, w, 3, device=device)
292
+ grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device)
293
+ grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1)
294
+ grid[:, :, :, 2] = 1
295
+
296
+ return grid
297
+
298
+
299
+ def affine_grid(grid, mat):
300
+ n, h, w, _ = grid.shape
301
+ return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2)
302
+
303
+
304
+ def get_padding(G, height, width, kernel_size):
305
+ device = G.device
306
+
307
+ cx = (width - 1) / 2
308
+ cy = (height - 1) / 2
309
+ cp = torch.tensor(
310
+ [(-cx, -cy, 1), (cx, -cy, 1), (cx, cy, 1), (-cx, cy, 1)], device=device
311
+ )
312
+ cp = G @ cp.T
313
+
314
+ pad_k = kernel_size // 4
315
+
316
+ pad = cp[:, :2, :].permute(1, 0, 2).flatten(1)
317
+ pad = torch.cat((-pad, pad)).max(1).values
318
+ pad = pad + torch.tensor([pad_k * 2 - cx, pad_k * 2 - cy] * 2, device=device)
319
+ pad = pad.max(torch.tensor([0, 0] * 2, device=device))
320
+ pad = pad.min(torch.tensor([width - 1, height - 1] * 2, device=device))
321
+
322
+ pad_x1, pad_y1, pad_x2, pad_y2 = pad.ceil().to(torch.int32)
323
+
324
+ return pad_x1, pad_x2, pad_y1, pad_y2
325
+
326
+
327
+ def try_sample_affine_and_pad(img, p, kernel_size, G=None):
328
+ batch, _, height, width = img.shape
329
+
330
+ G_try = G
331
+
332
+ if G is None:
333
+ G_try = torch.inverse(sample_affine(p, batch, height, width))
334
+
335
+ pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(G_try, height, width, kernel_size)
336
+
337
+ img_pad = F.pad(img, (pad_x1, pad_x2, pad_y1, pad_y2), mode="reflect")
338
+
339
+ return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2)
340
+
341
+
342
+ class GridSampleForward(autograd.Function):
343
+ @staticmethod
344
+ def forward(ctx, input, grid):
345
+ out = F.grid_sample(
346
+ input, grid, mode="bilinear", padding_mode="zeros", align_corners=False
347
+ )
348
+ ctx.save_for_backward(input, grid)
349
+
350
+ return out
351
+
352
+ @staticmethod
353
+ def backward(ctx, grad_output):
354
+ input, grid = ctx.saved_tensors
355
+ grad_input, grad_grid = GridSampleBackward.apply(grad_output, input, grid)
356
+
357
+ return grad_input, grad_grid
358
+
359
+
360
+ class GridSampleBackward(autograd.Function):
361
+ @staticmethod
362
+ def forward(ctx, grad_output, input, grid):
363
+ op = torch._C._jit_get_operation("aten::grid_sampler_2d_backward")
364
+ grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)
365
+ ctx.save_for_backward(grid)
366
+
367
+ return grad_input, grad_grid
368
+
369
+ @staticmethod
370
+ def backward(ctx, grad_grad_input, grad_grad_grid):
371
+ grid, = ctx.saved_tensors
372
+ grad_grad_output = None
373
+
374
+ if ctx.needs_input_grad[0]:
375
+ grad_grad_output = GridSampleForward.apply(grad_grad_input, grid)
376
+
377
+ return grad_grad_output, None, None
378
+
379
+
380
+ grid_sample = GridSampleForward.apply
381
+
382
+
383
+ def scale_mat_single(s_x, s_y):
384
+ return torch.tensor(((s_x, 0, 0), (0, s_y, 0), (0, 0, 1)), dtype=torch.float32)
385
+
386
+
387
+ def translate_mat_single(t_x, t_y):
388
+ return torch.tensor(((1, 0, t_x), (0, 1, t_y), (0, 0, 1)), dtype=torch.float32)
389
+
390
+
391
+ def random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6):
392
+ kernel = antialiasing_kernel
393
+ len_k = len(kernel)
394
+
395
+ kernel = torch.as_tensor(kernel).to(img)
396
+ # kernel = torch.ger(kernel, kernel).to(img)
397
+ kernel_flip = torch.flip(kernel, (0,))
398
+
399
+ img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad(
400
+ img, p, len_k, G
401
+ )
402
+
403
+ G_inv = (
404
+ translate_mat_single((pad_x1 - pad_x2).item() / 2, (pad_y1 - pad_y2).item() / 2)
405
+ @ G
406
+ )
407
+ up_pad = (
408
+ (len_k + 2 - 1) // 2,
409
+ (len_k - 2) // 2,
410
+ (len_k + 2 - 1) // 2,
411
+ (len_k - 2) // 2,
412
+ )
413
+ img_2x = upfirdn2d(img_pad, kernel.unsqueeze(0), up=(2, 1), pad=(*up_pad[:2], 0, 0))
414
+ img_2x = upfirdn2d(img_2x, kernel.unsqueeze(1), up=(1, 2), pad=(0, 0, *up_pad[2:]))
415
+ G_inv = scale_mat_single(2, 2) @ G_inv @ scale_mat_single(1 / 2, 1 / 2)
416
+ G_inv = translate_mat_single(-0.5, -0.5) @ G_inv @ translate_mat_single(0.5, 0.5)
417
+ batch_size, channel, height, width = img.shape
418
+ pad_k = len_k // 4
419
+ shape = (batch_size, channel, (height + pad_k * 2) * 2, (width + pad_k * 2) * 2)
420
+ G_inv = (
421
+ scale_mat_single(2 / img_2x.shape[3], 2 / img_2x.shape[2])
422
+ @ G_inv
423
+ @ scale_mat_single(1 / (2 / shape[3]), 1 / (2 / shape[2]))
424
+ )
425
+ grid = F.affine_grid(G_inv[:, :2, :].to(img_2x), shape, align_corners=False)
426
+ img_affine = grid_sample(img_2x, grid)
427
+ d_p = -pad_k * 2
428
+ down_pad = (
429
+ d_p + (len_k - 2 + 1) // 2,
430
+ d_p + (len_k - 2) // 2,
431
+ d_p + (len_k - 2 + 1) // 2,
432
+ d_p + (len_k - 2) // 2,
433
+ )
434
+ img_down = upfirdn2d(
435
+ img_affine, kernel_flip.unsqueeze(0), down=(2, 1), pad=(*down_pad[:2], 0, 0)
436
+ )
437
+ img_down = upfirdn2d(
438
+ img_down, kernel_flip.unsqueeze(1), down=(1, 2), pad=(0, 0, *down_pad[2:])
439
+ )
440
+
441
+ return img_down, G
442
+
443
+
444
+ def apply_color(img, mat):
445
+ batch = img.shape[0]
446
+ img = img.permute(0, 2, 3, 1)
447
+ mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3)
448
+ mat_add = mat[:, :3, 3].view(batch, 1, 1, 3)
449
+ img = img @ mat_mul + mat_add
450
+ img = img.permute(0, 3, 1, 2)
451
+
452
+ return img
453
+
454
+
455
+ def random_apply_color(img, p, C=None):
456
+ if C is None:
457
+ C = sample_color(p, img.shape[0])
458
+
459
+ img = apply_color(img, C.to(img))
460
+
461
+ return img, C
462
+
463
+
464
+ def augment(img, p, transform_matrix=(None, None)):
465
+ img, G = random_apply_affine(img, p, transform_matrix[0])
466
+ img, C = random_apply_color(img, p, transform_matrix[1])
467
+
468
+ return img, (G, C)