random_erasing.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  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. #This code is adapted from https://github.com/zhunzhong07/Random-Erasing, and refer to Timm.
  15. from functools import partial
  16. import math
  17. import random
  18. import numpy as np
  19. class Pixels(object):
  20. def __init__(self, mode="const", mean=[0., 0., 0.]):
  21. self._mode = mode
  22. self._mean = mean
  23. def __call__(self, h=224, w=224, c=3):
  24. if self._mode == "rand":
  25. return np.random.normal(size=(1, 1, 3))
  26. elif self._mode == "pixel":
  27. return np.random.normal(size=(h, w, c))
  28. elif self._mode == "const":
  29. return self._mean
  30. else:
  31. raise Exception(
  32. "Invalid mode in RandomErasing, only support \"const\", \"rand\", \"pixel\""
  33. )
  34. class RandomErasing(object):
  35. """RandomErasing.
  36. """
  37. def __init__(self,
  38. EPSILON=0.5,
  39. sl=0.02,
  40. sh=0.4,
  41. r1=0.3,
  42. mean=[0., 0., 0.],
  43. attempt=100,
  44. use_log_aspect=False,
  45. mode='const'):
  46. self.EPSILON = eval(EPSILON) if isinstance(EPSILON, str) else EPSILON
  47. self.sl = eval(sl) if isinstance(sl, str) else sl
  48. self.sh = eval(sh) if isinstance(sh, str) else sh
  49. r1 = eval(r1) if isinstance(r1, str) else r1
  50. self.r1 = (math.log(r1), math.log(1 / r1)) if use_log_aspect else (
  51. r1, 1 / r1)
  52. self.use_log_aspect = use_log_aspect
  53. self.attempt = attempt
  54. self.get_pixels = Pixels(mode, mean)
  55. def __call__(self, img):
  56. if random.random() > self.EPSILON:
  57. return img
  58. for _ in range(self.attempt):
  59. area = img.shape[0] * img.shape[1]
  60. target_area = random.uniform(self.sl, self.sh) * area
  61. aspect_ratio = random.uniform(*self.r1)
  62. if self.use_log_aspect:
  63. aspect_ratio = math.exp(aspect_ratio)
  64. h = int(round(math.sqrt(target_area * aspect_ratio)))
  65. w = int(round(math.sqrt(target_area / aspect_ratio)))
  66. if w < img.shape[1] and h < img.shape[0]:
  67. pixels = self.get_pixels(h, w, img.shape[2])
  68. x1 = random.randint(0, img.shape[0] - h)
  69. y1 = random.randint(0, img.shape[1] - w)
  70. if img.shape[2] == 3:
  71. img[x1:x1 + h, y1:y1 + w, :] = pixels
  72. else:
  73. img[x1:x1 + h, y1:y1 + w, 0] = pixels[0]
  74. return img
  75. return img