fmix.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217
  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. import math
  15. import random
  16. import numpy as np
  17. from scipy.stats import beta
  18. def fftfreqnd(h, w=None, z=None):
  19. """ Get bin values for discrete fourier transform of size (h, w, z)
  20. :param h: Required, first dimension size
  21. :param w: Optional, second dimension size
  22. :param z: Optional, third dimension size
  23. """
  24. fz = fx = 0
  25. fy = np.fft.fftfreq(h)
  26. if w is not None:
  27. fy = np.expand_dims(fy, -1)
  28. if w % 2 == 1:
  29. fx = np.fft.fftfreq(w)[:w // 2 + 2]
  30. else:
  31. fx = np.fft.fftfreq(w)[:w // 2 + 1]
  32. if z is not None:
  33. fy = np.expand_dims(fy, -1)
  34. if z % 2 == 1:
  35. fz = np.fft.fftfreq(z)[:, None]
  36. else:
  37. fz = np.fft.fftfreq(z)[:, None]
  38. return np.sqrt(fx * fx + fy * fy + fz * fz)
  39. def get_spectrum(freqs, decay_power, ch, h, w=0, z=0):
  40. """ Samples a fourier image with given size and frequencies decayed by decay power
  41. :param freqs: Bin values for the discrete fourier transform
  42. :param decay_power: Decay power for frequency decay prop 1/f**d
  43. :param ch: Number of channels for the resulting mask
  44. :param h: Required, first dimension size
  45. :param w: Optional, second dimension size
  46. :param z: Optional, third dimension size
  47. """
  48. scale = np.ones(1) / (np.maximum(freqs, np.array([1. / max(w, h, z)]))
  49. **decay_power)
  50. param_size = [ch] + list(freqs.shape) + [2]
  51. param = np.random.randn(*param_size)
  52. scale = np.expand_dims(scale, -1)[None, :]
  53. return scale * param
  54. def make_low_freq_image(decay, shape, ch=1):
  55. """ Sample a low frequency image from fourier space
  56. :param decay_power: Decay power for frequency decay prop 1/f**d
  57. :param shape: Shape of desired mask, list up to 3 dims
  58. :param ch: Number of channels for desired mask
  59. """
  60. freqs = fftfreqnd(*shape)
  61. spectrum = get_spectrum(freqs, decay, ch,
  62. *shape) #.reshape((1, *shape[:-1], -1))
  63. spectrum = spectrum[:, 0] + 1j * spectrum[:, 1]
  64. mask = np.real(np.fft.irfftn(spectrum, shape))
  65. if len(shape) == 1:
  66. mask = mask[:1, :shape[0]]
  67. if len(shape) == 2:
  68. mask = mask[:1, :shape[0], :shape[1]]
  69. if len(shape) == 3:
  70. mask = mask[:1, :shape[0], :shape[1], :shape[2]]
  71. mask = mask
  72. mask = (mask - mask.min())
  73. mask = mask / mask.max()
  74. return mask
  75. def sample_lam(alpha, reformulate=False):
  76. """ Sample a lambda from symmetric beta distribution with given alpha
  77. :param alpha: Alpha value for beta distribution
  78. :param reformulate: If True, uses the reformulation of [1].
  79. """
  80. if reformulate:
  81. lam = beta.rvs(alpha + 1, alpha)
  82. else:
  83. lam = beta.rvs(alpha, alpha)
  84. return lam
  85. def binarise_mask(mask, lam, in_shape, max_soft=0.0):
  86. """ Binarises a given low frequency image such that it has mean lambda.
  87. :param mask: Low frequency image, usually the result of `make_low_freq_image`
  88. :param lam: Mean value of final mask
  89. :param in_shape: Shape of inputs
  90. :param max_soft: Softening value between 0 and 0.5 which smooths hard edges in the mask.
  91. :return:
  92. """
  93. idx = mask.reshape(-1).argsort()[::-1]
  94. mask = mask.reshape(-1)
  95. num = math.ceil(lam * mask.size) if random.random() > 0.5 else math.floor(
  96. lam * mask.size)
  97. eff_soft = max_soft
  98. if max_soft > lam or max_soft > (1 - lam):
  99. eff_soft = min(lam, 1 - lam)
  100. soft = int(mask.size * eff_soft)
  101. num_low = int(num - soft)
  102. num_high = int(num + soft)
  103. mask[idx[:num_high]] = 1
  104. mask[idx[num_low:]] = 0
  105. mask[idx[num_low:num_high]] = np.linspace(1, 0, (num_high - num_low))
  106. mask = mask.reshape((1, 1, in_shape[0], in_shape[1]))
  107. return mask
  108. def sample_mask(alpha, decay_power, shape, max_soft=0.0, reformulate=False):
  109. """ Samples a mean lambda from beta distribution parametrised by alpha, creates a low frequency image and binarises
  110. it based on this lambda
  111. :param alpha: Alpha value for beta distribution from which to sample mean of mask
  112. :param decay_power: Decay power for frequency decay prop 1/f**d
  113. :param shape: Shape of desired mask, list up to 3 dims
  114. :param max_soft: Softening value between 0 and 0.5 which smooths hard edges in the mask.
  115. :param reformulate: If True, uses the reformulation of [1].
  116. """
  117. if isinstance(shape, int):
  118. shape = (shape, )
  119. # Choose lambda
  120. lam = sample_lam(alpha, reformulate)
  121. # Make mask, get mean / std
  122. mask = make_low_freq_image(decay_power, shape)
  123. mask = binarise_mask(mask, lam, shape, max_soft)
  124. return float(lam), mask
  125. def sample_and_apply(x,
  126. alpha,
  127. decay_power,
  128. shape,
  129. max_soft=0.0,
  130. reformulate=False):
  131. """
  132. :param x: Image batch on which to apply fmix of shape [b, c, shape*]
  133. :param alpha: Alpha value for beta distribution from which to sample mean of mask
  134. :param decay_power: Decay power for frequency decay prop 1/f**d
  135. :param shape: Shape of desired mask, list up to 3 dims
  136. :param max_soft: Softening value between 0 and 0.5 which smooths hard edges in the mask.
  137. :param reformulate: If True, uses the reformulation of [1].
  138. :return: mixed input, permutation indices, lambda value of mix,
  139. """
  140. lam, mask = sample_mask(alpha, decay_power, shape, max_soft, reformulate)
  141. index = np.random.permutation(x.shape[0])
  142. x1, x2 = x * mask, x[index] * (1 - mask)
  143. return x1 + x2, index, lam
  144. class FMixBase:
  145. """ FMix augmentation
  146. Args:
  147. decay_power (float): Decay power for frequency decay prop 1/f**d
  148. alpha (float): Alpha value for beta distribution from which to sample mean of mask
  149. size ([int] | [int, int] | [int, int, int]): Shape of desired mask, list up to 3 dims
  150. max_soft (float): Softening value between 0 and 0.5 which smooths hard edges in the mask.
  151. reformulate (bool): If True, uses the reformulation of [1].
  152. """
  153. def __init__(self,
  154. decay_power=3,
  155. alpha=1,
  156. size=(32, 32),
  157. max_soft=0.0,
  158. reformulate=False):
  159. super().__init__()
  160. self.decay_power = decay_power
  161. self.reformulate = reformulate
  162. self.size = size
  163. self.alpha = alpha
  164. self.max_soft = max_soft
  165. self.index = None
  166. self.lam = None
  167. def __call__(self, x):
  168. raise NotImplementedError
  169. def loss(self, *args, **kwargs):
  170. raise NotImplementedError