def batch_assign(self,
                   pred_box_batch,
                   gt_box_batch,
                   pred_class_batch,
                   gt_class_targets_batch,
                   gt_weights_batch=None,
                   unmatched_class_label_batch=None):
    """Batched assignment of classification and regression targets.

    Args:
      pred_box_batch: a tensor of shape [batch_size, num_queries, 4]
        representing predicted bounding boxes.
      gt_box_batch: a tensor of shape [batch_size, num_queries, 4]
        representing groundtruth bounding boxes.
      pred_class_batch: A list of tensors with length batch_size, where each
        each tensor has shape [num_queries, num_classes] to be used
        by certain similarity calculators.
      gt_class_targets_batch: a list of tensors with length batch_size, where
        each tensor has shape [num_gt_boxes_i, num_classes] and
        num_gt_boxes_i is the number of boxes in the ith boxlist of
        gt_box_batch.
      gt_weights_batch: A list of 1-D tf.float32 tensors of shape
        [num_boxes] containing weights for groundtruth boxes.
      unmatched_class_label_batch: a float32 tensor with shape
        [d_1, d_2, ..., d_k] which is consistent with the classification target
        for each anchor (and can be empty for scalar targets).  This shape must
        thus be compatible with the `gt_class_targets_batch`.

    Returns:
      batch_cls_targets: a tensor with shape [batch_size, num_pred_boxes,
        num_classes],
      batch_cls_weights: a tensor with shape [batch_size, num_pred_boxes,
        num_classes],
      batch_reg_targets: a tensor with shape [batch_size, num_pred_boxes,
        box_code_dimension]
      batch_reg_weights: a tensor with shape [batch_size, num_pred_boxes].
    """
    pred_box_batch = [
        box_list.BoxList(pred_box)
        for pred_box in tf.unstack(pred_box_batch)]
    gt_box_batch = [
        box_list.BoxList(gt_box)
        for gt_box in tf.unstack(gt_box_batch)]

    cls_targets_list = []
    cls_weights_list = []
    reg_targets_list = []
    reg_weights_list = []
    if gt_weights_batch is None:
      gt_weights_batch = [None] * len(gt_class_targets_batch)
    if unmatched_class_label_batch is None:
      unmatched_class_label_batch = [None] * len(gt_class_targets_batch)
    pred_class_batch = tf.unstack(pred_class_batch)
    for (pred_boxes, gt_boxes, pred_class_batch, gt_class_targets, gt_weights,
         unmatched_class_label) in zip(pred_box_batch, gt_box_batch,
                                       pred_class_batch, gt_class_targets_batch,
                                       gt_weights_batch,
                                       unmatched_class_label_batch):
      (cls_targets, cls_weights, reg_targets,
       reg_weights) = self.assign(pred_boxes, gt_boxes, pred_class_batch,
                                  gt_class_targets, gt_weights,
                                  unmatched_class_label)
      cls_targets_list.append(cls_targets)
      cls_weights_list.append(cls_weights)
      reg_targets_list.append(reg_targets)
      reg_weights_list.append(reg_weights)
    batch_cls_targets = tf.stack(cls_targets_list)
    batch_cls_weights = tf.stack(cls_weights_list)
    batch_reg_targets = tf.stack(reg_targets_list)
    batch_reg_weights = tf.stack(reg_weights_list)
    return (batch_cls_targets, batch_cls_weights, batch_reg_targets,
            batch_reg_weights)
Batched assignment of classification and regression targets.

    Args:
      pred_box_batch: a tensor of shape [batch_size, num_queries, 4]
        representing predicted bounding boxes.
      gt_box_batch: a tensor of shape [batch_size, num_queries, 4]
        representing groundtruth bounding boxes.
      pred_class_batch: A list of tensors with length batch_size, where each
        each tensor has shape [num_queries, num_classes] to be used
        by certain similarity calculators.
      gt_class_targets_batch: a list of tensors with length batch_size, where
        each tensor has shape [num_gt_boxes_i, num_classes] and
        num_gt_boxes_i is the number of boxes in the ith boxlist of
        gt_box_batch.
      gt_weights_batch: A list of 1-D tf.float32 tensors of shape
        [num_boxes] containing weights for groundtruth boxes.
      unmatched_class_label_batch: a float32 tensor with shape
        [d_1, d_2, ..., d_k] which is consistent with the classification target
        for each anchor (and can be empty for scalar targets).  This shape must
        thus be compatible with the `gt_class_targets_batch`.

    Returns:
      batch_cls_targets: a tensor with shape [batch_size, num_pred_boxes,
        num_classes],
      batch_cls_weights: a tensor with shape [batch_size, num_pred_boxes,
        num_classes],
      batch_reg_targets: a tensor with shape [batch_size, num_pred_boxes,
        box_code_dimension]
      batch_reg_weights: a tensor with shape [batch_size, num_pred_boxes].
    
inference_api/src/main/object_detection/core/target_assigner.py