text: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  import torch
  
  from maskrcnn_benchmark.modeling.box_coder import BoxCoder
  from maskrcnn_benchmark.structures.bounding_box import BoxList
  from maskrcnn_benchmark.structures.boxlist_ops import cat_boxlist
  from maskrcnn_benchmark.structures.boxlist_ops import boxlist_nms
  from maskrcnn_benchmark.structures.boxlist_ops import remove_small_boxes
  
  from ..utils import cat
  from .utils import permute_and_flatten
  
  import scipy.ndimage
  import numpy as np
  
  
  def extract_bboxes(mask):
      """Compute bounding boxes from masks.
      mask: [height, width, num_instances]. Mask pixels are either 1 or 0.
  
      Returns: bbox array [num_instances, (y1, x1, y2, x2)].
      """
      if len(mask.shape) == 2:
          mask = mask[:, :, np.newaxis]
      boxes = np.zeros([mask.shape[-1], 4], dtype=np.float32)
      for i in range(mask.shape[-1]):
          m = mask[:, :, i]
          # Bounding box.
          horizontal_indicies = np.where(np.any(m, axis=0))[0]
          vertical_indicies = np.where(np.any(m, axis=1))[0]
          if horizontal_indicies.shape[0]:
              x1, x2 = horizontal_indicies[[0, -1]]
              y1, y2 = vertical_indicies[[0, -1]]
              # x2 and y2 should not be part of the box. Increment by 1.
              x2 += 1
              y2 += 1
              w = x2 - x1
              h = y2 - y1
              ctr_x = (x1 + 0.5 * w) * 4.0
              ctr_y = (y1 + 0.5 * h) * 4.0
              w = w * 4.0 / 0.3
              h = h * 4.0 / 0.3
              x1 = ctr_x - w * 0.5
              y1 = ctr_y - h * 0.5
              x2 = ctr_x + w * 0.5
              y2 = ctr_y + h * 0.5
          else:
              assert False, (horizontal_indicies, vertical_indicies)
              # No mask for this instance. Might happen due to
              # resizing or cropping. Set bbox to zeros
              x1, x2, y1, y2 = 0, 0, 0, 0
          boxes[i] = np.array([x1, y1, x2, y2])
      return boxes.astype(np.float32)
  
  
  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,
          fpn_post_nms_per_batch=True,
          pred_targets=False,
          pred_targets_as_true=False,
      ):
          """
          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 = BoxCoder(weights=(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
          self.fpn_post_nms_per_batch = fpn_post_nms_per_batch
  
          self.pred_targets = pred_targets
          self.pred_targets_as_true = pred_targets_as_true
  
      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)
          ]
  
          return proposals
  
      def add_pred_proposals(self, proposals, pred_targets):
          """
          Arguments:
              proposals: list[BoxList]
              targets: list[BoxList]
          """
          # Get the device we're operating on
          device = proposals[0].bbox.device
  
          for idx, proposal in enumerate(proposals):
              if pred_targets and pred_targets[idx]:
                  gt_box = pred_targets[idx].copy_with_fields([])
                  gt_box.add_field("objectness", torch.ones(len(gt_box), device=device) * 2)
                  proposals[idx] = cat_boxlist([proposal, gt_box])
  
          return proposals
  
      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 * 4, H, W
          """
          device = objectness.device
          N, A, H, W = objectness.shape
  
          # put in the same format as anchors
          objectness = permute_and_flatten(objectness, N, A, 1, H, W).view(N, -1)
          objectness = objectness.sigmoid()
  
          box_regression = permute_and_flatten(box_regression, N, A, 4, H, W)
  
          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, 4)[batch_idx, topk_idx]
  
          proposals = self.box_coder.decode(
              box_regression.view(-1, 4), concat_anchors.view(-1, 4)
          )
  
          proposals = proposals.view(N, -1, 4)
  
          result = []
          for proposal, score, im_shape in zip(proposals, objectness, image_shapes):
              boxlist = BoxList(proposal, im_shape, mode="xyxy")
              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",
              )
              result.append(boxlist)
          return result
  
      def forward(self, anchors, objectness, box_regression, targets=None,
                  centerness=None, rpn_center_box_regression=None, centerness_pack=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)
  
          if self.pred_targets:
              pred_targets = []
              if True:
                  for img_centerness, center_box_reg in zip(centerness, rpn_center_box_regression):
                      # gt_centerness, gt_bbox, anchor_bbox = center_target
                      # print(rpn_center_box_regression, anchor_bbox)
                      # gt_mask = gt_centerness.detach().cpu().numpy() > 0.0
                      img_centerness = img_centerness[0, :, :]
  
                      center_box_reg = center_box_reg[:, :, :].permute(1, 2, 0)
  
                      anchor_bbox = np.zeros(shape=(center_box_reg.shape[0], center_box_reg.shape[1], 4))
                      for xx in range(anchor_bbox.shape[1]):
                          for yy in range(anchor_bbox.shape[0]):
                              anchor_bbox[yy, xx, :] = [max(0.0, xx * 4 - 16), max(0.0, yy * 4 - 16),
                                                        min(xx * 4 + 16, boxlists[0].size[0]),
                                                        min(yy * 4 + 16, boxlists[0].size[1])]
                      anchor_bbox = torch.as_tensor(anchor_bbox, device=center_box_reg.device)
  
                      # print(center_box_reg.shape, anchor_bbox.shape)
                      boxes = self.box_coder.decode(center_box_reg.reshape(-1, 4), anchor_bbox.view(-1, 4))
  
                      pred_target = None
                      pred_score = torch.sigmoid(img_centerness.detach()).cpu().numpy()
                      pred_mask = pred_score > 0.95
                      # print(gt_mask.shape, pred_mask.shape)
                      imllabel, numlabel = scipy.ndimage.label(pred_mask)
                      if numlabel > 0:
                          valid = np.zeros(shape=(numlabel,), dtype=np.bool)
                          box_inds = []
                          for ano in range(1, numlabel + 1):
                              mask = imllabel == ano
                              valid[ano - 1] = True  #  gt_mask[mask].sum() == 0
                              box_inds.append(np.argmax(pred_score * mask))
                          if np.any(valid):
                              boxes = boxes[box_inds, :]
                              # print(box_inds, boxes, anchor_bbox.view(-1, 4)[box_inds, :], gt_bbox.view(-1, 4)[box_inds, :])
                              pred_target = BoxList(torch.as_tensor(boxes), boxlists[0].size, mode="xyxy")
                              pred_target.clip_to_image()
                              pred_target = pred_target.to(img_centerness.device)
                              # print(img_centerness.device, pred_target.bbox.device)
                      pred_targets.append(pred_target)
              else:
                  for img_centerness in centerness:
                      pred_target = None
                      pred_mask = torch.sigmoid(img_centerness[0,:,:].detach()).cpu().numpy() > 0.95
                      # print(gt_mask.shape, pred_mask.shape)
                      imllabel, numlabel = scipy.ndimage.label(pred_mask)
                      if numlabel > 0:
                          masks = np.zeros(shape=(pred_mask.shape[0], pred_mask.shape[1], numlabel), dtype=np.uint8)
                          valid = np.zeros(shape=(numlabel,), dtype=np.bool)
                          for ano in range(1, numlabel + 1):
                              mask = imllabel == ano
                              valid[ano - 1] = True
                              masks[:, :, ano - 1] = mask
                          if np.any(valid):
                              masks = masks[:, :, valid]
                              boxes = extract_bboxes(masks)
                              pred_target = BoxList(torch.as_tensor(boxes), boxlists[0].size, mode="xyxy")
                              pred_target.clip_to_image()
                              pred_target = pred_target.to(img_centerness.device)
                              # print(img_centerness.device, pred_target.bbox.device)
                      pred_targets.append(pred_target)
  
              if True:
                  if not self.training:
                      print('add', [len(pred_target) for pred_target in pred_targets if pred_target], 'proposals')
                  boxlists = self.add_pred_proposals(boxlists, pred_targets)
          else:
              pred_targets = None
  
          return boxlists, pred_targets
  
      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
          # NOTE: it should be per image, and not per batch. However, to be consistent 
          # with Detectron, the default is per batch (see Issue #672)
          if self.training and self.fpn_post_nms_per_batch:
              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.bool)
              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
      fpn_post_nms_per_batch = config.MODEL.RPN.FPN_POST_NMS_PER_BATCH
      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,
          fpn_post_nms_per_batch=fpn_post_nms_per_batch,
          pred_targets=config.MODEL.RPN.PRED_TARGETS,
      )
      return box_selector
  