text: # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Ops used to post-process raw detections."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  import box_utils
  
  
  def generate_detections_per_image_tpu(cls_outputs,
                                        box_outputs,
                                        anchor_boxes,
                                        image_info,
                                        pre_nms_num_detections=1000,
                                        post_nms_num_detections=100,
                                        nms_threshold=0.3,
                                        bbox_reg_weights=(10., 10., 5., 5.)):
    """Generate the final detections per image given the model outputs.
  
    Args:
      cls_outputs: a tensor with shape [N, num_classes], which stacks class
        logit outputs 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 cls_outputs should be the output of softmax().
      box_outputs: a tensor with shape [N, num_classes*4], which stacks box
        regression outputs on all feature levels. The N is the number of total
        anchors on all levels.
      anchor_boxes: a tensor with shape [N, 4], which stacks anchors on all
        feature levels. The N is the number of total anchors on all levels.
      image_info: a tensor of shape [5] which encodes the input image's [height,
        width, scale, original_height, original_width]
      pre_nms_num_detections: an integer that specifies the number of candidates
        before NMS.
      post_nms_num_detections: an integer that specifies the number of candidates
        after NMS.
      nms_threshold: a float number to specify the IOU threshold of NMS.
      bbox_reg_weights: a list of 4 float scalars, which are default weights on
        (dx, dy, dw, dh) for normalizing bbox regression targets.
  
    Returns:
      detections: Tuple of tensors corresponding to number of valid boxes,
      box coordinates, object categories for each boxes, and box scores
      -- respectively.
    """
    num_boxes, num_classes = cls_outputs.get_shape().as_list()
  
    # Remove background class scores.
    cls_outputs = cls_outputs[:, 1:num_classes]
    top_k_scores, top_k_indices_with_classes = tf.nn.top_k(
        tf.reshape(cls_outputs, [-1]),
        k=pre_nms_num_detections,
        sorted=False)
    classes = tf.mod(top_k_indices_with_classes, num_classes - 1)
    top_k_indices = tf.floordiv(top_k_indices_with_classes, num_classes - 1)
  
    anchor_boxes = tf.gather(anchor_boxes, top_k_indices)
    box_outputs = tf.reshape(
        box_outputs, [num_boxes, num_classes, 4])[:, 1:num_classes, :]
    class_indices = classes
    box_outputs = tf.gather_nd(box_outputs,
                               tf.stack([top_k_indices, class_indices], axis=1))
  
    # apply bounding box regression to anchors
    boxes = box_utils.decode_boxes(
        box_outputs, anchor_boxes, bbox_reg_weights)
    boxes = box_utils.clip_boxes(
        boxes, image_info[0], image_info[1])
  
    list_of_all_boxes = []
    list_of_all_scores = []
    list_of_all_classes = []
    # Skip background class.
    for class_i in range(num_classes):
      # Compute bitmask for the given classes.
      class_i_bitmask = tf.cast(tf.equal(classes, class_i), top_k_scores.dtype)
      # This works because score is in [0, 1].
      class_i_scores = top_k_scores * class_i_bitmask
      # The TPU and CPU have different behaviors for
      # tf.image.non_max_suppression_padded (b/116754376).
      (class_i_post_nms_indices,
       class_i_nms_num_valid) = tf.image.non_max_suppression_padded(
           tf.to_float(boxes),
           tf.to_float(class_i_scores),
           post_nms_num_detections,
           iou_threshold=nms_threshold,
           score_threshold=0.05,
           pad_to_max_output_size=True,
           name='nms_detections_' + str(class_i))
      class_i_post_nms_boxes = tf.gather(boxes, class_i_post_nms_indices)
      class_i_post_nms_scores = tf.gather(class_i_scores,
                                          class_i_post_nms_indices)
      mask = tf.less(tf.range(post_nms_num_detections), [class_i_nms_num_valid])
      class_i_post_nms_scores = tf.where(
          mask, class_i_post_nms_scores, tf.zeros_like(class_i_post_nms_scores))
      class_i_classes = tf.fill(tf.shape(class_i_post_nms_scores), class_i+1)
      list_of_all_boxes.append(class_i_post_nms_boxes)
      list_of_all_scores.append(class_i_post_nms_scores)
      list_of_all_classes.append(class_i_classes)
  
    post_nms_boxes = tf.concat(list_of_all_boxes, axis=0)
    post_nms_scores = tf.concat(list_of_all_scores, axis=0)
    post_nms_classes = tf.concat(list_of_all_classes, axis=0)
  
    # sort all results.
    post_nms_scores, sorted_indices = tf.nn.top_k(
        tf.to_float(post_nms_scores),
        k=post_nms_num_detections,
        sorted=True)
    post_nms_boxes = tf.gather(post_nms_boxes, sorted_indices)
    post_nms_classes = tf.gather(post_nms_classes, sorted_indices)
  
    valid_mask = tf.where(
        tf.greater(post_nms_scores, 0), tf.ones_like(post_nms_scores),
        tf.zeros_like(post_nms_scores))
    num_valid_boxes = tf.reduce_sum(valid_mask, axis=-1)
    box_classes = tf.to_float(post_nms_classes)
    return num_valid_boxes, post_nms_boxes, box_classes, post_nms_scores
  
  
  def generate_detections_tpu(class_outputs,
                              box_outputs,
                              anchor_boxes,
                              image_info,
                              pre_nms_num_detections=1000,
                              post_nms_num_detections=100,
                              nms_threshold=0.3,
                              bbox_reg_weights=(10., 10., 5., 5.)):
    """Generate the final detections given the model outputs (TPU version).
  
    Args:
      class_outputs: a tensor with shape [batch_size, N, num_classes], which
        stacks class logit outputs 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.
      box_outputs: a tensor with shape [batch_size, N, num_classes*4], which
        stacks box regression outputs on all feature levels. The N is the number
        of total anchors on all levels.
      anchor_boxes: a tensor with shape [batch_size, N, 4], which stacks anchors
        on all feature levels. The N is the number of total anchors on all levels.
      image_info: a tensor of shape [batch_size, 5] which encodes each image's
        [height, width, scale, original_height, original_width].
      pre_nms_num_detections: an integer that specifies the number of candidates
        before NMS.
      post_nms_num_detections: an integer that specifies the number of candidates
        after NMS.
      nms_threshold: a float number to specify the IOU threshold of NMS.
      bbox_reg_weights: a list of 4 float scalars, which are default weights on
        (dx, dy, dw, dh) for normalizing bbox regression targets.
  
    Returns:
      a tuple of tensors corresponding to number of valid boxes,
      box coordinates, object categories for each boxes, and box scores stacked
      in batch_size.
    """
    with tf.name_scope('generate_detections'):
      batch_size, _, _ = class_outputs.get_shape().as_list()
      softmax_class_outputs = tf.nn.softmax(class_outputs)
  
      num_valid_boxes, box_coordinates, box_classes, box_scores = ([], [], [], [])
      for i in range(batch_size):
        result = generate_detections_per_image_tpu(
            softmax_class_outputs[i], box_outputs[i], anchor_boxes[i],
            image_info[i], pre_nms_num_detections, post_nms_num_detections,
            nms_threshold, bbox_reg_weights)
  
        num_valid_boxes.append(result[0])
        box_coordinates.append(result[1])
        box_classes.append(result[2])
        box_scores.append(result[3])
  
      num_valid_boxes = tf.stack(num_valid_boxes)
      box_coordinates = tf.stack(box_coordinates)
      box_classes = tf.stack(box_classes)
      box_scores = tf.stack(box_scores)
      return num_valid_boxes, box_coordinates, box_classes, box_scores
  
  
  def generate_detections_gpu(class_outputs,
                              box_outputs,
                              anchor_boxes,
                              image_info,
                              pre_nms_num_detections=1000,
                              post_nms_num_detections=100,
                              nms_threshold=0.3,
                              bbox_reg_weights=(10., 10., 5., 5.)):
    """Generate the final detections given the model outputs (GPU version).
  
    Args:
      class_outputs: a tensor with shape [batch_size, N, num_classes], which
        stacks class logit outputs 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.
      box_outputs: a tensor with shape [batch_size, N, num_classes*4], which
        stacks box regression outputs on all feature levels. The N is the number
        of total anchors on all levels.
      anchor_boxes: a tensor with shape [batch_size, N, 4], which stacks anchors
        on all feature levels. The N is the number of total anchors on all levels.
      image_info: a tensor of shape [batch_size, 5] which encodes each image's
        [height, width, scale, original_height, original_width].
      pre_nms_num_detections: an integer that specifies the number of candidates
        before NMS.
      post_nms_num_detections: an integer that specifies the number of candidates
        after NMS.
      nms_threshold: a float number to specify the IOU threshold of NMS.
      bbox_reg_weights: a list of 4 float scalars, which are default weights on
        (dx, dy, dw, dh) for normalizing bbox regression targets.
  
    Returns:
      a tuple of tensors corresponding to number of valid boxes,
      box coordinates, object categories for each boxes, and box scores stacked
      in batch_size.
    """
    with tf.name_scope('generate_detections'):
      batch_size, num_boxes, num_classes = class_outputs.get_shape().as_list()
      softmax_class_outputs = tf.nn.softmax(class_outputs)
  
      # Remove background
      scores = tf.slice(softmax_class_outputs, [0, 0, 1], [-1, -1, -1])
      boxes = tf.slice(
          tf.reshape(box_outputs, [batch_size, num_boxes, num_classes, 4]),
          [0, 0, 1, 0], [-1, -1, -1, -1])
  
      anchor_boxes = (tf.expand_dims(anchor_boxes, axis=2) *
                      tf.ones([1, 1, num_classes - 1, 1]))
  
      num_detections = num_boxes * (num_classes - 1)
  
      boxes = tf.reshape(boxes, [batch_size, num_detections, 4])
      scores = tf.reshape(scores, [batch_size, num_detections, 1])
      anchor_boxes = tf.reshape(anchor_boxes, [batch_size, num_detections, 4])
  
      # Decode
      boxes = box_utils.decode_boxes(
          boxes, anchor_boxes, bbox_reg_weights)
  
      # Clip boxes
      height = tf.expand_dims(image_info[:, 0:1], axis=-1)
      width = tf.expand_dims(image_info[:, 1:2], axis=-1)
      boxes = box_utils.clip_boxes(boxes, height, width)
  
      # NMS
      pre_nms_boxes = box_utils.to_normalized_coordinates(
          boxes, height, width)
      pre_nms_boxes = tf.reshape(
          pre_nms_boxes, [batch_size, num_boxes, num_classes - 1, 4])
      pre_nms_scores = tf.reshape(
          scores, [batch_size, num_boxes, num_classes - 1])
      (post_nms_boxes, post_nms_scores, post_nms_classes,
       post_nms_num_valid_boxes) = (
           tf.image.combined_non_max_suppression(
               pre_nms_boxes,
               pre_nms_scores,
               max_output_size_per_class=pre_nms_num_detections,
               max_total_size=post_nms_num_detections,
               iou_threshold=nms_threshold,
               score_threshold=0.0,
               pad_per_class=False))
      post_nms_classes = post_nms_classes + 1
      post_nms_boxes = box_utils.to_absolute_coordinates(
          post_nms_boxes, height, width)
      return (post_nms_num_valid_boxes, post_nms_boxes,
              tf.to_float(post_nms_classes), post_nms_scores)
  