code: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  import torch
  
  from maskrcnn_benchmark.modeling.box_coder import BoxCoder
  from maskrcnn_benchmark.modeling.rbox_coder import RBoxCoder
  from maskrcnn_benchmark.structures.bounding_box import BoxList, RBoxList
  from maskrcnn_benchmark.structures.rboxlist_ops import cat_boxlist
  from maskrcnn_benchmark.structures.rboxlist_ops import boxlist_nms #
  from maskrcnn_benchmark.structures.rboxlist_ops import remove_small_boxes
  
  from ..utils import cat
  
  
  class RPNPostProcessor(torch.nn.Module):
      """
      Performs post-processing on the outputs of the RPN boxes, before feeding the
      proposals to the heads
      """
  
      def __init__(
          self,
          pre_nms_top_n,
          post_nms_top_n,
          nms_thresh,
          min_size,
          box_coder=None,
          fpn_post_nms_top_n=None,
      ):
          """
          Arguments:
              pre_nms_top_n (int)
              post_nms_top_n (int)
              nms_thresh (float)
              min_size (int)
              box_coder (BoxCoder)
              fpn_post_nms_top_n (int)
          """
          super(RPNPostProcessor, self).__init__()
          self.pre_nms_top_n = pre_nms_top_n
          self.post_nms_top_n = post_nms_top_n
          self.nms_thresh = nms_thresh
          self.min_size = min_size
  
          if box_coder is None:
              box_coder = RBoxCoder(weights=(1.0, 1.0, 1.0, 1.0, 1.0))
          self.box_coder = box_coder
  
          if fpn_post_nms_top_n is None:
              fpn_post_nms_top_n = post_nms_top_n
          self.fpn_post_nms_top_n = fpn_post_nms_top_n
  
      def add_gt_proposals(self, proposals, targets):
          """
          Arguments:
              proposals: list[BoxList]
              targets: list[BoxList]
          """
          # Get the device we're operating on
          device = proposals[0].bbox.device
  
          gt_boxes = [target.copy_with_fields([]) for target in targets]
  
          # later cat of bbox requires all fields to be present for all bbox
          # so we need to add a dummy for objectness that's missing
          for gt_box in gt_boxes:
              gt_box.add_field("objectness", torch.ones(len(gt_box), device=device))
  
          proposals = [
              cat_boxlist((proposal, gt_box))
              for proposal, gt_box in zip(proposals, gt_boxes)
          ]
          # print('rrpn_proposal:', proposals[0].bbox.size(), proposals[0].bbox[:, 2:4])
          return proposals
  
      # proposal_target_layer
      def forward_for_single_feature_map(self, anchors, objectness, box_regression):
          """
          Arguments:
              anchors: list[BoxList]
              objectness: tensor of size N, A, H, W
              box_regression: tensor of size N, A * 5, H, W
          """
          device = objectness.device
          N, A, H, W = objectness.shape
  
          # put in the same format as anchors
          objectness = objectness.permute(0, 2, 3, 1).reshape(N, -1)
          objectness = objectness.sigmoid()
          box_regression = box_regression.view(N, -1, 5, H, W).permute(0, 3, 4, 1, 2)
          box_regression = box_regression.reshape(N, -1, 5)
  
          num_anchors = A * H * W
  
          pre_nms_top_n = min(self.pre_nms_top_n, num_anchors)
          objectness, topk_idx = objectness.topk(pre_nms_top_n, dim=1, sorted=True)
  
          batch_idx = torch.arange(N, device=device)[:, None]
          box_regression = box_regression[batch_idx, topk_idx]
  
          image_shapes = [box.size for box in anchors]
          concat_anchors = torch.cat([a.bbox for a in anchors], dim=0)
          concat_anchors = concat_anchors.reshape(N, -1, 5)[batch_idx, topk_idx]
  
          # print('concat_anchors:', concat_anchors.size(), concat_anchors[:, 2:4])
  
          proposals = self.box_coder.decode(
              box_regression.view(-1, 5), concat_anchors.view(-1, 5)
          )
  
          proposals = proposals.view(N, -1, 5)
          # print('outsider:', proposals.size(), proposals[:, 2:4], 'box_regression:', box_regression)
  
          #-------
          result = []
          for proposal, score, im_shape in zip(proposals, objectness, image_shapes):
              boxlist = RBoxList(proposal, im_shape, mode="xywha")
  
              # print('before nms:', boxlist.bbox.size(), boxlist.bbox[:, 2:4])
  
              boxlist.add_field("objectness", score)
              # boxlist = boxlist.clip_to_image(remove_empty=False)
              boxlist = remove_small_boxes(boxlist, self.min_size)
              boxlist = boxlist_nms(
                  boxlist,
                  self.nms_thresh,
                  max_proposals=self.post_nms_top_n,
                  score_field="objectness",
              )
  
              # print('after nms:', boxlist.bbox.size(), boxlist.bbox[:, 2:4])
  
              result.append(boxlist)
          return result
  
      def forward(self, anchors, objectness, box_regression, targets=None):
          """
          Arguments:
              anchors: list[list[BoxList]]
              objectness: list[tensor]
              box_regression: list[tensor]
  
          Returns:
              boxlists (list[BoxList]): the post-processed anchors, after
                  applying box decoding and NMS
          """
          sampled_boxes = []
          num_levels = len(objectness)
          anchors = list(zip(*anchors))
          for a, o, b in zip(anchors, objectness, box_regression):
              sampled_boxes.append(self.forward_for_single_feature_map(a, o, b))
  
          boxlists = list(zip(*sampled_boxes))
          boxlists = [cat_boxlist(boxlist) for boxlist in boxlists]
  
          if num_levels > 1:
              boxlists = self.select_over_all_levels(boxlists)
  
          # append ground-truth bboxes to proposals
          if self.training and targets is not None:
              boxlists = self.add_gt_proposals(boxlists, targets)
  
          return boxlists
  
      def select_over_all_levels(self, boxlists):
          num_images = len(boxlists)
          # different behavior during training and during testing:
          # during training, post_nms_top_n is over *all* the proposals combined, while
          # during testing, it is over the proposals for each image
          # TODO resolve this difference and make it consistent. It should be per image,
          # and not per batch
          if self.training:
              objectness = torch.cat(
                  [boxlist.get_field("objectness") for boxlist in boxlists], dim=0
              )
              box_sizes = [len(boxlist) for boxlist in boxlists]
              post_nms_top_n = min(self.fpn_post_nms_top_n, len(objectness))
              _, inds_sorted = torch.topk(objectness, post_nms_top_n, dim=0, sorted=True)
              inds_mask = torch.zeros_like(objectness, dtype=torch.uint8)
              inds_mask[inds_sorted] = 1
              inds_mask = inds_mask.split(box_sizes)
              for i in range(num_images):
                  boxlists[i] = boxlists[i][inds_mask[i]]
          else:
              for i in range(num_images):
                  objectness = boxlists[i].get_field("objectness")
                  post_nms_top_n = min(self.fpn_post_nms_top_n, len(objectness))
                  _, inds_sorted = torch.topk(
                      objectness, post_nms_top_n, dim=0, sorted=True
                  )
                  boxlists[i] = boxlists[i][inds_sorted]
          return boxlists
  
  
  def make_rpn_postprocessor(config, rpn_box_coder, is_train):
      fpn_post_nms_top_n = config.MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN
      if not is_train:
          fpn_post_nms_top_n = config.MODEL.RPN.FPN_POST_NMS_TOP_N_TEST
  
      pre_nms_top_n = config.MODEL.RPN.PRE_NMS_TOP_N_TRAIN
      post_nms_top_n = config.MODEL.RPN.POST_NMS_TOP_N_TRAIN
      if not is_train:
          pre_nms_top_n = config.MODEL.RPN.PRE_NMS_TOP_N_TEST
          post_nms_top_n = config.MODEL.RPN.POST_NMS_TOP_N_TEST
      nms_thresh = config.MODEL.RPN.NMS_THRESH
      min_size = config.MODEL.RPN.MIN_SIZE
      box_selector = RPNPostProcessor(
          pre_nms_top_n=pre_nms_top_n,
          post_nms_top_n=post_nms_top_n,
          nms_thresh=nms_thresh,
          min_size=min_size,
          box_coder=rpn_box_coder,
          fpn_post_nms_top_n=fpn_post_nms_top_n,
      )
      return box_selector
  
apis: torch.topk torch.arange torch.cat torch.zeros_like