functional.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  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. # encoding: utf-8
  15. import numpy as np
  16. from PIL import Image, ImageOps, ImageEnhance
  17. def int_parameter(level, maxval):
  18. """Helper function to scale `val` between 0 and maxval .
  19. Args:
  20. level: Level of the operation that will be between [0, `PARAMETER_MAX`].
  21. maxval: Maximum value that the operation can have. This will be scaled to
  22. level/PARAMETER_MAX.
  23. Returns:
  24. An int that results from scaling `maxval` according to `level`.
  25. """
  26. return int(level * maxval / 10)
  27. def float_parameter(level, maxval):
  28. """Helper function to scale `val` between 0 and maxval.
  29. Args:
  30. level: Level of the operation that will be between [0, `PARAMETER_MAX`].
  31. maxval: Maximum value that the operation can have. This will be scaled to
  32. level/PARAMETER_MAX.
  33. Returns:
  34. A float that results from scaling `maxval` according to `level`.
  35. """
  36. return float(level) * maxval / 10.
  37. def sample_level(n):
  38. return np.random.uniform(low=0.1, high=n)
  39. def autocontrast(pil_img, *args):
  40. return ImageOps.autocontrast(pil_img)
  41. def equalize(pil_img, *args):
  42. return ImageOps.equalize(pil_img)
  43. def posterize(pil_img, level, *args):
  44. level = int_parameter(sample_level(level), 4)
  45. return ImageOps.posterize(pil_img, 4 - level)
  46. def rotate(pil_img, level, *args):
  47. degrees = int_parameter(sample_level(level), 30)
  48. if np.random.uniform() > 0.5:
  49. degrees = -degrees
  50. return pil_img.rotate(degrees, resample=Image.BILINEAR)
  51. def solarize(pil_img, level, *args):
  52. level = int_parameter(sample_level(level), 256)
  53. return ImageOps.solarize(pil_img, 256 - level)
  54. def shear_x(pil_img, level):
  55. level = float_parameter(sample_level(level), 0.3)
  56. if np.random.uniform() > 0.5:
  57. level = -level
  58. return pil_img.transform(
  59. pil_img.size,
  60. Image.AFFINE, (1, level, 0, 0, 1, 0),
  61. resample=Image.BILINEAR)
  62. def shear_y(pil_img, level):
  63. level = float_parameter(sample_level(level), 0.3)
  64. if np.random.uniform() > 0.5:
  65. level = -level
  66. return pil_img.transform(
  67. pil_img.size,
  68. Image.AFFINE, (1, 0, 0, level, 1, 0),
  69. resample=Image.BILINEAR)
  70. def translate_x(pil_img, level):
  71. level = int_parameter(sample_level(level), pil_img.size[0] / 3)
  72. if np.random.random() > 0.5:
  73. level = -level
  74. return pil_img.transform(
  75. pil_img.size,
  76. Image.AFFINE, (1, 0, level, 0, 1, 0),
  77. resample=Image.BILINEAR)
  78. def translate_y(pil_img, level):
  79. level = int_parameter(sample_level(level), pil_img.size[1] / 3)
  80. if np.random.random() > 0.5:
  81. level = -level
  82. return pil_img.transform(
  83. pil_img.size,
  84. Image.AFFINE, (1, 0, 0, 0, 1, level),
  85. resample=Image.BILINEAR)
  86. # operation that overlaps with ImageNet-C's test set
  87. def color(pil_img, level, *args):
  88. level = float_parameter(sample_level(level), 1.8) + 0.1
  89. return ImageEnhance.Color(pil_img).enhance(level)
  90. # operation that overlaps with ImageNet-C's test set
  91. def contrast(pil_img, level, *args):
  92. level = float_parameter(sample_level(level), 1.8) + 0.1
  93. return ImageEnhance.Contrast(pil_img).enhance(level)
  94. # operation that overlaps with ImageNet-C's test set
  95. def brightness(pil_img, level, *args):
  96. level = float_parameter(sample_level(level), 1.8) + 0.1
  97. return ImageEnhance.Brightness(pil_img).enhance(level)
  98. # operation that overlaps with ImageNet-C's test set
  99. def sharpness(pil_img, level, *args):
  100. level = float_parameter(sample_level(level), 1.8) + 0.1
  101. return ImageEnhance.Sharpness(pil_img).enhance(level)
  102. augmentations = [
  103. autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
  104. translate_x, translate_y
  105. ]