def crop_and_threshold_masks(elems, input_height, input_width, mask_height=256, mask_width=256, score_threshold=0.5, densepose_class_index=(- 1)):
    'Crops and thresholds masks based on detection boxes.\n\n  Args:\n    elems: A tuple of\n      boxes - float32 tensor of shape [max_detections, 4]\n      classes - int32 tensor of shape [max_detections] (0-indexed)\n      masks - float32 tensor of shape [output_height, output_width, num_classes]\n      part_heatmap - float32 tensor of shape [output_height, output_width,\n        num_parts]\n      surf_coords - float32 tensor of shape [output_height, output_width,\n        2 * num_parts]\n      true_height - scalar int tensor\n      true_width - scalar int tensor\n    input_height: Input height to network.\n    input_width: Input width to network.\n    mask_height: Height for resizing mask crops.\n    mask_width: Width for resizing mask crops.\n    score_threshold: The threshold at which to convert predicted mask\n      into foreground pixels.\n    densepose_class_index: scalar int tensor with the class index (0-indexed)\n      for DensePose.\n\n  Returns:\n    A tuple of\n    all_instances: A [max_detections, mask_height, mask_width] uint8 tensor\n      with a predicted foreground mask for each instance. Background is encoded\n      as 0, and foreground is encoded as a positive integer. Specific part\n      indices are encoded as 1-indexed parts (for classes that have part\n      information).\n    surface_coords: A [max_detections, mask_height, mask_width, 2]\n      float32 tensor with (v, u) coordinates. for each part.\n  '
    (boxes, classes, masks, part_heatmap, surf_coords, true_height, true_width) = elems
    boxlist = box_list.BoxList(boxes)
    y_scale = (true_height / input_height)
    x_scale = (true_width / input_width)
    boxlist = box_list_ops.scale(boxlist, y_scale, x_scale)
    boxes = boxlist.get()
    num_classes = tf.shape(masks)[(- 1)]
    masks_4d = tf.transpose(masks, perm=[2, 0, 1])[:, :, :, tf.newaxis]
    part_heatmap_4d = tf.tile(part_heatmap[tf.newaxis, :, :, :], multiples=[num_classes, 1, 1, 1])
    surf_coords_4d = tf.tile(surf_coords[tf.newaxis, :, :, :], multiples=[num_classes, 1, 1, 1])
    feature_maps_concat = tf.concat([masks_4d, part_heatmap_4d, surf_coords_4d], axis=(- 1))
    cropped_masks = tf2.image.crop_and_resize(feature_maps_concat, boxes=boxes, box_indices=classes, crop_size=[mask_height, mask_width], method='bilinear')
    num_parts = tf.shape(part_heatmap)[(- 1)]
    (instance_masks, part_heatmap_cropped, surface_coords_cropped) = tf.split(cropped_masks, [1, num_parts, (2 * num_parts)], axis=(- 1))
    instance_masks_int = tf.cast(tf.math.greater_equal(instance_masks, score_threshold), dtype=tf.int32)
    det_with_parts = tf.equal(classes, densepose_class_index)
    det_with_parts = tf.cast(tf.reshape(det_with_parts, [(- 1), 1, 1, 1]), dtype=tf.int32)
    instance_masks_with_parts = tf.math.multiply(instance_masks_int, det_with_parts)
    det_without_parts = (1 - det_with_parts)
    instance_masks_without_parts = tf.math.multiply(instance_masks_int, det_without_parts)
    part_mask_int_zero_indexed = tf.math.argmax(part_heatmap_cropped, axis=(- 1), output_type=tf.int32)[:, :, :, tf.newaxis]
    part_mask_int_one_indexed = (part_mask_int_zero_indexed + 1)
    all_instances = (instance_masks_without_parts + (instance_masks_with_parts * part_mask_int_one_indexed))
    surface_coords_cropped = tf.reshape(surface_coords_cropped, [(- 1), mask_height, mask_width, num_parts, 2])
    surface_coords = gather_surface_coords_for_parts(surface_coords_cropped, part_mask_int_zero_indexed)
    surface_coords = (surface_coords * tf.cast(instance_masks_with_parts, tf.float32))
    return [tf.squeeze(all_instances, axis=3), surface_coords]
Crops and thresholds masks based on detection boxes.

Args:
  elems: A tuple of
    boxes - float32 tensor of shape [max_detections, 4]
    classes - int32 tensor of shape [max_detections] (0-indexed)
    masks - float32 tensor of shape [output_height, output_width, num_classes]
    part_heatmap - float32 tensor of shape [output_height, output_width,
      num_parts]
    surf_coords - float32 tensor of shape [output_height, output_width,
      2 * num_parts]
    true_height - scalar int tensor
    true_width - scalar int tensor
  input_height: Input height to network.
  input_width: Input width to network.
  mask_height: Height for resizing mask crops.
  mask_width: Width for resizing mask crops.
  score_threshold: The threshold at which to convert predicted mask
    into foreground pixels.
  densepose_class_index: scalar int tensor with the class index (0-indexed)
    for DensePose.

Returns:
  A tuple of
  all_instances: A [max_detections, mask_height, mask_width] uint8 tensor
    with a predicted foreground mask for each instance. Background is encoded
    as 0, and foreground is encoded as a positive integer. Specific part
    indices are encoded as 1-indexed parts (for classes that have part
    information).
  surface_coords: A [max_detections, mask_height, mask_width, 2]
    float32 tensor with (v, u) coordinates. for each part.
def crop_and_threshold_masks(elems, input_height, input_width, mask_height=256, mask_width=256, score_threshold=0.5, densepose_class_index=(- 1)):
    'Crops and thresholds masks based on detection boxes.\n\n  Args:\n    elems: A tuple of\n      boxes - float32 tensor of shape [max_detections, 4]\n      classes - int32 tensor of shape [max_detections] (0-indexed)\n      masks - float32 tensor of shape [output_height, output_width, num_classes]\n      part_heatmap - float32 tensor of shape [output_height, output_width,\n        num_parts]\n      surf_coords - float32 tensor of shape [output_height, output_width,\n        2 * num_parts]\n      true_height - scalar int tensor\n      true_width - scalar int tensor\n    input_height: Input height to network.\n    input_width: Input width to network.\n    mask_height: Height for resizing mask crops.\n    mask_width: Width for resizing mask crops.\n    score_threshold: The threshold at which to convert predicted mask\n      into foreground pixels.\n    densepose_class_index: scalar int tensor with the class index (0-indexed)\n      for DensePose.\n\n  Returns:\n    A tuple of\n    all_instances: A [max_detections, mask_height, mask_width] uint8 tensor\n      with a predicted foreground mask for each instance. Background is encoded\n      as 0, and foreground is encoded as a positive integer. Specific part\n      indices are encoded as 1-indexed parts (for classes that have part\n      information).\n    surface_coords: A [max_detections, mask_height, mask_width, 2]\n      float32 tensor with (v, u) coordinates. for each part.\n  '
    (boxes, classes, masks, part_heatmap, surf_coords, true_height, true_width) = elems
    boxlist = box_list.BoxList(boxes)
    y_scale = (true_height / input_height)
    x_scale = (true_width / input_width)
    boxlist = box_list_ops.scale(boxlist, y_scale, x_scale)
    boxes = boxlist.get()
    num_classes = tf.shape(masks)[(- 1)]
    masks_4d = tf.transpose(masks, perm=[2, 0, 1])[:, :, :, tf.newaxis]
    part_heatmap_4d = tf.tile(part_heatmap[tf.newaxis, :, :, :], multiples=[num_classes, 1, 1, 1])
    surf_coords_4d = tf.tile(surf_coords[tf.newaxis, :, :, :], multiples=[num_classes, 1, 1, 1])
    feature_maps_concat = tf.concat([masks_4d, part_heatmap_4d, surf_coords_4d], axis=(- 1))
    cropped_masks = tf2.image.crop_and_resize(feature_maps_concat, boxes=boxes, box_indices=classes, crop_size=[mask_height, mask_width], method='bilinear')
    num_parts = tf.shape(part_heatmap)[(- 1)]
    (instance_masks, part_heatmap_cropped, surface_coords_cropped) = tf.split(cropped_masks, [1, num_parts, (2 * num_parts)], axis=(- 1))
    instance_masks_int = tf.cast(tf.math.greater_equal(instance_masks, score_threshold), dtype=tf.int32)
    det_with_parts = tf.equal(classes, densepose_class_index)
    det_with_parts = tf.cast(tf.reshape(det_with_parts, [(- 1), 1, 1, 1]), dtype=tf.int32)
    instance_masks_with_parts = tf.math.multiply(instance_masks_int, det_with_parts)
    det_without_parts = (1 - det_with_parts)
    instance_masks_without_parts = tf.math.multiply(instance_masks_int, det_without_parts)
    part_mask_int_zero_indexed = tf.math.argmax(part_heatmap_cropped, axis=(- 1), output_type=tf.int32)[:, :, :, tf.newaxis]
    part_mask_int_one_indexed = (part_mask_int_zero_indexed + 1)
    all_instances = (instance_masks_without_parts + (instance_masks_with_parts * part_mask_int_one_indexed))
    surface_coords_cropped = tf.reshape(surface_coords_cropped, [(- 1), mask_height, mask_width, num_parts, 2])
    surface_coords = gather_surface_coords_for_parts(surface_coords_cropped, part_mask_int_zero_indexed)
    surface_coords = (surface_coords * tf.cast(instance_masks_with_parts, tf.float32))
    return [tf.squeeze(all_instances, axis=3), surface_coords]