Spaces:
Running
Running
Upload transform/randaugment.py
Browse files- transform/randaugment.py +340 -0
transform/randaugment.py
ADDED
@@ -0,0 +1,340 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import numpy as np
|
3 |
+
|
4 |
+
|
5 |
+
## aug functions
|
6 |
+
def identity_func(img):
|
7 |
+
return img
|
8 |
+
|
9 |
+
|
10 |
+
def autocontrast_func(img, cutoff=0):
|
11 |
+
'''
|
12 |
+
same output as PIL.ImageOps.autocontrast
|
13 |
+
'''
|
14 |
+
n_bins = 256
|
15 |
+
|
16 |
+
def tune_channel(ch):
|
17 |
+
n = ch.size
|
18 |
+
cut = cutoff * n // 100
|
19 |
+
if cut == 0:
|
20 |
+
high, low = ch.max(), ch.min()
|
21 |
+
else:
|
22 |
+
hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
|
23 |
+
low = np.argwhere(np.cumsum(hist) > cut)
|
24 |
+
low = 0 if low.shape[0] == 0 else low[0]
|
25 |
+
high = np.argwhere(np.cumsum(hist[::-1]) > cut)
|
26 |
+
high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
|
27 |
+
if high <= low:
|
28 |
+
table = np.arange(n_bins)
|
29 |
+
else:
|
30 |
+
scale = (n_bins - 1) / (high - low)
|
31 |
+
offset = -low * scale
|
32 |
+
table = np.arange(n_bins) * scale + offset
|
33 |
+
table[table < 0] = 0
|
34 |
+
table[table > n_bins - 1] = n_bins - 1
|
35 |
+
table = table.clip(0, 255).astype(np.uint8)
|
36 |
+
return table[ch]
|
37 |
+
|
38 |
+
channels = [tune_channel(ch) for ch in cv2.split(img)]
|
39 |
+
out = cv2.merge(channels)
|
40 |
+
return out
|
41 |
+
|
42 |
+
|
43 |
+
def equalize_func(img):
|
44 |
+
'''
|
45 |
+
same output as PIL.ImageOps.equalize
|
46 |
+
PIL's implementation is different from cv2.equalize
|
47 |
+
'''
|
48 |
+
n_bins = 256
|
49 |
+
|
50 |
+
def tune_channel(ch):
|
51 |
+
hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
|
52 |
+
non_zero_hist = hist[hist != 0].reshape(-1)
|
53 |
+
step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
|
54 |
+
if step == 0: return ch
|
55 |
+
n = np.empty_like(hist)
|
56 |
+
n[0] = step // 2
|
57 |
+
n[1:] = hist[:-1]
|
58 |
+
table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
|
59 |
+
return table[ch]
|
60 |
+
|
61 |
+
channels = [tune_channel(ch) for ch in cv2.split(img)]
|
62 |
+
out = cv2.merge(channels)
|
63 |
+
return out
|
64 |
+
|
65 |
+
|
66 |
+
def rotate_func(img, degree, fill=(0, 0, 0)):
|
67 |
+
'''
|
68 |
+
like PIL, rotate by degree, not radians
|
69 |
+
'''
|
70 |
+
H, W = img.shape[0], img.shape[1]
|
71 |
+
center = W / 2, H / 2
|
72 |
+
M = cv2.getRotationMatrix2D(center, degree, 1)
|
73 |
+
out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
|
74 |
+
return out
|
75 |
+
|
76 |
+
|
77 |
+
def solarize_func(img, thresh=128):
|
78 |
+
'''
|
79 |
+
same output as PIL.ImageOps.posterize
|
80 |
+
'''
|
81 |
+
table = np.array([el if el < thresh else 255 - el for el in range(256)])
|
82 |
+
table = table.clip(0, 255).astype(np.uint8)
|
83 |
+
out = table[img]
|
84 |
+
return out
|
85 |
+
|
86 |
+
|
87 |
+
def color_func(img, factor):
|
88 |
+
'''
|
89 |
+
same output as PIL.ImageEnhance.Color
|
90 |
+
'''
|
91 |
+
## implementation according to PIL definition, quite slow
|
92 |
+
# degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
|
93 |
+
# out = blend(degenerate, img, factor)
|
94 |
+
# M = (
|
95 |
+
# np.eye(3) * factor
|
96 |
+
# + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
|
97 |
+
# )[np.newaxis, np.newaxis, :]
|
98 |
+
M = (
|
99 |
+
np.float32([
|
100 |
+
[0.886, -0.114, -0.114],
|
101 |
+
[-0.587, 0.413, -0.587],
|
102 |
+
[-0.299, -0.299, 0.701]]) * factor
|
103 |
+
+ np.float32([[0.114], [0.587], [0.299]])
|
104 |
+
)
|
105 |
+
out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
|
106 |
+
return out
|
107 |
+
|
108 |
+
|
109 |
+
def contrast_func(img, factor):
|
110 |
+
"""
|
111 |
+
same output as PIL.ImageEnhance.Contrast
|
112 |
+
"""
|
113 |
+
mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
|
114 |
+
table = np.array([(
|
115 |
+
el - mean) * factor + mean
|
116 |
+
for el in range(256)
|
117 |
+
]).clip(0, 255).astype(np.uint8)
|
118 |
+
out = table[img]
|
119 |
+
return out
|
120 |
+
|
121 |
+
|
122 |
+
def brightness_func(img, factor):
|
123 |
+
'''
|
124 |
+
same output as PIL.ImageEnhance.Contrast
|
125 |
+
'''
|
126 |
+
table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)
|
127 |
+
out = table[img]
|
128 |
+
return out
|
129 |
+
|
130 |
+
|
131 |
+
def sharpness_func(img, factor):
|
132 |
+
'''
|
133 |
+
The differences the this result and PIL are all on the 4 boundaries, the center
|
134 |
+
areas are same
|
135 |
+
'''
|
136 |
+
kernel = np.ones((3, 3), dtype=np.float32)
|
137 |
+
kernel[1][1] = 5
|
138 |
+
kernel /= 13
|
139 |
+
degenerate = cv2.filter2D(img, -1, kernel)
|
140 |
+
if factor == 0.0:
|
141 |
+
out = degenerate
|
142 |
+
elif factor == 1.0:
|
143 |
+
out = img
|
144 |
+
else:
|
145 |
+
out = img.astype(np.float32)
|
146 |
+
degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
|
147 |
+
out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)
|
148 |
+
out = out.astype(np.uint8)
|
149 |
+
return out
|
150 |
+
|
151 |
+
|
152 |
+
def shear_x_func(img, factor, fill=(0, 0, 0)):
|
153 |
+
H, W = img.shape[0], img.shape[1]
|
154 |
+
M = np.float32([[1, factor, 0], [0, 1, 0]])
|
155 |
+
out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
|
156 |
+
return out
|
157 |
+
|
158 |
+
|
159 |
+
def translate_x_func(img, offset, fill=(0, 0, 0)):
|
160 |
+
'''
|
161 |
+
same output as PIL.Image.transform
|
162 |
+
'''
|
163 |
+
H, W = img.shape[0], img.shape[1]
|
164 |
+
M = np.float32([[1, 0, -offset], [0, 1, 0]])
|
165 |
+
out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
|
166 |
+
return out
|
167 |
+
|
168 |
+
|
169 |
+
def translate_y_func(img, offset, fill=(0, 0, 0)):
|
170 |
+
'''
|
171 |
+
same output as PIL.Image.transform
|
172 |
+
'''
|
173 |
+
H, W = img.shape[0], img.shape[1]
|
174 |
+
M = np.float32([[1, 0, 0], [0, 1, -offset]])
|
175 |
+
out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
|
176 |
+
return out
|
177 |
+
|
178 |
+
|
179 |
+
def posterize_func(img, bits):
|
180 |
+
'''
|
181 |
+
same output as PIL.ImageOps.posterize
|
182 |
+
'''
|
183 |
+
out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
|
184 |
+
return out
|
185 |
+
|
186 |
+
|
187 |
+
def shear_y_func(img, factor, fill=(0, 0, 0)):
|
188 |
+
H, W = img.shape[0], img.shape[1]
|
189 |
+
M = np.float32([[1, 0, 0], [factor, 1, 0]])
|
190 |
+
out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
|
191 |
+
return out
|
192 |
+
|
193 |
+
|
194 |
+
def cutout_func(img, pad_size, replace=(0, 0, 0)):
|
195 |
+
replace = np.array(replace, dtype=np.uint8)
|
196 |
+
H, W = img.shape[0], img.shape[1]
|
197 |
+
rh, rw = np.random.random(2)
|
198 |
+
pad_size = pad_size // 2
|
199 |
+
ch, cw = int(rh * H), int(rw * W)
|
200 |
+
x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
|
201 |
+
y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
|
202 |
+
out = img.copy()
|
203 |
+
out[x1:x2, y1:y2, :] = replace
|
204 |
+
return out
|
205 |
+
|
206 |
+
|
207 |
+
### level to args
|
208 |
+
def enhance_level_to_args(MAX_LEVEL):
|
209 |
+
def level_to_args(level):
|
210 |
+
return ((level / MAX_LEVEL) * 1.8 + 0.1,)
|
211 |
+
return level_to_args
|
212 |
+
|
213 |
+
|
214 |
+
def shear_level_to_args(MAX_LEVEL, replace_value):
|
215 |
+
def level_to_args(level):
|
216 |
+
level = (level / MAX_LEVEL) * 0.3
|
217 |
+
if np.random.random() > 0.5: level = -level
|
218 |
+
return (level, replace_value)
|
219 |
+
|
220 |
+
return level_to_args
|
221 |
+
|
222 |
+
|
223 |
+
def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
|
224 |
+
def level_to_args(level):
|
225 |
+
level = (level / MAX_LEVEL) * float(translate_const)
|
226 |
+
if np.random.random() > 0.5: level = -level
|
227 |
+
return (level, replace_value)
|
228 |
+
|
229 |
+
return level_to_args
|
230 |
+
|
231 |
+
|
232 |
+
def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
|
233 |
+
def level_to_args(level):
|
234 |
+
level = int((level / MAX_LEVEL) * cutout_const)
|
235 |
+
return (level, replace_value)
|
236 |
+
|
237 |
+
return level_to_args
|
238 |
+
|
239 |
+
|
240 |
+
def solarize_level_to_args(MAX_LEVEL):
|
241 |
+
def level_to_args(level):
|
242 |
+
level = int((level / MAX_LEVEL) * 256)
|
243 |
+
return (level, )
|
244 |
+
return level_to_args
|
245 |
+
|
246 |
+
|
247 |
+
def none_level_to_args(level):
|
248 |
+
return ()
|
249 |
+
|
250 |
+
|
251 |
+
def posterize_level_to_args(MAX_LEVEL):
|
252 |
+
def level_to_args(level):
|
253 |
+
level = int((level / MAX_LEVEL) * 4)
|
254 |
+
return (level, )
|
255 |
+
return level_to_args
|
256 |
+
|
257 |
+
|
258 |
+
def rotate_level_to_args(MAX_LEVEL, replace_value):
|
259 |
+
def level_to_args(level):
|
260 |
+
level = (level / MAX_LEVEL) * 30
|
261 |
+
if np.random.random() < 0.5:
|
262 |
+
level = -level
|
263 |
+
return (level, replace_value)
|
264 |
+
|
265 |
+
return level_to_args
|
266 |
+
|
267 |
+
|
268 |
+
func_dict = {
|
269 |
+
'Identity': identity_func,
|
270 |
+
'AutoContrast': autocontrast_func,
|
271 |
+
'Equalize': equalize_func,
|
272 |
+
'Rotate': rotate_func,
|
273 |
+
'Solarize': solarize_func,
|
274 |
+
'Color': color_func,
|
275 |
+
'Contrast': contrast_func,
|
276 |
+
'Brightness': brightness_func,
|
277 |
+
'Sharpness': sharpness_func,
|
278 |
+
'ShearX': shear_x_func,
|
279 |
+
'TranslateX': translate_x_func,
|
280 |
+
'TranslateY': translate_y_func,
|
281 |
+
'Posterize': posterize_func,
|
282 |
+
'ShearY': shear_y_func,
|
283 |
+
}
|
284 |
+
|
285 |
+
translate_const = 10
|
286 |
+
MAX_LEVEL = 10
|
287 |
+
replace_value = (128, 128, 128)
|
288 |
+
arg_dict = {
|
289 |
+
'Identity': none_level_to_args,
|
290 |
+
'AutoContrast': none_level_to_args,
|
291 |
+
'Equalize': none_level_to_args,
|
292 |
+
'Rotate': rotate_level_to_args(MAX_LEVEL, replace_value),
|
293 |
+
'Solarize': solarize_level_to_args(MAX_LEVEL),
|
294 |
+
'Color': enhance_level_to_args(MAX_LEVEL),
|
295 |
+
'Contrast': enhance_level_to_args(MAX_LEVEL),
|
296 |
+
'Brightness': enhance_level_to_args(MAX_LEVEL),
|
297 |
+
'Sharpness': enhance_level_to_args(MAX_LEVEL),
|
298 |
+
'ShearX': shear_level_to_args(MAX_LEVEL, replace_value),
|
299 |
+
'TranslateX': translate_level_to_args(
|
300 |
+
translate_const, MAX_LEVEL, replace_value
|
301 |
+
),
|
302 |
+
'TranslateY': translate_level_to_args(
|
303 |
+
translate_const, MAX_LEVEL, replace_value
|
304 |
+
),
|
305 |
+
'Posterize': posterize_level_to_args(MAX_LEVEL),
|
306 |
+
'ShearY': shear_level_to_args(MAX_LEVEL, replace_value),
|
307 |
+
}
|
308 |
+
|
309 |
+
|
310 |
+
class RandomAugment(object):
|
311 |
+
|
312 |
+
def __init__(self, N=2, M=10, isPIL=False, augs=[]):
|
313 |
+
self.N = N
|
314 |
+
self.M = M
|
315 |
+
self.isPIL = isPIL
|
316 |
+
if augs:
|
317 |
+
self.augs = augs
|
318 |
+
else:
|
319 |
+
self.augs = list(arg_dict.keys())
|
320 |
+
|
321 |
+
def get_random_ops(self):
|
322 |
+
sampled_ops = np.random.choice(self.augs, self.N)
|
323 |
+
return [(op, 0.5, self.M) for op in sampled_ops]
|
324 |
+
|
325 |
+
def __call__(self, img):
|
326 |
+
if self.isPIL:
|
327 |
+
img = np.array(img)
|
328 |
+
ops = self.get_random_ops()
|
329 |
+
for name, prob, level in ops:
|
330 |
+
if np.random.random() > prob:
|
331 |
+
continue
|
332 |
+
args = arg_dict[name](level)
|
333 |
+
img = func_dict[name](img, *args)
|
334 |
+
return img
|
335 |
+
|
336 |
+
|
337 |
+
if __name__ == '__main__':
|
338 |
+
a = RandomAugment()
|
339 |
+
img = np.random.randn(32, 32, 3)
|
340 |
+
a(img)
|