def refine_keypoints(regressed_keypoints, keypoint_candidates, keypoint_scores, num_keypoint_candidates, bboxes=None, unmatched_keypoint_score=0.1, box_scale=1.2, candidate_search_scale=0.3, candidate_ranking_mode='min_distance'):
    'Refines regressed keypoints by snapping to the nearest candidate keypoints.\n\n  The initial regressed keypoints represent a full set of keypoints regressed\n  from the centers of the objects. The keypoint candidates are estimated\n  independently from heatmaps, and are not associated with any object instances.\n  This function refines the regressed keypoints by "snapping" to the\n  nearest/highest score/highest score-distance ratio (depending on the\n  candidate_ranking_mode) candidate of the same keypoint type (e.g. "nose").\n  If no candidates are nearby, the regressed keypoint remains unchanged.\n\n  In order to snap a regressed keypoint to a candidate keypoint, the following\n  must be satisfied:\n  - the candidate keypoint must be of the same type as the regressed keypoint\n  - the candidate keypoint must not lie outside the predicted boxes (or the\n    boxes which encloses the regressed keypoints for the instance if `bboxes` is\n    not provided). Note that the box is scaled by\n    `regressed_box_scale` in height and width, to provide some margin around the\n    keypoints\n  - the distance to the closest candidate keypoint cannot exceed\n    candidate_search_scale * max(height, width), where height and width refer to\n    the bounding box for the instance.\n\n  Note that the same candidate keypoint is allowed to snap to regressed\n  keypoints in difference instances.\n\n  Args:\n    regressed_keypoints: A float tensor of shape\n      [batch_size, num_instances, num_keypoints, 2] with the initial regressed\n      keypoints.\n    keypoint_candidates: A tensor of shape\n      [batch_size, max_candidates, num_keypoints, 2] holding the location of\n      keypoint candidates in [y, x] format (expressed in absolute coordinates in\n      the output coordinate frame).\n    keypoint_scores: A float tensor of shape\n      [batch_size, max_candidates, num_keypoints] indicating the scores for\n      keypoint candidates.\n    num_keypoint_candidates: An integer tensor of shape\n      [batch_size, num_keypoints] indicating the number of valid candidates for\n      each keypoint type, as there may be padding (dim 1) of\n      `keypoint_candidates` and `keypoint_scores`.\n    bboxes: A tensor of shape [batch_size, num_instances, 4] with predicted\n      bounding boxes for each instance, expressed in the output coordinate\n      frame. If not provided, boxes will be computed from regressed keypoints.\n    unmatched_keypoint_score: float, the default score to use for regressed\n      keypoints that are not successfully snapped to a nearby candidate.\n    box_scale: float, the multiplier to expand the bounding boxes (either the\n      provided boxes or those which tightly cover the regressed keypoints) for\n      an instance. This scale is typically larger than 1.0 when not providing\n      `bboxes`.\n    candidate_search_scale: float, the scale parameter that multiplies the\n      largest dimension of a bounding box. The resulting distance becomes a\n      search radius for candidates in the vicinity of each regressed keypoint.\n    candidate_ranking_mode: A string as one of [\'min_distance\',\n     \'score_distance_ratio\'] indicating how to select the candidate. If invalid\n      value is provided, an ValueError will be raised.\n\n  Returns:\n    A tuple with:\n    refined_keypoints: A float tensor of shape\n      [batch_size, num_instances, num_keypoints, 2] with the final, refined\n      keypoints.\n    refined_scores: A float tensor of shape\n      [batch_size, num_instances, num_keypoints] with scores associated with all\n      instances and keypoints in `refined_keypoints`.\n\n  Raises:\n    ValueError: if provided candidate_ranking_mode is not one of\n      [\'min_distance\', \'score_distance_ratio\']\n  '
    (batch_size, num_instances, num_keypoints, _) = shape_utils.combined_static_and_dynamic_shape(regressed_keypoints)
    max_candidates = keypoint_candidates.shape[1]
    range_tiled = tf.tile(tf.reshape(tf.range(max_candidates), [1, max_candidates, 1]), [batch_size, 1, num_keypoints])
    num_candidates_tiled = tf.tile(tf.expand_dims(num_keypoint_candidates, 1), [1, max_candidates, 1])
    invalid_candidates = (range_tiled >= num_candidates_tiled)
    nan_mask = tf.where(invalid_candidates, (np.nan * tf.ones_like(invalid_candidates, dtype=tf.float32)), tf.ones_like(invalid_candidates, dtype=tf.float32))
    keypoint_candidates_with_nans = tf.math.multiply(keypoint_candidates, tf.expand_dims(nan_mask, (- 1)))
    regressed_keypoint_expanded = tf.expand_dims(regressed_keypoints, axis=2)
    keypoint_candidates_expanded = tf.expand_dims(keypoint_candidates_with_nans, axis=1)
    sqrd_distances = tf.math.reduce_sum(tf.math.squared_difference(regressed_keypoint_expanded, keypoint_candidates_expanded), axis=(- 1))
    distances = tf.math.sqrt(sqrd_distances)
    min_distances = tf.math.reduce_min(distances, axis=2)
    if (candidate_ranking_mode == 'min_distance'):
        nearby_candidate_inds = tf.math.argmin(distances, axis=2)
    elif (candidate_ranking_mode == 'score_distance_ratio'):
        tiled_keypoint_scores = tf.tile(tf.expand_dims(keypoint_scores, axis=1), multiples=[1, num_instances, 1, 1])
        ranking_scores = (tiled_keypoint_scores / (distances + 1e-06))
        nearby_candidate_inds = tf.math.argmax(ranking_scores, axis=2)
    else:
        raise ValueError(('Not recognized candidate_ranking_mode: %s' % candidate_ranking_mode))
    (nearby_candidate_coords, nearby_candidate_scores) = _gather_candidates_at_indices(keypoint_candidates, keypoint_scores, nearby_candidate_inds)
    if (bboxes is None):
        regressed_keypoints_flattened = tf.reshape(regressed_keypoints, [(- 1), num_keypoints, 2])
        bboxes_flattened = keypoint_ops.keypoints_to_enclosing_bounding_boxes(regressed_keypoints_flattened)
    else:
        bboxes_flattened = tf.reshape(bboxes, [(- 1), 4])
    boxlist = box_list.BoxList(bboxes_flattened)
    boxlist_scaled = box_list_ops.scale_height_width(boxlist, box_scale, box_scale)
    bboxes_scaled = boxlist_scaled.get()
    bboxes = tf.reshape(bboxes_scaled, [batch_size, num_instances, 4])
    bboxes_tiled = tf.tile(tf.expand_dims(bboxes, 2), [1, 1, num_keypoints, 1])
    (ymin, xmin, ymax, xmax) = tf.unstack(bboxes_tiled, axis=3)
    search_radius = (tf.math.maximum((ymax - ymin), (xmax - xmin)) * candidate_search_scale)
    mask = (((((tf.cast((nearby_candidate_coords[:, :, :, 0] < ymin), tf.int32) + tf.cast((nearby_candidate_coords[:, :, :, 0] > ymax), tf.int32)) + tf.cast((nearby_candidate_coords[:, :, :, 1] < xmin), tf.int32)) + tf.cast((nearby_candidate_coords[:, :, :, 1] > xmax), tf.int32)) + tf.cast((nearby_candidate_scores < unmatched_keypoint_score), tf.int32)) + tf.cast((min_distances > search_radius), tf.int32))
    mask = (mask > 0)
    refined_keypoints = tf.where(tf.tile(tf.expand_dims(mask, (- 1)), [1, 1, 1, 2]), regressed_keypoints, nearby_candidate_coords)
    refined_scores = tf.where(mask, (unmatched_keypoint_score * tf.ones_like(nearby_candidate_scores)), nearby_candidate_scores)
    return (refined_keypoints, refined_scores)
Refines regressed keypoints by snapping to the nearest candidate keypoints.

The initial regressed keypoints represent a full set of keypoints regressed
from the centers of the objects. The keypoint candidates are estimated
independently from heatmaps, and are not associated with any object instances.
This function refines the regressed keypoints by "snapping" to the
nearest/highest score/highest score-distance ratio (depending on the
candidate_ranking_mode) candidate of the same keypoint type (e.g. "nose").
If no candidates are nearby, the regressed keypoint remains unchanged.

In order to snap a regressed keypoint to a candidate keypoint, the following
must be satisfied:
- the candidate keypoint must be of the same type as the regressed keypoint
- the candidate keypoint must not lie outside the predicted boxes (or the
  boxes which encloses the regressed keypoints for the instance if `bboxes` is
  not provided). Note that the box is scaled by
  `regressed_box_scale` in height and width, to provide some margin around the
  keypoints
- the distance to the closest candidate keypoint cannot exceed
  candidate_search_scale * max(height, width), where height and width refer to
  the bounding box for the instance.

Note that the same candidate keypoint is allowed to snap to regressed
keypoints in difference instances.

Args:
  regressed_keypoints: A float tensor of shape
    [batch_size, num_instances, num_keypoints, 2] with the initial regressed
    keypoints.
  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] indicating the scores for
    keypoint candidates.
  num_keypoint_candidates: An integer tensor of shape
    [batch_size, num_keypoints] indicating the number of valid candidates for
    each keypoint type, as there may be padding (dim 1) of
    `keypoint_candidates` and `keypoint_scores`.
  bboxes: A tensor of shape [batch_size, num_instances, 4] with predicted
    bounding boxes for each instance, expressed in the output coordinate
    frame. If not provided, boxes will be computed from regressed keypoints.
  unmatched_keypoint_score: float, the default score to use for regressed
    keypoints that are not successfully snapped to a nearby candidate.
  box_scale: float, the multiplier to expand the bounding boxes (either the
    provided boxes or those which tightly cover the regressed keypoints) for
    an instance. This scale is typically larger than 1.0 when not providing
    `bboxes`.
  candidate_search_scale: float, the scale parameter that multiplies the
    largest dimension of a bounding box. The resulting distance becomes a
    search radius for candidates in the vicinity of each regressed keypoint.
  candidate_ranking_mode: A string as one of ['min_distance',
   'score_distance_ratio'] indicating how to select the candidate. If invalid
    value is provided, an ValueError will be raised.

Returns:
  A tuple with:
  refined_keypoints: A float tensor of shape
    [batch_size, num_instances, num_keypoints, 2] with the final, refined
    keypoints.
  refined_scores: A float tensor of shape
    [batch_size, num_instances, num_keypoints] with scores associated with all
    instances and keypoints in `refined_keypoints`.

Raises:
  ValueError: if provided candidate_ranking_mode is not one of
    ['min_distance', 'score_distance_ratio']
def refine_keypoints(regressed_keypoints, keypoint_candidates, keypoint_scores, num_keypoint_candidates, bboxes=None, unmatched_keypoint_score=0.1, box_scale=1.2, candidate_search_scale=0.3, candidate_ranking_mode='min_distance'):
    'Refines regressed keypoints by snapping to the nearest candidate keypoints.\n\n  The initial regressed keypoints represent a full set of keypoints regressed\n  from the centers of the objects. The keypoint candidates are estimated\n  independently from heatmaps, and are not associated with any object instances.\n  This function refines the regressed keypoints by "snapping" to the\n  nearest/highest score/highest score-distance ratio (depending on the\n  candidate_ranking_mode) candidate of the same keypoint type (e.g. "nose").\n  If no candidates are nearby, the regressed keypoint remains unchanged.\n\n  In order to snap a regressed keypoint to a candidate keypoint, the following\n  must be satisfied:\n  - the candidate keypoint must be of the same type as the regressed keypoint\n  - the candidate keypoint must not lie outside the predicted boxes (or the\n    boxes which encloses the regressed keypoints for the instance if `bboxes` is\n    not provided). Note that the box is scaled by\n    `regressed_box_scale` in height and width, to provide some margin around the\n    keypoints\n  - the distance to the closest candidate keypoint cannot exceed\n    candidate_search_scale * max(height, width), where height and width refer to\n    the bounding box for the instance.\n\n  Note that the same candidate keypoint is allowed to snap to regressed\n  keypoints in difference instances.\n\n  Args:\n    regressed_keypoints: A float tensor of shape\n      [batch_size, num_instances, num_keypoints, 2] with the initial regressed\n      keypoints.\n    keypoint_candidates: A tensor of shape\n      [batch_size, max_candidates, num_keypoints, 2] holding the location of\n      keypoint candidates in [y, x] format (expressed in absolute coordinates in\n      the output coordinate frame).\n    keypoint_scores: A float tensor of shape\n      [batch_size, max_candidates, num_keypoints] indicating the scores for\n      keypoint candidates.\n    num_keypoint_candidates: An integer tensor of shape\n      [batch_size, num_keypoints] indicating the number of valid candidates for\n      each keypoint type, as there may be padding (dim 1) of\n      `keypoint_candidates` and `keypoint_scores`.\n    bboxes: A tensor of shape [batch_size, num_instances, 4] with predicted\n      bounding boxes for each instance, expressed in the output coordinate\n      frame. If not provided, boxes will be computed from regressed keypoints.\n    unmatched_keypoint_score: float, the default score to use for regressed\n      keypoints that are not successfully snapped to a nearby candidate.\n    box_scale: float, the multiplier to expand the bounding boxes (either the\n      provided boxes or those which tightly cover the regressed keypoints) for\n      an instance. This scale is typically larger than 1.0 when not providing\n      `bboxes`.\n    candidate_search_scale: float, the scale parameter that multiplies the\n      largest dimension of a bounding box. The resulting distance becomes a\n      search radius for candidates in the vicinity of each regressed keypoint.\n    candidate_ranking_mode: A string as one of [\'min_distance\',\n     \'score_distance_ratio\'] indicating how to select the candidate. If invalid\n      value is provided, an ValueError will be raised.\n\n  Returns:\n    A tuple with:\n    refined_keypoints: A float tensor of shape\n      [batch_size, num_instances, num_keypoints, 2] with the final, refined\n      keypoints.\n    refined_scores: A float tensor of shape\n      [batch_size, num_instances, num_keypoints] with scores associated with all\n      instances and keypoints in `refined_keypoints`.\n\n  Raises:\n    ValueError: if provided candidate_ranking_mode is not one of\n      [\'min_distance\', \'score_distance_ratio\']\n  '
    (batch_size, num_instances, num_keypoints, _) = shape_utils.combined_static_and_dynamic_shape(regressed_keypoints)
    max_candidates = keypoint_candidates.shape[1]
    range_tiled = tf.tile(tf.reshape(tf.range(max_candidates), [1, max_candidates, 1]), [batch_size, 1, num_keypoints])
    num_candidates_tiled = tf.tile(tf.expand_dims(num_keypoint_candidates, 1), [1, max_candidates, 1])
    invalid_candidates = (range_tiled >= num_candidates_tiled)
    nan_mask = tf.where(invalid_candidates, (np.nan * tf.ones_like(invalid_candidates, dtype=tf.float32)), tf.ones_like(invalid_candidates, dtype=tf.float32))
    keypoint_candidates_with_nans = tf.math.multiply(keypoint_candidates, tf.expand_dims(nan_mask, (- 1)))
    regressed_keypoint_expanded = tf.expand_dims(regressed_keypoints, axis=2)
    keypoint_candidates_expanded = tf.expand_dims(keypoint_candidates_with_nans, axis=1)
    sqrd_distances = tf.math.reduce_sum(tf.math.squared_difference(regressed_keypoint_expanded, keypoint_candidates_expanded), axis=(- 1))
    distances = tf.math.sqrt(sqrd_distances)
    min_distances = tf.math.reduce_min(distances, axis=2)
    if (candidate_ranking_mode == 'min_distance'):
        nearby_candidate_inds = tf.math.argmin(distances, axis=2)
    elif (candidate_ranking_mode == 'score_distance_ratio'):
        tiled_keypoint_scores = tf.tile(tf.expand_dims(keypoint_scores, axis=1), multiples=[1, num_instances, 1, 1])
        ranking_scores = (tiled_keypoint_scores / (distances + 1e-06))
        nearby_candidate_inds = tf.math.argmax(ranking_scores, axis=2)
    else:
        raise ValueError(('Not recognized candidate_ranking_mode: %s' % candidate_ranking_mode))
    (nearby_candidate_coords, nearby_candidate_scores) = _gather_candidates_at_indices(keypoint_candidates, keypoint_scores, nearby_candidate_inds)
    if (bboxes is None):
        regressed_keypoints_flattened = tf.reshape(regressed_keypoints, [(- 1), num_keypoints, 2])
        bboxes_flattened = keypoint_ops.keypoints_to_enclosing_bounding_boxes(regressed_keypoints_flattened)
    else:
        bboxes_flattened = tf.reshape(bboxes, [(- 1), 4])
    boxlist = box_list.BoxList(bboxes_flattened)
    boxlist_scaled = box_list_ops.scale_height_width(boxlist, box_scale, box_scale)
    bboxes_scaled = boxlist_scaled.get()
    bboxes = tf.reshape(bboxes_scaled, [batch_size, num_instances, 4])
    bboxes_tiled = tf.tile(tf.expand_dims(bboxes, 2), [1, 1, num_keypoints, 1])
    (ymin, xmin, ymax, xmax) = tf.unstack(bboxes_tiled, axis=3)
    search_radius = (tf.math.maximum((ymax - ymin), (xmax - xmin)) * candidate_search_scale)
    mask = (((((tf.cast((nearby_candidate_coords[:, :, :, 0] < ymin), tf.int32) + tf.cast((nearby_candidate_coords[:, :, :, 0] > ymax), tf.int32)) + tf.cast((nearby_candidate_coords[:, :, :, 1] < xmin), tf.int32)) + tf.cast((nearby_candidate_coords[:, :, :, 1] > xmax), tf.int32)) + tf.cast((nearby_candidate_scores < unmatched_keypoint_score), tf.int32)) + tf.cast((min_distances > search_radius), tf.int32))
    mask = (mask > 0)
    refined_keypoints = tf.where(tf.tile(tf.expand_dims(mask, (- 1)), [1, 1, 1, 2]), regressed_keypoints, nearby_candidate_coords)
    refined_scores = tf.where(mask, (unmatched_keypoint_score * tf.ones_like(nearby_candidate_scores)), nearby_candidate_scores)
    return (refined_keypoints, refined_scores)