text: import numpy as np
  import torch
  from torchvision.ops import nms, boxes
  
  def yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image):
      #-----------------------------------------------------------------#
      #   把y轴放前面是因为方便预测框和图像的宽高进行相乘
      #-----------------------------------------------------------------#
      box_yx = box_xy[..., ::-1]
      box_hw = box_wh[..., ::-1]
      input_shape = np.array(input_shape)
      image_shape = np.array(image_shape)
  
      if letterbox_image:
          #-----------------------------------------------------------------#
          #   这里求出来的offset是图像有效区域相对于图像左上角的偏移情况
          #   new_shape指的是宽高缩放情况
          #-----------------------------------------------------------------#
          new_shape = np.round(image_shape * np.min(input_shape/image_shape))
          offset  = (input_shape - new_shape)/2./input_shape
          scale   = input_shape/new_shape
  
          box_yx  = (box_yx - offset) * scale
          box_hw *= scale
  
      box_mins    = box_yx - (box_hw / 2.)
      box_maxes   = box_yx + (box_hw / 2.)
      boxes  = np.concatenate([box_mins[..., 0:1], box_mins[..., 1:2], box_maxes[..., 0:1], box_maxes[..., 1:2]], axis=-1)
      boxes *= np.concatenate([image_shape, image_shape], axis=-1)
      return boxes
  
  def decode_outputs(outputs, input_shape):
      grids   = []
      strides = []
      hw      = [x.shape[-2:] for x in outputs]
      #---------------------------------------------------#
      #   outputs输入前代表每个特征层的预测结果
      #   batch_size, 4 + 1 + num_classes, 80, 80 => batch_size, 4 + 1 + num_classes, 6400
      #   batch_size, 5 + num_classes, 40, 40
      #   batch_size, 5 + num_classes, 20, 20
      #   batch_size, 4 + 1 + num_classes, 6400 + 1600 + 400 -> batch_size, 4 + 1 + num_classes, 8400
      #   堆叠后为batch_size, 8400, 5 + num_classes
      #---------------------------------------------------#
      outputs = torch.cat([x.flatten(start_dim=2) for x in outputs], dim=2).permute(0, 2, 1)
      #---------------------------------------------------#
      #   获得每一个特征点属于每一个种类的概率
      #---------------------------------------------------#
      outputs[:, :, 4:] = torch.sigmoid(outputs[:, :, 4:])
      for h, w in hw:
          #---------------------------#
          #   根据特征层的高宽生成网格点
          #---------------------------#   
          grid_y, grid_x  = torch.meshgrid([torch.arange(h), torch.arange(w)])
          #---------------------------#
          #   1, 6400, 2
          #   1, 1600, 2
          #   1, 400, 2
          #---------------------------#   
          grid            = torch.stack((grid_x, grid_y), 2).view(1, -1, 2)
          shape           = grid.shape[:2]
  
          grids.append(grid)
          strides.append(torch.full((shape[0], shape[1], 1), input_shape[0] / h))
      #---------------------------#
      #   将网格点堆叠到一起
      #   1, 6400, 2
      #   1, 1600, 2
      #   1, 400, 2
      #
      #   1, 8400, 2
      #---------------------------#
      grids               = torch.cat(grids, dim=1).type(outputs.type())
      strides             = torch.cat(strides, dim=1).type(outputs.type())
      #------------------------#
      #   根据网格点进行解码
      #------------------------#
      outputs[..., :2]    = (outputs[..., :2] + grids) * strides
      outputs[..., 2:4]   = torch.exp(outputs[..., 2:4]) * strides
      #-----------------#
      #   归一化
      #-----------------#
      outputs[..., [0,2]] = outputs[..., [0,2]] / input_shape[1]
      outputs[..., [1,3]] = outputs[..., [1,3]] / input_shape[0]
      return outputs
  
  def non_max_suppression(prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):
      #----------------------------------------------------------#
      #   将预测结果的格式转换成左上角右下角的格式。
      #   prediction  [batch_size, num_anchors, 85]
      #----------------------------------------------------------#
      box_corner          = prediction.new(prediction.shape)
      box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
      box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
      box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
      box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
      prediction[:, :, :4] = box_corner[:, :, :4]
      
      output = [None for _ in range(len(prediction))]
      #----------------------------------------------------------#
      #   对输入图片进行循环，一般只会进行一次
      #----------------------------------------------------------#
      for i, image_pred in enumerate(prediction):
          #----------------------------------------------------------#
          #   对种类预测部分取max。
          #   class_conf  [num_anchors, 1]    种类置信度
          #   class_pred  [num_anchors, 1]    种类
          #----------------------------------------------------------#
          class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True)
  
          #----------------------------------------------------------#
          #   利用置信度进行第一轮筛选
          #----------------------------------------------------------#
          conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()
  
          if not image_pred.size(0):
              continue
          #-------------------------------------------------------------------------#
          #   detections  [num_anchors, 7]
          #   7的内容为：x1, y1, x2, y2, obj_conf, class_conf, class_pred
          #-------------------------------------------------------------------------#
          detections = torch.cat((image_pred[:, :5], class_conf, class_pred.float()), 1)
          detections = detections[conf_mask]
          
          nms_out_index = boxes.batched_nms(
              detections[:, :4],
              detections[:, 4] * detections[:, 5],
              detections[:, 6],
              nms_thres,
          )
  
          output[i]   = detections[nms_out_index]
  
          # #------------------------------------------#
          # #   获得预测结果中包含的所有种类
          # #------------------------------------------#
          # unique_labels = detections[:, -1].cpu().unique()
  
          # if prediction.is_cuda:
          #     unique_labels = unique_labels.cuda()
          #     detections = detections.cuda()
  
          # for c in unique_labels:
          #     #------------------------------------------#
          #     #   获得某一类得分筛选后全部的预测结果
          #     #------------------------------------------#
          #     detections_class = detections[detections[:, -1] == c]
  
          #     #------------------------------------------#
          #     #   使用官方自带的非极大抑制会速度更快一些！
          #     #------------------------------------------#
          #     keep = nms(
          #         detections_class[:, :4],
          #         detections_class[:, 4] * detections_class[:, 5],
          #         nms_thres
          #     )
          #     max_detections = detections_class[keep]
              
          #     # # 按照存在物体的置信度排序
          #     # _, conf_sort_index = torch.sort(detections_class[:, 4]*detections_class[:, 5], descending=True)
          #     # detections_class = detections_class[conf_sort_index]
          #     # # 进行非极大抑制
          #     # max_detections = []
          #     # while detections_class.size(0):
          #     #     # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
          #     #     max_detections.append(detections_class[0].unsqueeze(0))
          #     #     if len(detections_class) == 1:
          #     #         break
          #     #     ious = bbox_iou(max_detections[-1], detections_class[1:])
          #     #     detections_class = detections_class[1:][ious < nms_thres]
          #     # # 堆叠
          #     # max_detections = torch.cat(max_detections).data
              
          #     # Add max detections to outputs
          #     output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))
          
          if output[i] is not None:
              output[i]           = output[i].cpu().numpy()
              box_xy, box_wh      = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]
              output[i][:, :4]    = yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)
      return output
  