text: # ---------------------------------------------------------------
  # bbox_nms_with_reppoints.py
  # Set-up time: 2020/4/25 上午11:13
  # Copyright (c) 2020 ICT
  # Licensed under The MIT License [see LICENSE for details]
  # Written by Kenneth-Wong (Wenbin-Wang) @ VIPL.ICT
  # Contact: wenbin.wang@vipl.ict.ac.cn [OR] nkwangwenbin@gmail.com
  # ---------------------------------------------------------------
  
  import torch
  
  from mmdet.ops.nms import nms_wrapper
  
  
  def multiclass_nms_with_reppoints(multi_bboxes,
                                    multi_reppoints,
                                    multi_scores,
                                    score_thr,
                                    nms_cfg,
                                    max_num=-1,
                                    score_factors=None):
      """NMS for multi-class bboxes with reppoints.
  
      Args:
          multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
          multi_scores (Tensor): shape (n, #class), where the 0th column
              contains scores of the background class, but this will be ignored.
          score_thr (float): bbox threshold, bboxes with scores lower than it
              will not be considered.
          nms_thr (float): NMS IoU threshold
          max_num (int): if there are more than max_num bboxes after NMS,
              only top max_num will be kept.
          score_factors (Tensor): The factors multiplied to scores before
              applying NMS
  
      Returns:
          tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels
              are 0-based.
      """
      num_classes = multi_scores.size(1) - 1
      # exclude background category
      if multi_bboxes.shape[1] > 4:
          bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)[:, 1:]
          reppoints = multi_reppoints.view(multi_scores.size(0), -1, multi_reppoints.shape[1])[:, 1:]
      else:
          bboxes = multi_bboxes[:, None].expand(-1, num_classes, 4)
          reppoints = multi_reppoints[:, None].expand(-1, num_classes, multi_reppoints.shape[1])
      scores = multi_scores[:, 1:]
  
      # filter out boxes with low scores
      valid_mask = scores > score_thr
      bboxes = bboxes[valid_mask]
      if score_factors is not None:
          scores = scores * score_factors[:, None]
      scores = scores[valid_mask]
      reppoints = reppoints[valid_mask]
      labels = valid_mask.nonzero()[:, 1]
  
      if bboxes.numel() == 0:
          bboxes = multi_bboxes.new_zeros((0, 5))
          reppoints = multi_reppoints.new_zeros((0, multi_reppoints.shape[1]))
          labels = multi_bboxes.new_zeros((0,), dtype=torch.long)
          return bboxes, reppoints, labels
  
      # Modified from https://github.com/pytorch/vision/blob
      # /505cd6957711af790211896d32b40291bea1bc21/torchvision/ops/boxes.py#L39.
      # strategy: in order to perform NMS independently per class.
      # we add an offset to all the boxes. The offset is dependent
      # only on the class idx, and is large enough so that boxes
      # from different classes do not overlap
      max_coordinate = bboxes.max()
      offsets = labels.to(bboxes) * (max_coordinate + 1)
      bboxes_for_nms = bboxes + offsets[:, None]
      nms_cfg_ = nms_cfg.copy()
      nms_type = nms_cfg_.pop('type', 'nms')
      nms_op = getattr(nms_wrapper, nms_type)
      dets, keep = nms_op(
          torch.cat([bboxes_for_nms, scores[:, None]], 1), **nms_cfg_)
      bboxes = bboxes[keep]
      reppoints = reppoints[keep]
      scores = dets[:, -1]  # soft_nms will modify scores
      labels = labels[keep]
  
      if keep.size(0) > max_num:
          _, inds = scores.sort(descending=True)
          inds = inds[:max_num]
          bboxes = bboxes[inds]
          reppoints = reppoints[inds]
          scores = scores[inds]
          labels = labels[inds]
  
      return torch.cat([bboxes, scores[:, None]], 1), reppoints, labels
  