def prune_outside_window(dp_num_points, dp_part_ids, dp_surface_coords, window,
                         scope=None):
  """Prunes DensePose points that fall outside a given window.

  This function replaces points that fall outside the given window with zeros.
  See also clip_to_window which clips any DensePose points that fall outside the
  given window.

  Note that this operation uses dynamic shapes, and therefore is not currently
  suitable for TPU.

  Args:
    dp_num_points: a tensor of shape [num_instances] that indicates how many
      (non-padded) DensePose points there are per instance.
    dp_part_ids: a tensor of shape [num_instances, num_points] with DensePose
      part ids. These part_ids are 0-indexed, where the first non-background
      part has index 0.
    dp_surface_coords: a tensor of shape [num_instances, num_points, 4] with
      DensePose surface coordinates in (y, x, v, u) format.
    window: a tensor of shape [4] representing the [y_min, x_min, y_max, x_max]
      window outside of which the op should prune the points.
    scope: name scope.

  Returns:
    new_dp_num_points: a tensor of shape [num_instances] that indicates how many
      (non-padded) DensePose points there are per instance after pruning.
    new_dp_part_ids: a tensor of shape [num_instances, num_points] with
      DensePose part ids. These part_ids are 0-indexed, where the first
      non-background part has index 0.
    new_dp_surface_coords: a tensor of shape [num_instances, num_points, 4] with
      DensePose surface coordinates after pruning.
  """
  with tf.name_scope(scope, 'DensePosePruneOutsideWindow'):
    y, x, _, _ = tf.unstack(dp_surface_coords, axis=-1)
    win_y_min, win_x_min, win_y_max, win_x_max = tf.unstack(window)

    num_instances, num_points = shape_utils.combined_static_and_dynamic_shape(
        dp_part_ids)
    dp_num_points_tiled = tf.tile(dp_num_points[:, tf.newaxis],
                                  multiples=[1, num_points])
    range_tiled = tf.tile(tf.range(num_points)[tf.newaxis, :],
                          multiples=[num_instances, 1])
    valid_initial = range_tiled < dp_num_points_tiled
    valid_in_window = tf.logical_and(
        tf.logical_and(y >= win_y_min, y <= win_y_max),
        tf.logical_and(x >= win_x_min, x <= win_x_max))
    valid_indices = tf.logical_and(valid_initial, valid_in_window)

    new_dp_num_points = tf.math.reduce_sum(
        tf.cast(valid_indices, tf.int32), axis=1)
    max_num_points = tf.math.reduce_max(new_dp_num_points)

    def gather_and_reshuffle(elems):
      dp_part_ids, dp_surface_coords, valid_indices = elems
      locs = tf.where(valid_indices)[:, 0]
      valid_part_ids = tf.gather(dp_part_ids, locs, axis=0)
      valid_part_ids_padded = shape_utils.pad_or_clip_nd(
          valid_part_ids, output_shape=[max_num_points])
      valid_surface_coords = tf.gather(dp_surface_coords, locs, axis=0)
      valid_surface_coords_padded = shape_utils.pad_or_clip_nd(
          valid_surface_coords, output_shape=[max_num_points, 4])
      return [valid_part_ids_padded, valid_surface_coords_padded]

    new_dp_part_ids, new_dp_surface_coords = (
        shape_utils.static_or_dynamic_map_fn(
            gather_and_reshuffle,
            elems=[dp_part_ids, dp_surface_coords, valid_indices],
            dtype=[tf.int32, tf.float32],
            back_prop=False))
    return new_dp_num_points, new_dp_part_ids, new_dp_surface_coords
Prunes DensePose points that fall outside a given window.

  This function replaces points that fall outside the given window with zeros.
  See also clip_to_window which clips any DensePose points that fall outside the
  given window.

  Note that this operation uses dynamic shapes, and therefore is not currently
  suitable for TPU.

  Args:
    dp_num_points: a tensor of shape [num_instances] that indicates how many
      (non-padded) DensePose points there are per instance.
    dp_part_ids: a tensor of shape [num_instances, num_points] with DensePose
      part ids. These part_ids are 0-indexed, where the first non-background
      part has index 0.
    dp_surface_coords: a tensor of shape [num_instances, num_points, 4] with
      DensePose surface coordinates in (y, x, v, u) format.
    window: a tensor of shape [4] representing the [y_min, x_min, y_max, x_max]
      window outside of which the op should prune the points.
    scope: name scope.

  Returns:
    new_dp_num_points: a tensor of shape [num_instances] that indicates how many
      (non-padded) DensePose points there are per instance after pruning.
    new_dp_part_ids: a tensor of shape [num_instances, num_points] with
      DensePose part ids. These part_ids are 0-indexed, where the first
      non-background part has index 0.
    new_dp_surface_coords: a tensor of shape [num_instances, num_points, 4] with
      DensePose surface coordinates after pruning.
  
inference_api/src/main/object_detection/core/densepose_ops.py