timm_autoaugment.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # Code was heavily based on https://github.com/rwightman/pytorch-image-models
  15. import random
  16. import math
  17. import re
  18. from PIL import Image, ImageOps, ImageEnhance, ImageChops
  19. import PIL
  20. import numpy as np
  21. IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
  22. _PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]])
  23. _FILL = (128, 128, 128)
  24. # This signifies the max integer that the controller RNN could predict for the
  25. # augmentation scheme.
  26. _MAX_LEVEL = 10.
  27. _HPARAMS_DEFAULT = dict(
  28. translate_const=250,
  29. img_mean=_FILL, )
  30. _RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
  31. def _pil_interp(method):
  32. if method == 'bicubic':
  33. return Image.BICUBIC
  34. elif method == 'lanczos':
  35. return Image.LANCZOS
  36. elif method == 'hamming':
  37. return Image.HAMMING
  38. else:
  39. # default bilinear, do we want to allow nearest?
  40. return Image.BILINEAR
  41. def _interpolation(kwargs):
  42. interpolation = kwargs.pop('resample', Image.BILINEAR)
  43. if isinstance(interpolation, (list, tuple)):
  44. return random.choice(interpolation)
  45. else:
  46. return interpolation
  47. def _check_args_tf(kwargs):
  48. if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
  49. kwargs.pop('fillcolor')
  50. kwargs['resample'] = _interpolation(kwargs)
  51. def shear_x(img, factor, **kwargs):
  52. _check_args_tf(kwargs)
  53. return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0),
  54. **kwargs)
  55. def shear_y(img, factor, **kwargs):
  56. _check_args_tf(kwargs)
  57. return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0),
  58. **kwargs)
  59. def translate_x_rel(img, pct, **kwargs):
  60. pixels = pct * img.size[0]
  61. _check_args_tf(kwargs)
  62. return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0),
  63. **kwargs)
  64. def translate_y_rel(img, pct, **kwargs):
  65. pixels = pct * img.size[1]
  66. _check_args_tf(kwargs)
  67. return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels),
  68. **kwargs)
  69. def translate_x_abs(img, pixels, **kwargs):
  70. _check_args_tf(kwargs)
  71. return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0),
  72. **kwargs)
  73. def translate_y_abs(img, pixels, **kwargs):
  74. _check_args_tf(kwargs)
  75. return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels),
  76. **kwargs)
  77. def rotate(img, degrees, **kwargs):
  78. _check_args_tf(kwargs)
  79. if _PIL_VER >= (5, 2):
  80. return img.rotate(degrees, **kwargs)
  81. elif _PIL_VER >= (5, 0):
  82. w, h = img.size
  83. post_trans = (0, 0)
  84. rotn_center = (w / 2.0, h / 2.0)
  85. angle = -math.radians(degrees)
  86. matrix = [
  87. round(math.cos(angle), 15),
  88. round(math.sin(angle), 15),
  89. 0.0,
  90. round(-math.sin(angle), 15),
  91. round(math.cos(angle), 15),
  92. 0.0,
  93. ]
  94. def transform(x, y, matrix):
  95. (a, b, c, d, e, f) = matrix
  96. return a * x + b * y + c, d * x + e * y + f
  97. matrix[2], matrix[5] = transform(-rotn_center[0] - post_trans[0],
  98. -rotn_center[1] - post_trans[1],
  99. matrix)
  100. matrix[2] += rotn_center[0]
  101. matrix[5] += rotn_center[1]
  102. return img.transform(img.size, Image.AFFINE, matrix, **kwargs)
  103. else:
  104. return img.rotate(degrees, resample=kwargs['resample'])
  105. def auto_contrast(img, **__):
  106. return ImageOps.autocontrast(img)
  107. def invert(img, **__):
  108. return ImageOps.invert(img)
  109. def equalize(img, **__):
  110. return ImageOps.equalize(img)
  111. def solarize(img, thresh, **__):
  112. return ImageOps.solarize(img, thresh)
  113. def solarize_add(img, add, thresh=128, **__):
  114. lut = []
  115. for i in range(256):
  116. if i < thresh:
  117. lut.append(min(255, i + add))
  118. else:
  119. lut.append(i)
  120. if img.mode in ("L", "RGB"):
  121. if img.mode == "RGB" and len(lut) == 256:
  122. lut = lut + lut + lut
  123. return img.point(lut)
  124. else:
  125. return img
  126. def posterize(img, bits_to_keep, **__):
  127. if bits_to_keep >= 8:
  128. return img
  129. return ImageOps.posterize(img, bits_to_keep)
  130. def contrast(img, factor, **__):
  131. return ImageEnhance.Contrast(img).enhance(factor)
  132. def color(img, factor, **__):
  133. return ImageEnhance.Color(img).enhance(factor)
  134. def brightness(img, factor, **__):
  135. return ImageEnhance.Brightness(img).enhance(factor)
  136. def sharpness(img, factor, **__):
  137. return ImageEnhance.Sharpness(img).enhance(factor)
  138. def _randomly_negate(v):
  139. """With 50% prob, negate the value"""
  140. return -v if random.random() > 0.5 else v
  141. def _rotate_level_to_arg(level, _hparams):
  142. # range [-30, 30]
  143. level = (level / _MAX_LEVEL) * 30.
  144. level = _randomly_negate(level)
  145. return level,
  146. def _enhance_level_to_arg(level, _hparams):
  147. # range [0.1, 1.9]
  148. return (level / _MAX_LEVEL) * 1.8 + 0.1,
  149. def _enhance_increasing_level_to_arg(level, _hparams):
  150. # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
  151. # range [0.1, 1.9]
  152. level = (level / _MAX_LEVEL) * .9
  153. level = 1.0 + _randomly_negate(level)
  154. return level,
  155. def _shear_level_to_arg(level, _hparams):
  156. # range [-0.3, 0.3]
  157. level = (level / _MAX_LEVEL) * 0.3
  158. level = _randomly_negate(level)
  159. return level,
  160. def _translate_abs_level_to_arg(level, hparams):
  161. translate_const = hparams['translate_const']
  162. level = (level / _MAX_LEVEL) * float(translate_const)
  163. level = _randomly_negate(level)
  164. return level,
  165. def _translate_rel_level_to_arg(level, hparams):
  166. # default range [-0.45, 0.45]
  167. translate_pct = hparams.get('translate_pct', 0.45)
  168. level = (level / _MAX_LEVEL) * translate_pct
  169. level = _randomly_negate(level)
  170. return level,
  171. def _posterize_level_to_arg(level, _hparams):
  172. # As per Tensorflow TPU EfficientNet impl
  173. # range [0, 4], 'keep 0 up to 4 MSB of original image'
  174. # intensity/severity of augmentation decreases with level
  175. return int((level / _MAX_LEVEL) * 4),
  176. def _posterize_increasing_level_to_arg(level, hparams):
  177. # As per Tensorflow models research and UDA impl
  178. # range [4, 0], 'keep 4 down to 0 MSB of original image',
  179. # intensity/severity of augmentation increases with level
  180. return 4 - _posterize_level_to_arg(level, hparams)[0],
  181. def _posterize_original_level_to_arg(level, _hparams):
  182. # As per original AutoAugment paper description
  183. # range [4, 8], 'keep 4 up to 8 MSB of image'
  184. # intensity/severity of augmentation decreases with level
  185. return int((level / _MAX_LEVEL) * 4) + 4,
  186. def _solarize_level_to_arg(level, _hparams):
  187. # range [0, 256]
  188. # intensity/severity of augmentation decreases with level
  189. return int((level / _MAX_LEVEL) * 256),
  190. def _solarize_increasing_level_to_arg(level, _hparams):
  191. # range [0, 256]
  192. # intensity/severity of augmentation increases with level
  193. return 256 - _solarize_level_to_arg(level, _hparams)[0],
  194. def _solarize_add_level_to_arg(level, _hparams):
  195. # range [0, 110]
  196. return int((level / _MAX_LEVEL) * 110),
  197. LEVEL_TO_ARG = {
  198. 'AutoContrast': None,
  199. 'Equalize': None,
  200. 'Invert': None,
  201. 'Rotate': _rotate_level_to_arg,
  202. # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
  203. 'Posterize': _posterize_level_to_arg,
  204. 'PosterizeIncreasing': _posterize_increasing_level_to_arg,
  205. 'PosterizeOriginal': _posterize_original_level_to_arg,
  206. 'Solarize': _solarize_level_to_arg,
  207. 'SolarizeIncreasing': _solarize_increasing_level_to_arg,
  208. 'SolarizeAdd': _solarize_add_level_to_arg,
  209. 'Color': _enhance_level_to_arg,
  210. 'ColorIncreasing': _enhance_increasing_level_to_arg,
  211. 'Contrast': _enhance_level_to_arg,
  212. 'ContrastIncreasing': _enhance_increasing_level_to_arg,
  213. 'Brightness': _enhance_level_to_arg,
  214. 'BrightnessIncreasing': _enhance_increasing_level_to_arg,
  215. 'Sharpness': _enhance_level_to_arg,
  216. 'SharpnessIncreasing': _enhance_increasing_level_to_arg,
  217. 'ShearX': _shear_level_to_arg,
  218. 'ShearY': _shear_level_to_arg,
  219. 'TranslateX': _translate_abs_level_to_arg,
  220. 'TranslateY': _translate_abs_level_to_arg,
  221. 'TranslateXRel': _translate_rel_level_to_arg,
  222. 'TranslateYRel': _translate_rel_level_to_arg,
  223. }
  224. NAME_TO_OP = {
  225. 'AutoContrast': auto_contrast,
  226. 'Equalize': equalize,
  227. 'Invert': invert,
  228. 'Rotate': rotate,
  229. 'Posterize': posterize,
  230. 'PosterizeIncreasing': posterize,
  231. 'PosterizeOriginal': posterize,
  232. 'Solarize': solarize,
  233. 'SolarizeIncreasing': solarize,
  234. 'SolarizeAdd': solarize_add,
  235. 'Color': color,
  236. 'ColorIncreasing': color,
  237. 'Contrast': contrast,
  238. 'ContrastIncreasing': contrast,
  239. 'Brightness': brightness,
  240. 'BrightnessIncreasing': brightness,
  241. 'Sharpness': sharpness,
  242. 'SharpnessIncreasing': sharpness,
  243. 'ShearX': shear_x,
  244. 'ShearY': shear_y,
  245. 'TranslateX': translate_x_abs,
  246. 'TranslateY': translate_y_abs,
  247. 'TranslateXRel': translate_x_rel,
  248. 'TranslateYRel': translate_y_rel,
  249. }
  250. class AugmentOp(object):
  251. def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
  252. hparams = hparams or _HPARAMS_DEFAULT
  253. self.aug_fn = NAME_TO_OP[name]
  254. self.level_fn = LEVEL_TO_ARG[name]
  255. self.prob = prob
  256. self.magnitude = magnitude
  257. self.hparams = hparams.copy()
  258. self.kwargs = dict(
  259. fillcolor=hparams['img_mean'] if 'img_mean' in hparams else _FILL,
  260. resample=hparams['interpolation']
  261. if 'interpolation' in hparams else _RANDOM_INTERPOLATION, )
  262. # If magnitude_std is > 0, we introduce some randomness
  263. # in the usually fixed policy and sample magnitude from a normal distribution
  264. # with mean `magnitude` and std-dev of `magnitude_std`.
  265. # NOTE This is my own hack, being tested, not in papers or reference impls.
  266. self.magnitude_std = self.hparams.get('magnitude_std', 0)
  267. def __call__(self, img):
  268. if self.prob < 1.0 and random.random() > self.prob:
  269. return img
  270. magnitude = self.magnitude
  271. if self.magnitude_std and self.magnitude_std > 0:
  272. magnitude = random.gauss(magnitude, self.magnitude_std)
  273. magnitude = min(_MAX_LEVEL, max(0, magnitude)) # clip to valid range
  274. level_args = self.level_fn(
  275. magnitude, self.hparams) if self.level_fn is not None else tuple()
  276. return self.aug_fn(img, *level_args, **self.kwargs)
  277. def auto_augment_policy_v0(hparams):
  278. # ImageNet v0 policy from TPU EfficientNet impl, cannot find a paper reference.
  279. policy = [
  280. [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
  281. [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
  282. [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
  283. [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
  284. [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
  285. [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
  286. [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
  287. [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
  288. [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
  289. [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
  290. [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
  291. [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
  292. [('Posterize', 0.4, 6), ('AutoContrast', 0.4, 7)],
  293. [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
  294. [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
  295. [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
  296. [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
  297. [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
  298. [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
  299. [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
  300. [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
  301. [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
  302. [('Posterize', 0.8, 2), ('Solarize', 0.6, 10)
  303. ], # This results in black image with Tpu posterize
  304. [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
  305. [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
  306. ]
  307. pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
  308. return pc
  309. def auto_augment_policy_v0r(hparams):
  310. # ImageNet v0 policy from TPU EfficientNet impl, with variation of Posterize used
  311. # in Google research implementation (number of bits discarded increases with magnitude)
  312. policy = [
  313. [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
  314. [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
  315. [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
  316. [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
  317. [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
  318. [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
  319. [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
  320. [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
  321. [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
  322. [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
  323. [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
  324. [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
  325. [('PosterizeIncreasing', 0.4, 6), ('AutoContrast', 0.4, 7)],
  326. [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
  327. [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
  328. [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
  329. [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
  330. [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
  331. [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
  332. [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
  333. [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
  334. [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
  335. [('PosterizeIncreasing', 0.8, 2), ('Solarize', 0.6, 10)],
  336. [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
  337. [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
  338. ]
  339. pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
  340. return pc
  341. def auto_augment_policy_original(hparams):
  342. # ImageNet policy from https://arxiv.org/abs/1805.09501
  343. policy = [
  344. [('PosterizeOriginal', 0.4, 8), ('Rotate', 0.6, 9)],
  345. [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
  346. [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
  347. [('PosterizeOriginal', 0.6, 7), ('PosterizeOriginal', 0.6, 6)],
  348. [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
  349. [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
  350. [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
  351. [('PosterizeOriginal', 0.8, 5), ('Equalize', 1.0, 2)],
  352. [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
  353. [('Equalize', 0.6, 8), ('PosterizeOriginal', 0.4, 6)],
  354. [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
  355. [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
  356. [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
  357. [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
  358. [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
  359. [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
  360. [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
  361. [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
  362. [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
  363. [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
  364. [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
  365. [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
  366. [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
  367. [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
  368. [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
  369. ]
  370. pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
  371. return pc
  372. def auto_augment_policy_originalr(hparams):
  373. # ImageNet policy from https://arxiv.org/abs/1805.09501 with research posterize variation
  374. policy = [
  375. [('PosterizeIncreasing', 0.4, 8), ('Rotate', 0.6, 9)],
  376. [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
  377. [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
  378. [('PosterizeIncreasing', 0.6, 7), ('PosterizeIncreasing', 0.6, 6)],
  379. [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
  380. [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
  381. [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
  382. [('PosterizeIncreasing', 0.8, 5), ('Equalize', 1.0, 2)],
  383. [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
  384. [('Equalize', 0.6, 8), ('PosterizeIncreasing', 0.4, 6)],
  385. [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
  386. [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
  387. [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
  388. [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
  389. [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
  390. [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
  391. [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
  392. [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
  393. [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
  394. [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
  395. [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
  396. [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
  397. [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
  398. [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
  399. [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
  400. ]
  401. pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
  402. return pc
  403. def auto_augment_policy(name='v0', hparams=None):
  404. hparams = hparams or _HPARAMS_DEFAULT
  405. if name == 'original':
  406. return auto_augment_policy_original(hparams)
  407. elif name == 'originalr':
  408. return auto_augment_policy_originalr(hparams)
  409. elif name == 'v0':
  410. return auto_augment_policy_v0(hparams)
  411. elif name == 'v0r':
  412. return auto_augment_policy_v0r(hparams)
  413. else:
  414. assert False, 'Unknown AA policy (%s)' % name
  415. class AutoAugment(object):
  416. def __init__(self, policy):
  417. self.policy = policy
  418. def __call__(self, img):
  419. sub_policy = random.choice(self.policy)
  420. for op in sub_policy:
  421. img = op(img)
  422. return img
  423. def auto_augment_transform(config_str, hparams):
  424. """
  425. Create a AutoAugment transform
  426. :param config_str: String defining configuration of auto augmentation. Consists of multiple sections separated by
  427. dashes ('-'). The first section defines the AutoAugment policy (one of 'v0', 'v0r', 'original', 'originalr').
  428. The remaining sections, not order sepecific determine
  429. 'mstd' - float std deviation of magnitude noise applied
  430. Ex 'original-mstd0.5' results in AutoAugment with original policy, magnitude_std 0.5
  431. :param hparams: Other hparams (kwargs) for the AutoAugmentation scheme
  432. :return: A callable Transform Op
  433. """
  434. config = config_str.split('-')
  435. policy_name = config[0]
  436. config = config[1:]
  437. for c in config:
  438. cs = re.split(r'(\d.*)', c)
  439. if len(cs) < 2:
  440. continue
  441. key, val = cs[:2]
  442. if key == 'mstd':
  443. # noise param injected via hparams for now
  444. hparams.setdefault('magnitude_std', float(val))
  445. else:
  446. assert False, 'Unknown AutoAugment config section'
  447. aa_policy = auto_augment_policy(policy_name, hparams=hparams)
  448. return AutoAugment(aa_policy)
  449. _RAND_TRANSFORMS = [
  450. 'AutoContrast',
  451. 'Equalize',
  452. 'Invert',
  453. 'Rotate',
  454. 'Posterize',
  455. 'Solarize',
  456. 'SolarizeAdd',
  457. 'Color',
  458. 'Contrast',
  459. 'Brightness',
  460. 'Sharpness',
  461. 'ShearX',
  462. 'ShearY',
  463. 'TranslateXRel',
  464. 'TranslateYRel',
  465. #'Cutout' # NOTE I've implement this as random erasing separately
  466. ]
  467. _RAND_INCREASING_TRANSFORMS = [
  468. 'AutoContrast',
  469. 'Equalize',
  470. 'Invert',
  471. 'Rotate',
  472. 'PosterizeIncreasing',
  473. 'SolarizeIncreasing',
  474. 'SolarizeAdd',
  475. 'ColorIncreasing',
  476. 'ContrastIncreasing',
  477. 'BrightnessIncreasing',
  478. 'SharpnessIncreasing',
  479. 'ShearX',
  480. 'ShearY',
  481. 'TranslateXRel',
  482. 'TranslateYRel',
  483. #'Cutout' # NOTE I've implement this as random erasing separately
  484. ]
  485. # These experimental weights are based loosely on the relative improvements mentioned in paper.
  486. # They may not result in increased performance, but could likely be tuned to so.
  487. _RAND_CHOICE_WEIGHTS_0 = {
  488. 'Rotate': 0.3,
  489. 'ShearX': 0.2,
  490. 'ShearY': 0.2,
  491. 'TranslateXRel': 0.1,
  492. 'TranslateYRel': 0.1,
  493. 'Color': .025,
  494. 'Sharpness': 0.025,
  495. 'AutoContrast': 0.025,
  496. 'Solarize': .005,
  497. 'SolarizeAdd': .005,
  498. 'Contrast': .005,
  499. 'Brightness': .005,
  500. 'Equalize': .005,
  501. 'Posterize': 0,
  502. 'Invert': 0,
  503. }
  504. def _select_rand_weights(weight_idx=0, transforms=None):
  505. transforms = transforms or _RAND_TRANSFORMS
  506. assert weight_idx == 0 # only one set of weights currently
  507. rand_weights = _RAND_CHOICE_WEIGHTS_0
  508. probs = [rand_weights[k] for k in transforms]
  509. probs /= np.sum(probs)
  510. return probs
  511. def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
  512. hparams = hparams or _HPARAMS_DEFAULT
  513. transforms = transforms or _RAND_TRANSFORMS
  514. return [
  515. AugmentOp(
  516. name, prob=0.5, magnitude=magnitude, hparams=hparams)
  517. for name in transforms
  518. ]
  519. class RandAugment(object):
  520. def __init__(self, ops, num_layers=2, choice_weights=None):
  521. self.ops = ops
  522. self.num_layers = num_layers
  523. self.choice_weights = choice_weights
  524. def __call__(self, img):
  525. # no replacement when using weighted choice
  526. ops = np.random.choice(
  527. self.ops,
  528. self.num_layers,
  529. replace=self.choice_weights is None,
  530. p=self.choice_weights)
  531. for op in ops:
  532. img = op(img)
  533. return img
  534. def rand_augment_transform(config_str, hparams):
  535. """
  536. Create a RandAugment transform
  537. :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
  538. dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
  539. sections, not order sepecific determine
  540. 'm' - integer magnitude of rand augment
  541. 'n' - integer num layers (number of transform ops selected per image)
  542. 'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
  543. 'mstd' - float std deviation of magnitude noise applied
  544. 'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
  545. Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
  546. 'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2
  547. :param hparams: Other hparams (kwargs) for the RandAugmentation scheme
  548. :return: A callable Transform Op
  549. """
  550. magnitude = _MAX_LEVEL # default to _MAX_LEVEL for magnitude (currently 10)
  551. num_layers = 2 # default to 2 ops per image
  552. weight_idx = None # default to no probability weights for op choice
  553. transforms = _RAND_TRANSFORMS
  554. config = config_str.split('-')
  555. assert config[0] == 'rand'
  556. config = config[1:]
  557. for c in config:
  558. cs = re.split(r'(\d.*)', c)
  559. if len(cs) < 2:
  560. continue
  561. key, val = cs[:2]
  562. if key == 'mstd':
  563. # noise param injected via hparams for now
  564. hparams.setdefault('magnitude_std', float(val))
  565. elif key == 'inc':
  566. if bool(val):
  567. transforms = _RAND_INCREASING_TRANSFORMS
  568. elif key == 'm':
  569. magnitude = int(val)
  570. elif key == 'n':
  571. num_layers = int(val)
  572. elif key == 'w':
  573. weight_idx = int(val)
  574. else:
  575. assert False, 'Unknown RandAugment config section'
  576. ra_ops = rand_augment_ops(
  577. magnitude=magnitude, hparams=hparams, transforms=transforms)
  578. choice_weights = None if weight_idx is None else _select_rand_weights(
  579. weight_idx)
  580. return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)
  581. _AUGMIX_TRANSFORMS = [
  582. 'AutoContrast',
  583. 'ColorIncreasing', # not in paper
  584. 'ContrastIncreasing', # not in paper
  585. 'BrightnessIncreasing', # not in paper
  586. 'SharpnessIncreasing', # not in paper
  587. 'Equalize',
  588. 'Rotate',
  589. 'PosterizeIncreasing',
  590. 'SolarizeIncreasing',
  591. 'ShearX',
  592. 'ShearY',
  593. 'TranslateXRel',
  594. 'TranslateYRel',
  595. ]
  596. def augmix_ops(magnitude=10, hparams=None, transforms=None):
  597. hparams = hparams or _HPARAMS_DEFAULT
  598. transforms = transforms or _AUGMIX_TRANSFORMS
  599. return [
  600. AugmentOp(
  601. name, prob=1.0, magnitude=magnitude, hparams=hparams)
  602. for name in transforms
  603. ]
  604. class AugMixAugment(object):
  605. """ AugMix Transform
  606. Adapted and improved from impl here: https://github.com/google-research/augmix/blob/master/imagenet.py
  607. From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
  608. https://arxiv.org/abs/1912.02781
  609. """
  610. def __init__(self, ops, alpha=1., width=3, depth=-1, blended=False):
  611. self.ops = ops
  612. self.alpha = alpha
  613. self.width = width
  614. self.depth = depth
  615. self.blended = blended # blended mode is faster but not well tested
  616. def _calc_blended_weights(self, ws, m):
  617. ws = ws * m
  618. cump = 1.
  619. rws = []
  620. for w in ws[::-1]:
  621. alpha = w / cump
  622. cump *= (1 - alpha)
  623. rws.append(alpha)
  624. return np.array(rws[::-1], dtype=np.float32)
  625. def _apply_blended(self, img, mixing_weights, m):
  626. # This is my first crack and implementing a slightly faster mixed augmentation. Instead
  627. # of accumulating the mix for each chain in a Numpy array and then blending with original,
  628. # it recomputes the blending coefficients and applies one PIL image blend per chain.
  629. # TODO the results appear in the right ballpark but they differ by more than rounding.
  630. img_orig = img.copy()
  631. ws = self._calc_blended_weights(mixing_weights, m)
  632. for w in ws:
  633. depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
  634. ops = np.random.choice(self.ops, depth, replace=True)
  635. img_aug = img_orig # no ops are in-place, deep copy not necessary
  636. for op in ops:
  637. img_aug = op(img_aug)
  638. img = Image.blend(img, img_aug, w)
  639. return img
  640. def _apply_basic(self, img, mixing_weights, m):
  641. # This is a literal adaptation of the paper/official implementation without normalizations and
  642. # PIL <-> Numpy conversions between every op. It is still quite CPU compute heavy compared to the
  643. # typical augmentation transforms, could use a GPU / Kornia implementation.
  644. img_shape = img.size[0], img.size[1], len(img.getbands())
  645. mixed = np.zeros(img_shape, dtype=np.float32)
  646. for mw in mixing_weights:
  647. depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
  648. ops = np.random.choice(self.ops, depth, replace=True)
  649. img_aug = img # no ops are in-place, deep copy not necessary
  650. for op in ops:
  651. img_aug = op(img_aug)
  652. mixed += mw * np.asarray(img_aug, dtype=np.float32)
  653. np.clip(mixed, 0, 255., out=mixed)
  654. mixed = Image.fromarray(mixed.astype(np.uint8))
  655. return Image.blend(img, mixed, m)
  656. def __call__(self, img):
  657. mixing_weights = np.float32(
  658. np.random.dirichlet([self.alpha] * self.width))
  659. m = np.float32(np.random.beta(self.alpha, self.alpha))
  660. if self.blended:
  661. mixed = self._apply_blended(img, mixing_weights, m)
  662. else:
  663. mixed = self._apply_basic(img, mixing_weights, m)
  664. return mixed
  665. def augment_and_mix_transform(config_str, hparams):
  666. """ Create AugMix transform
  667. :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
  668. dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
  669. sections, not order sepecific determine
  670. 'm' - integer magnitude (severity) of augmentation mix (default: 3)
  671. 'w' - integer width of augmentation chain (default: 3)
  672. 'd' - integer depth of augmentation chain (-1 is random [1, 3], default: -1)
  673. 'b' - integer (bool), blend each branch of chain into end result without a final blend, less CPU (default: 0)
  674. 'mstd' - float std deviation of magnitude noise applied (default: 0)
  675. Ex 'augmix-m5-w4-d2' results in AugMix with severity 5, chain width 4, chain depth 2
  676. :param hparams: Other hparams (kwargs) for the Augmentation transforms
  677. :return: A callable Transform Op
  678. """
  679. magnitude = 3
  680. width = 3
  681. depth = -1
  682. alpha = 1.
  683. blended = False
  684. config = config_str.split('-')
  685. assert config[0] == 'augmix'
  686. config = config[1:]
  687. for c in config:
  688. cs = re.split(r'(\d.*)', c)
  689. if len(cs) < 2:
  690. continue
  691. key, val = cs[:2]
  692. if key == 'mstd':
  693. # noise param injected via hparams for now
  694. hparams.setdefault('magnitude_std', float(val))
  695. elif key == 'm':
  696. magnitude = int(val)
  697. elif key == 'w':
  698. width = int(val)
  699. elif key == 'd':
  700. depth = int(val)
  701. elif key == 'a':
  702. alpha = float(val)
  703. elif key == 'b':
  704. blended = bool(val)
  705. else:
  706. assert False, 'Unknown AugMix config section'
  707. ops = augmix_ops(magnitude=magnitude, hparams=hparams)
  708. return AugMixAugment(
  709. ops, alpha=alpha, width=width, depth=depth, blended=blended)
  710. class RawTimmAutoAugment(object):
  711. """TimmAutoAugment API for PaddleClas."""
  712. def __init__(self,
  713. config_str="rand-m9-mstd0.5-inc1",
  714. interpolation="bicubic",
  715. img_size=224,
  716. mean=IMAGENET_DEFAULT_MEAN):
  717. if isinstance(img_size, (tuple, list)):
  718. img_size_min = min(img_size)
  719. else:
  720. img_size_min = img_size
  721. aa_params = dict(
  722. translate_const=int(img_size_min * 0.45),
  723. img_mean=tuple([min(255, round(255 * x)) for x in mean]), )
  724. if interpolation and interpolation != 'random':
  725. aa_params['interpolation'] = _pil_interp(interpolation)
  726. if config_str.startswith('rand'):
  727. self.augment_func = rand_augment_transform(config_str, aa_params)
  728. elif config_str.startswith('augmix'):
  729. aa_params['translate_pct'] = 0.3
  730. self.augment_func = augment_and_mix_transform(config_str,
  731. aa_params)
  732. elif config_str.startswith('auto'):
  733. self.augment_func = auto_augment_transform(config_str, aa_params)
  734. else:
  735. raise Exception(
  736. "ConfigError: The TimmAutoAugment Op only support RandAugment, AutoAugment, AugMix, and the config_str only starts with \"rand\", \"augmix\", \"auto\"."
  737. )
  738. def __call__(self, img):
  739. return self.augment_func(img)