def ravel_indices(feature_grid_y, feature_grid_x, num_levels, height, width,
                  box_levels):
  """Returns grid indices in a flattened feature map of shape [-1, channels].

  The returned 1-D array can be used to gather feature grid points from a
  feature map that has been flattened from [batch, num_levels, max_height,
  max_width, channels] to [batch * num_levels * max_height * max_width,
  channels].

  Args:
    feature_grid_y: An int32 tensor of shape [batch, num_boxes, size_y]
      containing y coordinate vector.
    feature_grid_x: An int32 tensor of shape [batch, num_boxes, size_x]
      containing x coordinate vector.
    num_levels: Number of feature levels.
    height: An integer indicating the padded height of feature maps.
    width: An integer indicating the padded width of feature maps.
    box_levels: An int32 tensor of shape [batch, num_boxes] indicating
      feature level assigned to each box.

  Returns:
    indices: A 1D int32 tensor containing feature point indices in a flattened
      feature grid.
  """
  num_boxes = tf.shape(feature_grid_y)[1]
  batch_size = tf.shape(feature_grid_y)[0]
  size_y = tf.shape(feature_grid_y)[2]
  size_x = tf.shape(feature_grid_x)[2]
  height_dim_offset = width
  level_dim_offset = height * height_dim_offset
  batch_dim_offset = num_levels * level_dim_offset

  batch_dim_indices = (
      tf.reshape(
          tf.range(batch_size) * batch_dim_offset, [batch_size, 1, 1, 1]) *
      tf.ones([1, num_boxes, size_y, size_x], dtype=tf.int32))
  box_level_indices = (
      tf.reshape(box_levels * level_dim_offset, [batch_size, num_boxes, 1, 1]) *
      tf.ones([1, 1, size_y, size_x], dtype=tf.int32))
  height_indices = (
      tf.reshape(feature_grid_y * height_dim_offset,
                 [batch_size, num_boxes, size_y, 1]) *
      tf.ones([1, 1, 1, size_x], dtype=tf.int32))
  width_indices = (
      tf.reshape(feature_grid_x, [batch_size, num_boxes, 1, size_x])
      * tf.ones([1, 1, size_y, 1], dtype=tf.int32))
  indices = (
      batch_dim_indices + box_level_indices + height_indices + width_indices)
  flattened_indices = tf.reshape(indices, [-1])
  return flattened_indices
Returns grid indices in a flattened feature map of shape [-1, channels].

  The returned 1-D array can be used to gather feature grid points from a
  feature map that has been flattened from [batch, num_levels, max_height,
  max_width, channels] to [batch * num_levels * max_height * max_width,
  channels].

  Args:
    feature_grid_y: An int32 tensor of shape [batch, num_boxes, size_y]
      containing y coordinate vector.
    feature_grid_x: An int32 tensor of shape [batch, num_boxes, size_x]
      containing x coordinate vector.
    num_levels: Number of feature levels.
    height: An integer indicating the padded height of feature maps.
    width: An integer indicating the padded width of feature maps.
    box_levels: An int32 tensor of shape [batch, num_boxes] indicating
      feature level assigned to each box.

  Returns:
    indices: A 1D int32 tensor containing feature point indices in a flattened
      feature grid.
  
inference_api/src/main/object_detection/utils/spatial_transform_ops.py