text: """CTPN positive and negative sample screening for RPN."""
  
  import numpy as np
  import mindspore.nn as nn
  from mindspore.ops import operations as P
  from mindspore.common.tensor import Tensor
  import mindspore.common.dtype as mstype
  from src.CTPN.BoundingBoxEncode import BoundingBoxEncode
  
  
  
  class BboxAssignSample(nn.Cell):
      """
      Bbox assigner and sampler definition.
  
      Args:
          config (dict): Config.
          batch_size (int): Batchsize.
          num_bboxes (int): The anchor nums.
          add_gt_as_proposals (bool): add gt bboxes as proposals flag.
  
      Returns:
          Tensor, output tensor.
          bbox_targets: bbox location, (batch_size, num_bboxes, 4)
          bbox_weights: bbox weights, (batch_size, num_bboxes, 1)
          labels: label for every bboxes, (batch_size, num_bboxes, 1)
          label_weights: label weight for every bboxes, (batch_size, num_bboxes, 1)
  
      Examples:
          BboxAssignSample(config, 2, 1024, True)
      """
  
      def __init__(self, config, batch_size, num_bboxes, add_gt_as_proposals):
          super(BboxAssignSample, self).__init__()
          cfg = config
          self.batch_size = batch_size
  
          self.neg_iou_thr = Tensor(cfg.neg_iou_thr, mstype.float16)
          self.pos_iou_thr = Tensor(cfg.pos_iou_thr, mstype.float16)
          self.min_pos_iou = Tensor(cfg.min_pos_iou, mstype.float16)
          self.zero_thr = Tensor(0.0, mstype.float16)
  
          self.num_bboxes = num_bboxes
          self.num_gts = cfg.num_gts
          self.num_expected_pos = cfg.num_expected_pos
          self.num_expected_neg = cfg.num_expected_neg
          self.add_gt_as_proposals = add_gt_as_proposals
  
          if self.add_gt_as_proposals:
              self.label_inds = Tensor(np.arange(1, self.num_gts + 1))
  
          self.concat = P.Concat(axis=0)
          self.max_gt = P.ArgMaxWithValue(axis=0)
          self.max_anchor = P.ArgMaxWithValue(axis=1)
          self.sum_inds = P.ReduceSum()
          self.iou = P.IOU()
          self.greaterequal = P.GreaterEqual()
          self.greater = P.Greater()
          self.select = P.Select()
          self.gatherND = P.GatherNd()
          self.squeeze = P.Squeeze()
          self.cast = P.Cast()
          self.logicaland = P.LogicalAnd()
          self.less = P.Less()
          self.random_choice_with_mask_pos = P.RandomChoiceWithMask(self.num_expected_pos)
          self.random_choice_with_mask_neg = P.RandomChoiceWithMask(self.num_expected_neg)
          self.reshape = P.Reshape()
          self.equal = P.Equal()
          self.bounding_box_encode = BoundingBoxEncode()
          self.scatterNdUpdate = P.ScatterNdUpdate()
          self.scatterNd = P.ScatterNd()
          self.logicalnot = P.LogicalNot()
          self.tile = P.Tile()
          self.zeros_like = P.ZerosLike()
  
          self.assigned_gt_inds = Tensor(np.array(-1 * np.ones(num_bboxes), dtype=np.int32))
          self.assigned_gt_zeros = Tensor(np.array(np.zeros(num_bboxes), dtype=np.int32))
          self.assigned_gt_ones = Tensor(np.array(np.ones(num_bboxes), dtype=np.int32))
          self.assigned_gt_ignores = Tensor(np.array(-1 * np.ones(num_bboxes), dtype=np.int32))
          self.assigned_pos_ones = Tensor(np.array(np.ones(self.num_expected_pos), dtype=np.int32))
  
          self.check_neg_mask = Tensor(np.array(np.ones(self.num_expected_neg - self.num_expected_pos), dtype=np.bool))
          self.range_pos_size = Tensor(np.arange(self.num_expected_pos).astype(np.float16))
          self.check_gt_one = Tensor(np.array(-1 * np.ones((self.num_gts, 4)), dtype=np.float16))
          self.check_anchor_two = Tensor(np.array(-2 * np.ones((self.num_bboxes, 4)), dtype=np.float16))
          self.print = P.Print()
  
  
      def construct(self, gt_bboxes_i, gt_labels_i, valid_mask, bboxes, gt_valids):
          gt_bboxes_i = self.select(self.cast(self.tile(self.reshape(self.cast(gt_valids, mstype.int32), \
                                    (self.num_gts, 1)), (1, 4)), mstype.bool_), gt_bboxes_i, self.check_gt_one)
          bboxes = self.select(self.cast(self.tile(self.reshape(self.cast(valid_mask, mstype.int32), \
                               (self.num_bboxes, 1)), (1, 4)), mstype.bool_), bboxes, self.check_anchor_two)
          overlaps = self.iou(bboxes, gt_bboxes_i)
          max_overlaps_w_gt_index, max_overlaps_w_gt = self.max_gt(overlaps)
          _, max_overlaps_w_ac = self.max_anchor(overlaps)
          neg_sample_iou_mask = self.logicaland(self.greaterequal(max_overlaps_w_gt, self.zero_thr), \
                                                self.less(max_overlaps_w_gt, self.neg_iou_thr))
          assigned_gt_inds2 = self.select(neg_sample_iou_mask, self.assigned_gt_zeros, self.assigned_gt_inds)
          pos_sample_iou_mask = self.greaterequal(max_overlaps_w_gt, self.pos_iou_thr)
          assigned_gt_inds3 = self.select(pos_sample_iou_mask, \
                                          max_overlaps_w_gt_index + self.assigned_gt_ones, assigned_gt_inds2)
          assigned_gt_inds4 = assigned_gt_inds3
          for j in range(self.num_gts):
              max_overlaps_w_ac_j = max_overlaps_w_ac[j:j+1:1]
              overlaps_w_gt_j = self.squeeze(overlaps[j:j+1:1, ::])
  
              pos_mask_j = self.logicaland(self.greaterequal(max_overlaps_w_ac_j, self.min_pos_iou), \
                                          self.equal(overlaps_w_gt_j, max_overlaps_w_ac_j))
              assigned_gt_inds4 = self.select(pos_mask_j, self.assigned_gt_ones + j, assigned_gt_inds4)
          assigned_gt_inds5 = self.select(valid_mask, assigned_gt_inds4, self.assigned_gt_ignores)
          pos_index, valid_pos_index = self.random_choice_with_mask_pos(self.greater(assigned_gt_inds5, 0))
          pos_check_valid = self.cast(self.greater(assigned_gt_inds5, 0), mstype.float16)
          pos_check_valid = self.sum_inds(pos_check_valid, -1)
          valid_pos_index = self.less(self.range_pos_size, pos_check_valid)
          pos_index = pos_index * self.reshape(self.cast(valid_pos_index, mstype.int32), (self.num_expected_pos, 1))
          pos_assigned_gt_index = self.gatherND(assigned_gt_inds5, pos_index) - self.assigned_pos_ones
          pos_assigned_gt_index = pos_assigned_gt_index * self.cast(valid_pos_index, mstype.int32)
          pos_assigned_gt_index = self.reshape(pos_assigned_gt_index, (self.num_expected_pos, 1))
          neg_index, valid_neg_index = self.random_choice_with_mask_neg(self.equal(assigned_gt_inds5, 0))
  
          num_pos = self.cast(self.logicalnot(valid_pos_index), mstype.float16)
          num_pos = self.sum_inds(num_pos, -1)
          unvalid_pos_index = self.less(self.range_pos_size, num_pos)
          valid_neg_index = self.logicaland(self.concat((self.check_neg_mask, unvalid_pos_index)), valid_neg_index)
          pos_bboxes_ = self.gatherND(bboxes, pos_index)
          pos_gt_bboxes_ = self.gatherND(gt_bboxes_i, pos_assigned_gt_index)
          pos_gt_labels = self.gatherND(gt_labels_i, pos_assigned_gt_index)
          pos_bbox_targets_ = self.bounding_box_encode(pos_bboxes_, pos_gt_bboxes_)
          valid_pos_index = self.cast(valid_pos_index, mstype.int32)
          valid_neg_index = self.cast(valid_neg_index, mstype.int32)
          bbox_targets_total = self.scatterNd(pos_index, pos_bbox_targets_, (self.num_bboxes, 4))
          bbox_weights_total = self.scatterNd(pos_index, valid_pos_index, (self.num_bboxes,))
          labels_total = self.scatterNd(pos_index, pos_gt_labels, (self.num_bboxes,))
          total_index = self.concat((pos_index, neg_index))
          total_valid_index = self.concat((valid_pos_index, valid_neg_index))
          label_weights_total = self.scatterNd(total_index, total_valid_index, (self.num_bboxes,))
          return bbox_targets_total, self.cast(bbox_weights_total, mstype.bool_), \
                 labels_total, self.cast(label_weights_total, mstype.bool_)
  