Spaces:
Runtime error
Runtime error
Arnaudding001
commited on
Commit
•
355bf1a
1
Parent(s):
15d7fed
Create simple_augment.py
Browse files- 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)
|