text: import math
  
  import cv2
  import numpy as np
  import torch
  import torchvision.transforms.functional as F
  
  from plasticorigins.detection.centernet.models import (
      create_model as create_base,
  )
  
  
  class ResizeForCenterNet:
      def __init__(self, fix_res=False):
          self.fix_res = fix_res
  
      def __call__(self, image):
          if self.fix_res:
              new_h = 512
              new_w = 512
          else:
              w, h = image.size
              new_h = (h | 31) + 1
              new_w = (w | 31) + 1
          image = F.resize(image, (new_h, new_w))
          return image
  
  
  def gaussian_radius(det_size, min_overlap=0.7):
      height, width = det_size
      a1 = 1
      b1 = height + width
      c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
      sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
      r1 = (b1 + sq1) / 2
  
      a2 = 4
      b2 = 2 * (height + width)
      c2 = (1 - min_overlap) * width * height
      sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
      r2 = (b2 + sq2) / 2
  
      a3 = 4 * min_overlap
      b3 = -2 * min_overlap * (height + width)
      c3 = (min_overlap - 1) * width * height
      sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
      r3 = (b3 + sq3) / 2
      return min(r1, r2, r3)
  
  
  def gaussian2D(shape, sigma=1):
      m, n = ((ss - 1.0) / 2.0 for ss in shape)
      y, x = np.ogrid[-m : m + 1, -n : n + 1]
  
      h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
      h[h < np.finfo(h.dtype).eps * h.max()] = 0
      return h
  
  
  def draw_umich_gaussian(heatmap, center, radius, k=1):
      diameter = 2 * radius + 1
      gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)
  
      x, y = int(center[0]), int(center[1])
  
      height, width = heatmap.shape[0:2]
  
      left, right = min(x, radius), min(width - x, radius + 1)
      top, bottom = min(y, radius), min(height - y, radius + 1)
  
      masked_heatmap = heatmap[y - top : y + bottom, x - left : x + right]
      masked_gaussian = gaussian[
          radius - top : radius + bottom, radius - left : radius + right
      ]
      if (
          min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0
      ):  # TODO debug
          np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
  
      return heatmap
  
  
  def blob_for_bbox(bbox, heatmap, downsampling_factor=None):
      if downsampling_factor is not None:
          left, top, w, h = (
              bbox_coord // downsampling_factor for bbox_coord in bbox
          )
      else:
          left, top, w, h = (bbox_coord for bbox_coord in bbox)
  
      right, bottom = left + w, top + h
      ct_int = None
      if h > 0 and w > 0:
          radius = gaussian_radius((math.ceil(h), math.ceil(w)))
          radius = max(0, int(radius))
          ct = np.array(
              [(left + right) / 2, (top + bottom) / 2], dtype=np.float32
          )
          ct_int = ct.astype(np.int32)
          heatmap = draw_umich_gaussian(heatmap, ct_int, radius)
      return heatmap, ct_int
  
  
  def pre_process_centernet(image, meta=None, fix_res=True):
      scale = 1.0
      mean = [0.408, 0.447, 0.47]
      std = [0.289, 0.274, 0.278]
      height, width = image.shape[0:2]
      new_height = int(height * scale)
      new_width = int(width * scale)
      if fix_res:
          inp_height, inp_width = 512, 512
          c = np.array([new_width / 2.0, new_height / 2.0], dtype=np.float32)
          s = max(height, width) * 1.0
      else:
          inp_height = (new_height | 31) + 1
          inp_width = (new_width | 31) + 1
          c = np.array([new_width // 2, new_height // 2], dtype=np.float32)
          s = np.array([inp_width, inp_height], dtype=np.float32)
  
      trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
      resized_image = cv2.resize(image, (new_width, new_height))
      inp_image = cv2.warpAffine(
          resized_image,
          trans_input,
          (inp_width, inp_height),
          flags=cv2.INTER_LINEAR,
      )
      inp_image = ((inp_image / 255.0 - mean) / std).astype(np.float32)
  
      images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)
      # if self.opt.flip_test:
      #     images = np.concatenate((images, images[:, :, :, ::-1]), axis=0)
      images = torch.from_numpy(images)
      # meta = {'c': c, 's': s,
      #         'out_height': inp_height // self.opt.down_ratio,
      #         'out_width': inp_width // self.opt.down_ratio}
      return images.squeeze()  # , meta
  
  
  def get_affine_transform(
      center,
      scale,
      rot,
      output_size,
      shift=np.array([0, 0], dtype=np.float32),
      inv=0,
  ):
      if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
          scale = np.array([scale, scale], dtype=np.float32)
  
      scale_tmp = scale
      src_w = scale_tmp[0]
      dst_w = output_size[0]
      dst_h = output_size[1]
  
      rot_rad = np.pi * rot / 180
      src_dir = get_dir([0, src_w * -0.5], rot_rad)
      dst_dir = np.array([0, dst_w * -0.5], np.float32)
  
      src = np.zeros((3, 2), dtype=np.float32)
      dst = np.zeros((3, 2), dtype=np.float32)
      src[0, :] = center + scale_tmp * shift
      src[1, :] = center + src_dir + scale_tmp * shift
      dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
      dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir
  
      src[2:, :] = get_3rd_point(src[0, :], src[1, :])
      dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
  
      if inv:
          trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
      else:
          trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
  
      return trans
  
  
  def get_dir(src_point, rot_rad):
      sn, cs = np.sin(rot_rad), np.cos(rot_rad)
  
      src_result = [0, 0]
      src_result[0] = src_point[0] * cs - src_point[1] * sn
      src_result[1] = src_point[0] * sn + src_point[1] * cs
  
      return src_result
  
  
  def get_3rd_point(a, b):
      direct = a - b
      return b + np.array([-direct[1], direct[0]], dtype=np.float32)
  
  
  def load_checkpoint(model, trained_model_weights_filename):
      checkpoint = torch.load(trained_model_weights_filename, map_location="cpu")
      model.load_state_dict(checkpoint["model"])
      return model
  
  
  def load_model(arch, model_weights, device):
  
      if model_weights is None:
          if arch == "mobilenet_v3_small":
              model_weights = "models/mobilenet_v3_pretrained.pth"
              arch = "mobilenetv3small"
          elif arch == "res_18":
              model_weights = "models/res18_pretrained.pth"
          elif arch == "dla_34":
              model_weights = "models/dla_34_pretrained.pth"
  
      heads = {"hm": 1} if arch != "dla_34" else {"hm": 1, "wh": 2}
  
      model = create_base(arch, heads=heads, head_conv=256).to(device)
      model = load_checkpoint(model, model_weights)
      for param in model.parameters():
          param.requires_grad = False
      model.eval()
  
      return model
  
  
  def _calculate_euclidean_similarity(distances, zero_distance):
      """Calculates the euclidean distance between two sets of detections, and then converts this into a similarity
      measure with values between 0 and 1 using the following formula: sim = max(0, 1 - dist/zero_distance).
      The default zero_distance of 2.0, corresponds to the default used in MOT15_3D, such that a 0.5 similarity
      threshold corresponds to a 1m distance threshold for TPs.
      """
      sim = np.maximum(0, 1 - distances / zero_distance)
      return sim
  