text: import torch
  from torch.jit.annotations import Tuple
  from torch import Tensor
  import torchvision
  
  
  def nms(boxes, scores, iou_threshold):
      # type: (Tensor, Tensor, float)
      """
      Performs non-maximum suppression (NMS) on the boxes according
      to their intersection-over-union (IoU).
  
      NMS iteratively removes lower scoring boxes which have an
      IoU greater than iou_threshold with another (higher scoring)
      box.
  
      If multiple boxes have the exact same score and satisfy the IoU
      criterion with respect to a reference box, the selected box is
      not guaranteed to be the same between CPU and GPU. This is similar
      to the behavior of argsort in PyTorch when repeated values are present.
  
      Parameters
      ----------
      boxes : Tensor[N, 4])
          boxes to perform NMS on. They
          are expected to be in (x1, y1, x2, y2) format
      scores : Tensor[N]
          scores for each one of the boxes
      iou_threshold : float
          discards all overlapping
          boxes with IoU > iou_threshold
  
      Returns
      -------
      keep : Tensor
          int64 tensor with the indices
          of the elements that have been kept
          by NMS, sorted in decreasing order of scores
      """
      return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
  
  
  def batched_nms(boxes, scores, idxs, iou_threshold):
      # type: (Tensor, Tensor, Tensor, float)
      """
      Performs non-maximum suppression in a batched fashion.
  
      Each index value correspond to a category, and NMS
      will not be applied between elements of different categories.
  
      Parameters
      ----------
      boxes : Tensor[N, 4]
          boxes where NMS will be performed. They
          are expected to be in (x1, y1, x2, y2) format
      scores : Tensor[N]
          scores for each one of the boxes
      idxs : Tensor[N]
          indices of the categories for each one of the boxes.
      iou_threshold : float
          discards all overlapping boxes
          with IoU > iou_threshold
  
      Returns
      -------
      keep : Tensor
          int64 tensor with the indices of
          the elements that have been kept by NMS, sorted
          in decreasing order of scores
      """
      if boxes.numel() == 0:
          return torch.empty((0,), dtype=torch.int64, device=boxes.device)
      # 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 = boxes.max()
      offsets = idxs.to(boxes) * (max_coordinate + 1)
      boxes_for_nms = boxes + offsets[:, None]
      keep = nms(boxes_for_nms, scores, iou_threshold)
      return keep
  
  
  def remove_small_boxes(boxes, min_size):
      # type: (Tensor, float)
      """
      Remove boxes which contains at least one side smaller than min_size.
  
      Arguments:
          boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format
          min_size (float): minimum size
  
      Returns:
          keep (Tensor[K]): indices of the boxes that have both sides
              larger than min_size
      """
      ws, hs = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]
      keep = (ws >= min_size) & (hs >= min_size)
      keep = keep.nonzero().squeeze(1)
      return keep
  
  
  def clip_boxes_to_image(boxes, size):
      # type: (Tensor, Tuple[int, int])
      """
      Clip boxes so that they lie inside an image of size `size`.
  
      Arguments:
          boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format
          size (Tuple[height, width]): size of the image
  
      Returns:
          clipped_boxes (Tensor[N, 4])
      """
      dim = boxes.dim()
      boxes_x = boxes[..., 0::2]
      boxes_y = boxes[..., 1::2]
      height, width = size
  
      if torchvision._is_tracing():
          boxes_x = torch.max(boxes_x, torch.tensor(0, dtype=boxes.dtype, device=boxes.device))
          boxes_x = torch.min(boxes_x, torch.tensor(width, dtype=boxes.dtype, device=boxes.device))
          boxes_y = torch.max(boxes_y, torch.tensor(0, dtype=boxes.dtype, device=boxes.device))
          boxes_y = torch.min(boxes_y, torch.tensor(height, dtype=boxes.dtype, device=boxes.device))
      else:
          boxes_x = boxes_x.clamp(min=0, max=width)
          boxes_y = boxes_y.clamp(min=0, max=height)
  
      clipped_boxes = torch.stack((boxes_x, boxes_y), dim=dim)
      return clipped_boxes.reshape(boxes.shape)
  
  
  def box_area(boxes):
      """
      Computes the area of a set of bounding boxes, which are specified by its
      (x1, y1, x2, y2) coordinates.
  
      Arguments:
          boxes (Tensor[N, 4]): boxes for which the area will be computed. They
              are expected to be in (x1, y1, x2, y2) format
  
      Returns:
          area (Tensor[N]): area for each box
      """
      return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
  
  
  # implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
  # with slight modifications
  def box_iou(boxes1, boxes2):
      """
      Return intersection-over-union (Jaccard index) of boxes.
  
      Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
  
      Arguments:
          boxes1 (Tensor[N, 4])
          boxes2 (Tensor[M, 4])
  
      Returns:
          iou (Tensor[N, M]): the NxM matrix containing the pairwise
              IoU values for every element in boxes1 and boxes2
      """
      area1 = box_area(boxes1)
      area2 = box_area(boxes2)
  
      lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
      rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
  
      wh = (rb - lt).clamp(min=0)  # [N,M,2]
      inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
  
      iou = inter / (area1[:, None] + area2 - inter)
      return iou
  