text: # encoding: utf-8
  """
  @author:  liaoxingyu
  @contact: liaoxingyu2@jd.com
  """
  
  import math
  import random
  import cv2
  from collections import deque
  from PIL import Image, ImageFilter
  import torch
  
  class RandomErasing(object):
      """ Randomly selects a rectangle region in an image and erases its pixels.
          'Random Erasing Data Augmentation' by Zhong et al.
          See https://arxiv.org/pdf/1708.04896.pdf
      Args:
           probability: The probability that the Random Erasing operation will be performed.
           sl: Minimum proportion of erased area against input image.
           sh: Maximum proportion of erased area against input image.
           r1: Minimum aspect ratio of erased area.
           mean: Erasing value.
      """
  
      def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)):
          self.probability = probability
          self.mean = mean
          self.sl = sl
          self.sh = sh
          self.r1 = r1
  
      def __call__(self, img):
  
          if random.uniform(0, 1) >= self.probability:
              return img
  
          for attempt in range(100):
              area = img.size()[1] * img.size()[2]
  
              target_area = random.uniform(self.sl, self.sh) * area
              aspect_ratio = random.uniform(self.r1, 1 / self.r1)
  
              h = int(round(math.sqrt(target_area * aspect_ratio)))
              w = int(round(math.sqrt(target_area / aspect_ratio)))
  
              if w < img.size()[2] and h < img.size()[1]:
                  x1 = random.randint(0, img.size()[1] - h)
                  y1 = random.randint(0, img.size()[2] - w)
                  if img.size()[0] == 3:
                      img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
                      img[1, x1:x1 + h, y1:y1 + w] = self.mean[1]
                      img[2, x1:x1 + h, y1:y1 + w] = self.mean[2]
                  else:
                      img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
                  return img
  
          return img
  
  class RandomPatch(object):
      """Random patch data augmentation.
      There is a patch pool that stores randomly extracted pathces from person images.
  
      For each input image, RandomPatch
          1) extracts a random patch and stores the patch in the patch pool;
          2) randomly selects a patch from the patch pool and pastes it on the
             input (at random position) to simulate occlusion.
      Reference:
          - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.
          - Zhou et al. Learning Generalisable Omni-Scale Representations
            for Person Re-Identification. arXiv preprint, 2019.
      """
  
      def __init__(
              self,
              prob_happen=0.5,
              pool_capacity=50000,
              min_sample_size=100,
              patch_min_area=0.01,
              patch_max_area=0.5,
              patch_min_ratio=0.1,
              prob_rotate=0.5,
              prob_flip_leftright=0.5,
      ):
          self.prob_happen = prob_happen
  
          self.patch_min_area = patch_min_area
          self.patch_max_area = patch_max_area
          self.patch_min_ratio = patch_min_ratio
  
          self.prob_rotate = prob_rotate
          self.prob_flip_leftright = prob_flip_leftright
  
          self.patchpool = deque(maxlen=pool_capacity)
          self.min_sample_size = min_sample_size
  
      def generate_wh(self, W, H):
          area = W * H
          for attempt in range(100):
              target_area = random.uniform(
                  self.patch_min_area, self.patch_max_area
              ) * area
              aspect_ratio = random.uniform(
                  self.patch_min_ratio, 1. / self.patch_min_ratio
              )
              h = int(round(math.sqrt(target_area * aspect_ratio)))
              w = int(round(math.sqrt(target_area / aspect_ratio)))
              if w < W and h < H:
                  return w, h
          return None, None
  
      def transform_patch(self, patch):
          if random.uniform(0, 1) > self.prob_flip_leftright:
              patch = patch.transpose(Image.FLIP_LEFT_RIGHT)
          if random.uniform(0, 1) > self.prob_rotate:
              patch = patch.rotate(random.randint(-10, 10))
          return patch
  
      def __call__(self, img):
          W, H = img.size  # original image size
  
          # collect new patch
          w, h = self.generate_wh(W, H)
          if w is not None and h is not None:
              x1 = random.randint(0, W - w)
              y1 = random.randint(0, H - h)
              new_patch = img.crop((x1, y1, x1 + w, y1 + h))
              self.patchpool.append(new_patch)
  
          if len(self.patchpool) < self.min_sample_size:
              return img
  
          if random.uniform(0, 1) > self.prob_happen:
              return img
  
          # paste a randomly selected patch on a random position
          patch = random.sample(self.patchpool, 1)[0]
          patchW, patchH = patch.size
          x1 = random.randint(0, W - patchW)
          y1 = random.randint(0, H - patchH)
          patch = self.transform_patch(patch)
          img.paste(patch, (x1, y1))
  
          return img
  
  