def draw_densepose_visualizations()

in object-detection/visualization_utils.py [0:0]


def draw_densepose_visualizations(eval_dict,
                                  max_boxes_to_draw=20,
                                  min_score_thresh=0.2,
                                  num_parts=24,
                                  dp_coord_to_visualize=0):
  """Draws DensePose visualizations.
  Args:
    eval_dict: The evaluation dictionary returned by
      eval_util.result_dict_for_batched_example().
    max_boxes_to_draw: The maximum number of boxes to draw for detections.
    min_score_thresh: The minimum score threshold for showing detections.
    num_parts: The number of different densepose parts.
    dp_coord_to_visualize: Whether to visualize v-coordinates (0) or
      u-coordinates (0) overlaid on the person masks.
  Returns:
    A list of [1, H, W, C] uint8 tensor, each element corresponding to an image
    in the batch.
  Raises:
    ValueError: If `dp_coord_to_visualize` is not 0 or 1.
  """
  if dp_coord_to_visualize not in (0, 1):
    raise ValueError('`dp_coord_to_visualize` must be either 0 for v '
                     'coordinates), or 1 for u coordinates, but instead got '
                     '{}'.format(dp_coord_to_visualize))
  detection_fields = fields.DetectionResultFields()
  input_data_fields = fields.InputDataFields()

  if detection_fields.detection_masks not in eval_dict:
    raise ValueError('Expected `detection_masks` in `eval_dict`.')
  if detection_fields.detection_surface_coords not in eval_dict:
    raise ValueError('Expected `detection_surface_coords` in `eval_dict`.')

  images_with_detections_list = []
  for indx in range(eval_dict[input_data_fields.original_image].shape[0]):
    # Note that detection masks have already been resized to the original image
    # shapes, but `original_image` has not.
    # TODO(ronnyvotel): Consider resizing `original_image` in
    # eval_util.result_dict_for_batched_example().
    true_shape = eval_dict[input_data_fields.true_image_shape][indx]
    original_shape = eval_dict[
        input_data_fields.original_image_spatial_shape][indx]
    image = eval_dict[input_data_fields.original_image][indx]
    image = shape_utils.pad_or_clip_nd(image, [true_shape[0], true_shape[1], 3])
    image = _resize_original_image(image, original_shape)

    scores = eval_dict[detection_fields.detection_scores][indx]
    detection_masks = eval_dict[detection_fields.detection_masks][indx]
    surface_coords = eval_dict[detection_fields.detection_surface_coords][indx]

    def draw_densepose_py_func(image, detection_masks, surface_coords, scores):
      """Overlays part masks and surface coords on original images."""
      surface_coord_image = np.copy(image)
      for i, (score, surface_coord, mask) in enumerate(
          zip(scores, surface_coords, detection_masks)):
        if i == max_boxes_to_draw:
          break
        if score > min_score_thresh:
          draw_part_mask_on_image_array(image, mask, num_parts=num_parts)
          draw_float_channel_on_image_array(
              surface_coord_image, surface_coord[:, :, dp_coord_to_visualize],
              mask)
      return np.concatenate([image, surface_coord_image], axis=1)

    image_with_densepose = tf.py_func(
        draw_densepose_py_func,
        [image, detection_masks, surface_coords, scores],
        tf.uint8)
    images_with_detections_list.append(
        image_with_densepose[tf.newaxis, :, :, :])
  return images_with_detections_list