text: import torch
  from mmcv.ops.nms import batched_nms
  
  from mmdet.core.bbox.iou_calculators import bbox_overlaps
  from mmcv.ops.nms import nms, soft_nms
  
  
  def multiclass_nms(multi_bboxes,
                     multi_scores,
                     score_thr,
                     nms_cfg,
                     max_num=-1,
                     score_factors=None):
      """NMS for multi-class bboxes.
  
      Args:
          multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
          multi_scores (Tensor): shape (n, #class), where the last 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)
      else:
          bboxes = multi_bboxes[:, None].expand(
              multi_scores.size(0), num_classes, 4)
  
      scores = multi_scores[:, :-1]
      if score_factors is not None:
          scores = scores * score_factors[:, None]
  
      labels = torch.arange(num_classes, dtype=torch.long)
      labels = labels.view(1, -1).expand_as(scores)
  
      bboxes = bboxes.reshape(-1, 4)
      scores = scores.reshape(-1)
      labels = labels.reshape(-1)
  
      # remove low scoring boxes
      valid_mask = scores > score_thr
      inds = valid_mask.nonzero(as_tuple=False).squeeze(1)
      bboxes, scores, labels = bboxes[inds], scores[inds], labels[inds]
      if inds.numel() == 0:
          if torch.onnx.is_in_onnx_export():
              raise RuntimeError('[ONNX Error] Can not record NMS '
                                 'as it has not been executed this time')
          return bboxes, labels
  
      # TODO: add size check before feed into batched_nms
      dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
  
      if max_num > 0:
          dets = dets[:max_num]
          keep = keep[:max_num]
  
      return dets, labels[keep]
  
  
  def multiclass_nms_with_mask(multi_bboxes,
                               multi_scores,
                               multi_masks,
                               score_thr,
                               nms_cfg,
                               max_num=-1,
                               score_factors=None):
      """NMS for multi-class bboxes.
  
      Args:
          multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
          multi_scores (Tensor): shape (n, #class)
          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.shape[1]
      bboxes, labels, masks = [], [], []
      nms_cfg_ = nms_cfg.copy()
      nms_type = nms_cfg_.pop('type', 'nms')
      iou_thr = nms_cfg_.pop('iou_thr', 0.5)
      if nms_type == 'nms':
          nms_op = nms
      else:
          nms_op = soft_nms
      for i in range(1, num_classes):
          cls_inds = multi_scores[:, i] > score_thr
          if not cls_inds.any():
              continue
          # get bboxes and scores of this class
          if multi_bboxes.shape[1] == 4:
              _bboxes = multi_bboxes[cls_inds, :]
              _masks = multi_masks[cls_inds, :]
          else:
              _bboxes = multi_bboxes[cls_inds, i * 4:(i + 1) * 4]
          _scores = multi_scores[cls_inds, i]
          if score_factors is not None:
              _scores *= score_factors[cls_inds]
          cls_dets, index = nms_op(_bboxes, _scores, iou_thr)
          cls_masks = _masks[index]
          cls_labels = multi_bboxes.new_full((cls_dets.shape[0], ),
                                             i - 1,
                                             dtype=torch.long)
          bboxes.append(cls_dets)
          labels.append(cls_labels)
          masks.append(cls_masks)
      if bboxes:
          bboxes = torch.cat(bboxes)
          labels = torch.cat(labels)
          masks = torch.cat(masks)
          if bboxes.shape[0] > max_num:
              _, inds = bboxes[:, -1].sort(descending=True)
              inds = inds[:max_num]
              bboxes = bboxes[inds]
              labels = labels[inds]
              masks = masks[inds]
      else:
          bboxes = multi_bboxes.new_zeros((0, 5))
          labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)
          masks = multi_bboxes.new_zeros((0, 2, 36))
  
      return bboxes, labels, masks
  
  
  def fast_nms(multi_bboxes,
               multi_scores,
               multi_coeffs,
               score_thr,
               iou_thr,
               top_k,
               max_num=-1):
      """Fast NMS in `YOLACT <https://arxiv.org/abs/1904.02689>`_.
  
      Fast NMS allows already-removed detections to suppress other detections so
      that every instance can be decided to be kept or discarded in parallel,
      which is not possible in traditional NMS. This relaxation allows us to
      implement Fast NMS entirely in standard GPU-accelerated matrix operations.
  
      Args:
          multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
          multi_scores (Tensor): shape (n, #class+1), where the last column
              contains scores of the background class, but this will be ignored.
          multi_coeffs (Tensor): shape (n, #class*coeffs_dim).
          score_thr (float): bbox threshold, bboxes with scores lower than it
              will not be considered.
          iou_thr (float): IoU threshold to be considered as conflicted.
          top_k (int): if there are more than top_k bboxes before NMS,
              only top top_k will be kept.
          max_num (int): if there are more than max_num bboxes after NMS,
              only top max_num will be kept. If -1, keep all the bboxes.
              Default: -1.
  
      Returns:
          tuple: (bboxes, labels, coefficients), tensors of shape (k, 5), (k, 1),
              and (k, coeffs_dim). Labels are 0-based.
      """
  
      scores = multi_scores[:, :-1].t()  # [#class, n]
      scores, idx = scores.sort(1, descending=True)
  
      idx = idx[:, :top_k].contiguous()
      scores = scores[:, :top_k]  # [#class, topk]
      num_classes, num_dets = idx.size()
      boxes = multi_bboxes[idx.view(-1), :].view(num_classes, num_dets, 4)
      coeffs = multi_coeffs[idx.view(-1), :].view(num_classes, num_dets, -1)
  
      iou = bbox_overlaps(boxes, boxes)  # [#class, topk, topk]
      iou.triu_(diagonal=1)
      iou_max, _ = iou.max(dim=1)
  
      # Now just filter out the ones higher than the threshold
      keep = iou_max <= iou_thr
  
      # Second thresholding introduces 0.2 mAP gain at negligible time cost
      keep *= scores > score_thr
  
      # Assign each kept detection to its corresponding class
      classes = torch.arange(
          num_classes, device=boxes.device)[:, None].expand_as(keep)
      classes = classes[keep]
  
      boxes = boxes[keep]
      coeffs = coeffs[keep]
      scores = scores[keep]
  
      # Only keep the top max_num highest scores across all classes
      scores, idx = scores.sort(0, descending=True)
      if max_num > 0:
          idx = idx[:max_num]
          scores = scores[:max_num]
  
      classes = classes[idx]
      boxes = boxes[idx]
      coeffs = coeffs[idx]
  
      cls_dets = torch.cat([boxes, scores[:, None]], dim=1)
      return cls_dets, classes, coeffs
  