text: # -*- coding: utf-8 -*- #
  # --------------------------------------------------------
  # @Author: Lintao
  # @Date: 3/31/21 3:35 PM
  # @Version	1.0
  # --------------------------------------------------------
  import time
  
  import numpy as np
  import torch
  import torchvision
  
  
  def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
      """Performs Non-Maximum Suppression (NMS) on inference results
  
      Returns:
           detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
      """
  
      nc = prediction.shape[2] - 5  # number of classes
      xc = prediction[..., 4] > conf_thres  # candidates
  
      # Settings
      min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
      max_det = 300  # maximum number of detections per image
      time_limit = 10.0  # seconds to quit after
      redundant = True  # require redundant detections
      multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)
      merge = False  # use merge-NMS
  
      t = time.time()
      output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
      for xi, x in enumerate(prediction):  # image index, image inference
          # Apply constraints
          # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
          x = x[xc[xi]]  # confidence
  
          # Cat apriori labels if autolabelling
          if labels and len(labels[xi]):
              l = labels[xi]
              v = torch.zeros((len(l), nc + 5), device=x.device)
              v[:, :4] = l[:, 1:5]  # box
              v[:, 4] = 1.0  # conf
              v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
              x = torch.cat((x, v), 0)
  
          # If none remain process next image
          if not x.shape[0]:
              continue
  
          # Compute conf
          x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
  
          # Box (center x, center y, width, height) to (x1, y1, x2, y2)
          box = xywh2xyxy(x[:, :4])
  
          # Detections matrix nx6 (xyxy, conf, cls)
          if multi_label:
              i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
              x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
          else:  # best class only
              conf, j = x[:, 5:].max(1, keepdim=True)
              x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
  
          # Filter by class
          if classes is not None:
              x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
  
          # Apply finite constraint
          # if not torch.isfinite(x).all():
          #     x = x[torch.isfinite(x).all(1)]
  
          # If none remain process next image
          n = x.shape[0]  # number of boxes
          if not n:
              continue
  
          # Sort by confidence
          # x = x[x[:, 4].argsort(descending=True)]
  
          # Batched NMS
          c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
          boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
          i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
          if i.shape[0] > max_det:  # limit detections
              i = i[:max_det]
          if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
              # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
              iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
              weights = iou * scores[None]  # box weights
              x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
              if redundant:
                  i = i[iou.sum(1) > 1]  # require redundancy
  
          output[xi] = x[i]
          if (time.time() - t) > time_limit:
              break  # time limit exceeded
  
      return output
  
  
  def xywh2xyxy(x):
      # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
      y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
      y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
      y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
      y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
      y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
      return y
  
  
  def IoU(bbox, gt):
      """
      :param bbox: (n, 4)
      :param gt: (m, 4)
      :return: (n, m)
      numpy 广播机制 从后向前对齐。 维度为1 的可以重复等价为任意维度
      eg: (4,3,2)   (3,2)  (3,2)会扩充为(4,3,2)
          (4,1,2)   (3,2) (4,1,2) 扩充为(4, 3, 2)  (3, 2)扩充为(4, 3,2) 扩充的方法为重复
      广播会在numpy的函数 如sum, maximun等函数中进行
      pytorch同理。
      扩充维度的方法：
      eg: a  a.shape: (3,2)  a[:, None, :] a.shape: (3, 1, 2) None 对应的维度相当于newaxis
      """
      lt = np.maximum(bbox[:, None, :2], gt[:, :2])  # left_top (x, y)
      rb = np.minimum(bbox[:, None, 2:], gt[:, 2:])  # right_bottom (x, y)
      wh = np.maximum(rb - lt + 1, 0)  # inter_area (w, h)
      inter_areas = wh[:, :, 0] * wh[:, :, 1]  # shape: (n, m)
      box_areas = (bbox[:, 2] - bbox[:, 0] + 1) * (bbox[:, 3] - bbox[:, 1] + 1)
      gt_areas = (gt[:, 2] - gt[:, 0] + 1) * (gt[:, 3] - gt[:, 1] + 1)
      IoU = inter_areas / (box_areas[:, None] + gt_areas - inter_areas)
      return IoU
  
  
  def box_iou(box1, box2):
      # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
      """
      Return intersection-over-union (Jaccard index) of boxes.
      Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
      Arguments:
          box1 (Tensor[N, 4])
          box2 (Tensor[M, 4])
      Returns:
          iou (Tensor[N, M]): the NxM matrix containing the pairwise
              IoU values for every element in boxes1 and boxes2
      """
  
      def box_area(box):
          # box = 4xn
          return (box[2] - box[0]) * (box[3] - box[1])
  
      area1 = box_area(box1.T)
      area2 = box_area(box2.T)
  
      # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
      inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
      return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)
  