def _generate_detections_batched(boxes: tf.Tensor, scores: tf.Tensor, pre_nms_score_threshold: float, nms_iou_threshold: float, max_num_detections: int):
    'Generates detected boxes with scores and classes for one-stage detector.\n\n  The function takes output of multi-level ConvNets and anchor boxes and\n  generates detected boxes. Note that this used batched nms, which is not\n  supported on TPU currently.\n\n  Args:\n    boxes: A `tf.Tensor` with shape `[batch_size, N, num_classes, 4]` or\n      `[batch_size, N, 1, 4]`, which box predictions on all feature levels. The\n      N is the number of total anchors on all levels.\n    scores: A `tf.Tensor` with shape `[batch_size, N, num_classes]`, which\n      stacks class probability on all feature levels. The N is the number of\n      total anchors on all levels. The num_classes is the number of classes\n      predicted by the model. Note that the class_outputs here is the raw score.\n    pre_nms_score_threshold: A `float` representing the threshold for deciding\n      when to remove boxes based on score.\n    nms_iou_threshold: A `float` representing the threshold for deciding whether\n      boxes overlap too much with respect to IOU.\n    max_num_detections: A `scalar` representing maximum number of boxes retained\n      over all classes.\n\n  Returns:\n    nms_boxes: A `float` tf.Tensor of shape [batch_size, max_num_detections, 4]\n      representing top detected boxes in [y1, x1, y2, x2].\n    nms_scores: A `float` tf.Tensor of shape [batch_size, max_num_detections]\n      representing sorted confidence scores for detected boxes. The values are\n      between [0, 1].\n    nms_classes: An `int` tf.Tensor of shape [batch_size, max_num_detections]\n      representing classes for detected boxes.\n    valid_detections: An `int` tf.Tensor of shape [batch_size] only the top\n      `valid_detections` boxes are valid detections.\n  '
    with tf.name_scope('generate_detections'):
        (nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections) = tf.image.combined_non_max_suppression(boxes, scores, max_output_size_per_class=max_num_detections, max_total_size=max_num_detections, iou_threshold=nms_iou_threshold, score_threshold=pre_nms_score_threshold, pad_per_class=False, clip_boxes=False)
        nmsed_classes = tf.cast(nmsed_classes, tf.int32)
    return (nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections)
Generates detected boxes with scores and classes for one-stage detector.

The function takes output of multi-level ConvNets and anchor boxes and
generates detected boxes. Note that this used batched nms, which is not
supported on TPU currently.

Args:
  boxes: A `tf.Tensor` with shape `[batch_size, N, num_classes, 4]` or
    `[batch_size, N, 1, 4]`, which box predictions on all feature levels. The
    N is the number of total anchors on all levels.
  scores: A `tf.Tensor` with shape `[batch_size, N, num_classes]`, which
    stacks class probability on all feature levels. The N is the number of
    total anchors on all levels. The num_classes is the number of classes
    predicted by the model. Note that the class_outputs here is the raw score.
  pre_nms_score_threshold: A `float` representing the threshold for deciding
    when to remove boxes based on score.
  nms_iou_threshold: A `float` representing the threshold for deciding whether
    boxes overlap too much with respect to IOU.
  max_num_detections: A `scalar` representing maximum number of boxes retained
    over all classes.

Returns:
  nms_boxes: A `float` tf.Tensor of shape [batch_size, max_num_detections, 4]
    representing top detected boxes in [y1, x1, y2, x2].
  nms_scores: A `float` tf.Tensor of shape [batch_size, max_num_detections]
    representing sorted confidence scores for detected boxes. The values are
    between [0, 1].
  nms_classes: An `int` tf.Tensor of shape [batch_size, max_num_detections]
    representing classes for detected boxes.
  valid_detections: An `int` tf.Tensor of shape [batch_size] only the top
    `valid_detections` boxes are valid detections.
def _generate_detections_batched(boxes: tf.Tensor, scores: tf.Tensor, pre_nms_score_threshold: float, nms_iou_threshold: float, max_num_detections: int):
    'Generates detected boxes with scores and classes for one-stage detector.\n\n  The function takes output of multi-level ConvNets and anchor boxes and\n  generates detected boxes. Note that this used batched nms, which is not\n  supported on TPU currently.\n\n  Args:\n    boxes: A `tf.Tensor` with shape `[batch_size, N, num_classes, 4]` or\n      `[batch_size, N, 1, 4]`, which box predictions on all feature levels. The\n      N is the number of total anchors on all levels.\n    scores: A `tf.Tensor` with shape `[batch_size, N, num_classes]`, which\n      stacks class probability on all feature levels. The N is the number of\n      total anchors on all levels. The num_classes is the number of classes\n      predicted by the model. Note that the class_outputs here is the raw score.\n    pre_nms_score_threshold: A `float` representing the threshold for deciding\n      when to remove boxes based on score.\n    nms_iou_threshold: A `float` representing the threshold for deciding whether\n      boxes overlap too much with respect to IOU.\n    max_num_detections: A `scalar` representing maximum number of boxes retained\n      over all classes.\n\n  Returns:\n    nms_boxes: A `float` tf.Tensor of shape [batch_size, max_num_detections, 4]\n      representing top detected boxes in [y1, x1, y2, x2].\n    nms_scores: A `float` tf.Tensor of shape [batch_size, max_num_detections]\n      representing sorted confidence scores for detected boxes. The values are\n      between [0, 1].\n    nms_classes: An `int` tf.Tensor of shape [batch_size, max_num_detections]\n      representing classes for detected boxes.\n    valid_detections: An `int` tf.Tensor of shape [batch_size] only the top\n      `valid_detections` boxes are valid detections.\n  '
    with tf.name_scope('generate_detections'):
        (nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections) = tf.image.combined_non_max_suppression(boxes, scores, max_output_size_per_class=max_num_detections, max_total_size=max_num_detections, iou_threshold=nms_iou_threshold, score_threshold=pre_nms_score_threshold, pad_per_class=False, clip_boxes=False)
        nmsed_classes = tf.cast(nmsed_classes, tf.int32)
    return (nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections)