content: # ------------------------------------------------------------------------
  # UP-DETR
  # Copyright (c) Tencent, Inc. and its affiliates. All Rights Reserved.
  # ------------------------------------------------------------------------
  # Modified from DETR (https://github.com/facebookresearch/detr)
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  # ------------------------------------------------------------------------
  """
  Transforms and data augmentation for both image + bbox.
  """
  import random
  
  import PIL
  import torch
  import torchvision.transforms as T
  import torchvision.transforms.functional as F
  
  from util.box_ops import box_xyxy_to_cxcywh
  from util.misc import interpolate
  
  
  def crop(image, target, region):
      cropped_image = F.crop(image, *region)
  
      target = target.copy()
      i, j, h, w = region
  
      # should we do something wrt the original size?
      target["size"] = torch.tensor([h, w])
  
      fields = ["labels", "area", "iscrowd"]
  
      if "boxes" in target:
          boxes = target["boxes"]
          max_size = torch.as_tensor([w, h], dtype=torch.float32)
          cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
          cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
          cropped_boxes = cropped_boxes.clamp(min=0)
          area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
          target["boxes"] = cropped_boxes.reshape(-1, 4)
          target["area"] = area
          fields.append("boxes")
  
      if "masks" in target:
          # FIXME should we update the area here if there are no boxes?
          target['masks'] = target['masks'][:, i:i + h, j:j + w]
          fields.append("masks")
  
      # remove elements for which the boxes or masks that have zero area
      if "boxes" in target or "masks" in target:
          # favor boxes selection when defining which elements to keep
          # this is compatible with previous implementation
          if "boxes" in target:
              cropped_boxes = target['boxes'].reshape(-1, 2, 2)
              keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
          else:
              keep = target['masks'].flatten(1).any(1)
  
          for field in fields:
              target[field] = target[field][keep]
  
      return cropped_image, target
  
  
  def hflip(image, target):
      flipped_image = F.hflip(image)
  
      w, h = image.size
  
      target = target.copy()
      if "boxes" in target:
          boxes = target["boxes"]
          boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0])
          target["boxes"] = boxes
  
      if "masks" in target:
          target['masks'] = target['masks'].flip(-1)
  
      return flipped_image, target
  
  
  def resize(image, target, size, max_size=None):
      # size can be min_size (scalar) or (w, h) tuple
  
      def get_size_with_aspect_ratio(image_size, size, max_size=None):
          w, h = image_size
          if max_size is not None:
              min_original_size = float(min((w, h)))
              max_original_size = float(max((w, h)))
              if max_original_size / min_original_size * size > max_size:
                  size = int(round(max_size * min_original_size / max_original_size))
  
          if (w <= h and w == size) or (h <= w and h == size):
              return (h, w)
  
          if w < h:
              ow = size
              oh = int(size * h / w)
          else:
              oh = size
              ow = int(size * w / h)
  
          return (oh, ow)
  
      def get_size(image_size, size, max_size=None):
          if isinstance(size, (list, tuple)):
              return size[::-1]
          else:
              return get_size_with_aspect_ratio(image_size, size, max_size)
  
      size = get_size(image.size, size, max_size)
      rescaled_image = F.resize(image, size)
  
      if target is None:
          return rescaled_image, None
  
      ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))
      ratio_width, ratio_height = ratios
  
      target = target.copy()
      if "boxes" in target:
          boxes = target["boxes"]
          scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
          target["boxes"] = scaled_boxes
  
      if "area" in target:
          area = target["area"]
          scaled_area = area * (ratio_width * ratio_height)
          target["area"] = scaled_area
  
      h, w = size
      target["size"] = torch.tensor([h, w])
  
      if "masks" in target:
          target['masks'] = interpolate(
              target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
  
      return rescaled_image, target
  
  
  def pad(image, target, padding):
      # assumes that we only pad on the bottom right corners
      padded_image = F.pad(image, (0, 0, padding[0], padding[1]))
      if target is None:
          return padded_image, None
      target = target.copy()
      # should we do something wrt the original size?
      target["size"] = torch.tensor(padded_image.size[::-1])
      if "masks" in target:
          target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[0], 0, padding[1]))
      return padded_image, target
  
  
  class RandomCrop(object):
      def __init__(self, size):
          self.size = size
  
      def __call__(self, img, target):
          region = T.RandomCrop.get_params(img, self.size)
          return crop(img, target, region)
  
  
  class RandomSizeCrop(object):
      def __init__(self, min_size: int, max_size: int):
          self.min_size = min_size
          self.max_size = max_size
  
      def __call__(self, img: PIL.Image.Image, target: dict):
          w = random.randint(self.min_size, min(img.width, self.max_size))
          h = random.randint(self.min_size, min(img.height, self.max_size))
          region = T.RandomCrop.get_params(img, [h, w])
          return crop(img, target, region)
  
  
  class CenterCrop(object):
      def __init__(self, size):
          self.size = size
  
      def __call__(self, img, target):
          image_width, image_height = img.size
          crop_height, crop_width = self.size
          crop_top = int(round((image_height - crop_height) / 2.))
          crop_left = int(round((image_width - crop_width) / 2.))
          return crop(img, target, (crop_top, crop_left, crop_height, crop_width))
  
  
  class RandomHorizontalFlip(object):
      def __init__(self, p=0.5):
          self.p = p
  
      def __call__(self, img, target):
          if random.random() < self.p:
              return hflip(img, target)
          return img, target
  
  
  class RandomResize(object):
      def __init__(self, sizes, max_size=None):
          assert isinstance(sizes, (list, tuple))
          self.sizes = sizes
          self.max_size = max_size
  
      def __call__(self, img, target=None):
          size = random.choice(self.sizes)
          return resize(img, target, size, self.max_size)
  
  
  class RandomPad(object):
      def __init__(self, max_pad):
          self.max_pad = max_pad
  
      def __call__(self, img, target):
          pad_x = random.randint(0, self.max_pad)
          pad_y = random.randint(0, self.max_pad)
          return pad(img, target, (pad_x, pad_y))
  
  
  class RandomSelect(object):
      """
      Randomly selects between transforms1 and transforms2,
      with probability p for transforms1 and (1 - p) for transforms2
      """
      def __init__(self, transforms1, transforms2, p=0.5):
          self.transforms1 = transforms1
          self.transforms2 = transforms2
          self.p = p
  
      def __call__(self, img, target):
          if random.random() < self.p:
              return self.transforms1(img, target)
          return self.transforms2(img, target)
  
  
  class ToTensor(object):
      def __call__(self, img, target):
          return F.to_tensor(img), target
  
  
  class RandomErasing(object):
  
      def __init__(self, *args, **kwargs):
          self.eraser = T.RandomErasing(*args, **kwargs)
  
      def __call__(self, img, target):
          return self.eraser(img), target
  
  
  class Normalize(object):
      def __init__(self, mean, std):
          self.mean = mean
          self.std = std
  
      def __call__(self, image, target=None):
          image = F.normalize(image, mean=self.mean, std=self.std)
          if target is None:
              return image, None
          target = target.copy()
          h, w = image.shape[-2:]
          if "boxes" in target:
              boxes = target["boxes"]
              boxes = box_xyxy_to_cxcywh(boxes)
              boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
              target["boxes"] = boxes
          return image, target
  
  
  class Compose(object):
      def __init__(self, transforms):
          self.transforms = transforms
  
      def __call__(self, image, target):
          for t in self.transforms:
              image, target = t(image, target)
          return image, target
  
      def __repr__(self):
          format_string = self.__class__.__name__ + "("
          for t in self.transforms:
              format_string += "\n"
              format_string += "    {0}".format(t)
          format_string += "\n)"
          return format_string
  