def prediction_tensors_to_keypoint_candidates(keypoint_heatmap_predictions, keypoint_heatmap_offsets, keypoint_score_threshold=0.1, max_pool_kernel_size=1, max_candidates=20):
    "Convert keypoint heatmap predictions and offsets to keypoint candidates.\n\n  Args:\n    keypoint_heatmap_predictions: A float tensor of shape [batch_size, height,\n      width, num_keypoints] representing the per-keypoint heatmaps.\n    keypoint_heatmap_offsets: A float tensor of shape [batch_size, height,\n      width, 2] (or [batch_size, height, width, 2 * num_keypoints] if\n      'per_keypoint_offset' is set True) representing the per-keypoint offsets.\n    keypoint_score_threshold: float, the threshold for considering a keypoint\n      a candidate.\n    max_pool_kernel_size: integer, the max pool kernel size to use to pull off\n      peak score locations in a neighborhood. For example, to make sure no two\n      neighboring values for the same keypoint are returned, set\n      max_pool_kernel_size=3. If None or 1, will not apply any local filtering.\n    max_candidates: integer, maximum number of keypoint candidates per\n      keypoint type.\n\n  Returns:\n    keypoint_candidates: A tensor of shape\n      [batch_size, max_candidates, num_keypoints, 2] holding the\n      location of keypoint candidates in [y, x] format (expressed in absolute\n      coordinates in the output coordinate frame).\n    keypoint_scores: A float tensor of shape\n      [batch_size, max_candidates, num_keypoints] with the scores for each\n      keypoint candidate. The scores come directly from the heatmap predictions.\n    num_keypoint_candidates: An integer tensor of shape\n      [batch_size, num_keypoints] with the number of candidates for each\n      keypoint type, as it's possible to filter some candidates due to the score\n      threshold.\n  "
    (batch_size, _, width, num_keypoints) = _get_shape(keypoint_heatmap_predictions, 4)
    (keypoint_scores, y_indices, x_indices, channel_indices) = top_k_feature_map_locations(keypoint_heatmap_predictions, max_pool_kernel_size=max_pool_kernel_size, k=max_candidates, per_channel=True)
    peak_spatial_indices = flattened_indices_from_row_col_indices(y_indices, x_indices, width)
    y_indices = _to_float32(y_indices)
    x_indices = _to_float32(x_indices)
    offsets_flat = _flatten_spatial_dimensions(keypoint_heatmap_offsets)
    selected_offsets = tf.gather(offsets_flat, peak_spatial_indices, batch_dims=1)
    (_, num_indices, num_channels) = _get_shape(selected_offsets, 3)
    if (num_channels > 2):
        reshaped_offsets = tf.reshape(selected_offsets, [batch_size, num_indices, (- 1), 2])
        offsets = tf.gather(reshaped_offsets, channel_indices, batch_dims=2)
    else:
        offsets = selected_offsets
    (y_offsets, x_offsets) = tf.unstack(offsets, axis=2)
    keypoint_candidates = tf.stack([(y_indices + y_offsets), (x_indices + x_offsets)], axis=2)
    keypoint_candidates = tf.reshape(keypoint_candidates, [batch_size, num_keypoints, max_candidates, 2])
    keypoint_candidates = tf.transpose(keypoint_candidates, [0, 2, 1, 3])
    keypoint_scores = tf.reshape(keypoint_scores, [batch_size, num_keypoints, max_candidates])
    keypoint_scores = tf.transpose(keypoint_scores, [0, 2, 1])
    num_candidates = tf.reduce_sum(tf.to_int32((keypoint_scores >= keypoint_score_threshold)), axis=1)
    return (keypoint_candidates, keypoint_scores, num_candidates)
Convert keypoint heatmap predictions and offsets to keypoint candidates.

Args:
  keypoint_heatmap_predictions: A float tensor of shape [batch_size, height,
    width, num_keypoints] representing the per-keypoint heatmaps.
  keypoint_heatmap_offsets: A float tensor of shape [batch_size, height,
    width, 2] (or [batch_size, height, width, 2 * num_keypoints] if
    'per_keypoint_offset' is set True) representing the per-keypoint offsets.
  keypoint_score_threshold: float, the threshold for considering a keypoint
    a candidate.
  max_pool_kernel_size: integer, the max pool kernel size to use to pull off
    peak score locations in a neighborhood. For example, to make sure no two
    neighboring values for the same keypoint are returned, set
    max_pool_kernel_size=3. If None or 1, will not apply any local filtering.
  max_candidates: integer, maximum number of keypoint candidates per
    keypoint type.

Returns:
  keypoint_candidates: A tensor of shape
    [batch_size, max_candidates, num_keypoints, 2] holding the
    location of keypoint candidates in [y, x] format (expressed in absolute
    coordinates in the output coordinate frame).
  keypoint_scores: A float tensor of shape
    [batch_size, max_candidates, num_keypoints] with the scores for each
    keypoint candidate. The scores come directly from the heatmap predictions.
  num_keypoint_candidates: An integer tensor of shape
    [batch_size, num_keypoints] with the number of candidates for each
    keypoint type, as it's possible to filter some candidates due to the score
    threshold.
def prediction_tensors_to_keypoint_candidates(keypoint_heatmap_predictions, keypoint_heatmap_offsets, keypoint_score_threshold=0.1, max_pool_kernel_size=1, max_candidates=20):
    "Convert keypoint heatmap predictions and offsets to keypoint candidates.\n\n  Args:\n    keypoint_heatmap_predictions: A float tensor of shape [batch_size, height,\n      width, num_keypoints] representing the per-keypoint heatmaps.\n    keypoint_heatmap_offsets: A float tensor of shape [batch_size, height,\n      width, 2] (or [batch_size, height, width, 2 * num_keypoints] if\n      'per_keypoint_offset' is set True) representing the per-keypoint offsets.\n    keypoint_score_threshold: float, the threshold for considering a keypoint\n      a candidate.\n    max_pool_kernel_size: integer, the max pool kernel size to use to pull off\n      peak score locations in a neighborhood. For example, to make sure no two\n      neighboring values for the same keypoint are returned, set\n      max_pool_kernel_size=3. If None or 1, will not apply any local filtering.\n    max_candidates: integer, maximum number of keypoint candidates per\n      keypoint type.\n\n  Returns:\n    keypoint_candidates: A tensor of shape\n      [batch_size, max_candidates, num_keypoints, 2] holding the\n      location of keypoint candidates in [y, x] format (expressed in absolute\n      coordinates in the output coordinate frame).\n    keypoint_scores: A float tensor of shape\n      [batch_size, max_candidates, num_keypoints] with the scores for each\n      keypoint candidate. The scores come directly from the heatmap predictions.\n    num_keypoint_candidates: An integer tensor of shape\n      [batch_size, num_keypoints] with the number of candidates for each\n      keypoint type, as it's possible to filter some candidates due to the score\n      threshold.\n  "
    (batch_size, _, width, num_keypoints) = _get_shape(keypoint_heatmap_predictions, 4)
    (keypoint_scores, y_indices, x_indices, channel_indices) = top_k_feature_map_locations(keypoint_heatmap_predictions, max_pool_kernel_size=max_pool_kernel_size, k=max_candidates, per_channel=True)
    peak_spatial_indices = flattened_indices_from_row_col_indices(y_indices, x_indices, width)
    y_indices = _to_float32(y_indices)
    x_indices = _to_float32(x_indices)
    offsets_flat = _flatten_spatial_dimensions(keypoint_heatmap_offsets)
    selected_offsets = tf.gather(offsets_flat, peak_spatial_indices, batch_dims=1)
    (_, num_indices, num_channels) = _get_shape(selected_offsets, 3)
    if (num_channels > 2):
        reshaped_offsets = tf.reshape(selected_offsets, [batch_size, num_indices, (- 1), 2])
        offsets = tf.gather(reshaped_offsets, channel_indices, batch_dims=2)
    else:
        offsets = selected_offsets
    (y_offsets, x_offsets) = tf.unstack(offsets, axis=2)
    keypoint_candidates = tf.stack([(y_indices + y_offsets), (x_indices + x_offsets)], axis=2)
    keypoint_candidates = tf.reshape(keypoint_candidates, [batch_size, num_keypoints, max_candidates, 2])
    keypoint_candidates = tf.transpose(keypoint_candidates, [0, 2, 1, 3])
    keypoint_scores = tf.reshape(keypoint_scores, [batch_size, num_keypoints, max_candidates])
    keypoint_scores = tf.transpose(keypoint_scores, [0, 2, 1])
    num_candidates = tf.reduce_sum(tf.to_int32((keypoint_scores >= keypoint_score_threshold)), axis=1)
    return (keypoint_candidates, keypoint_scores, num_candidates)