Open source MnasFPN and minor fixes to OD API (#8484)

310447280  by lzc:

    Internal change

310420845  by Zhichao Lu:

    Open source the internal Context RCNN code.

--
310362339  by Zhichao Lu:

    Internal change

310259448  by lzc:

    Update required TF version for OD API.

--
310252159  by Zhichao Lu:

    Port patch_ops_test to TF1/TF2 as TPUs.

--
310247180  by Zhichao Lu:

    Ignore keypoint heatmap loss in the regions/bounding boxes with target keypoint
    class but no valid keypoint annotations.

--
310178294  by Zhichao Lu:

    Opensource MnasFPN
    https://arxiv.org/abs/1912.01106

--
310094222  by lzc:

    Internal changes.

--
310085250  by lzc:

    Internal Change.

--
310016447  by huizhongc:

    Remove unrecognized classes from labeled_classes.

--
310009470  by rathodv:

    Mark batcher.py as TF1 only.

--
310001984  by rathodv:

    Update core/preprocessor.py to be compatible with TF1/TF2..

--
309455035  by Zhi...

research/object_detection/utils/visualization_utils.py

@@ -38,6 +38,7 @@ from six.moves import range
 from six.moves import zip
 import tensorflow as tf
 
+from object_detection.core import keypoint_ops
 from object_detection.core import standard_fields as fields
 from object_detection.utils import shape_utils
 
@@ -202,8 +203,11 @@ def draw_bounding_box_on_image(image,
                                   ymin * im_height, ymax * im_height)
   else:
     (left, right, top, bottom) = (xmin, xmax, ymin, ymax)
-  draw.line([(left, top), (left, bottom), (right, bottom),
-             (right, top), (left, top)], width=thickness, fill=color)
+  if thickness > 0:
+    draw.line([(left, top), (left, bottom), (right, bottom), (right, top),
+               (left, top)],
+              width=thickness,
+              fill=color)
   try:
     font = ImageFont.truetype('arial.ttf', 24)
   except IOError:
@@ -299,8 +303,11 @@ def draw_bounding_boxes_on_image(image,
                                boxes[i, 3], color, thickness, display_str_list)
 
 
-def create_visualization_fn(category_index, include_masks=False,
-                            include_keypoints=False, include_track_ids=False,
+def create_visualization_fn(category_index,
+                            include_masks=False,
+                            include_keypoints=False,
+                            include_keypoint_scores=False,
+                            include_track_ids=False,
                             **kwargs):
   """Constructs a visualization function that can be wrapped in a py_func.
 
@@ -311,9 +318,10 @@ def create_visualization_fn(category_index, include_masks=False,
   1: boxes
   2: classes
   3: scores
-  [4-6]: masks (optional)
-  [4-6]: keypoints (optional)
-  [4-6]: track_ids (optional)
+  [4]: masks (optional)
+  [4-5]: keypoints (optional)
+  [4-6]: keypoint_scores (optional)
+  [4-7]: track_ids (optional)
 
   -- Example 1 --
   vis_only_masks_fn = create_visualization_fn(category_index,
@@ -338,6 +346,8 @@ def create_visualization_fn(category_index, include_masks=False,
       the returned function.
     include_keypoints: Whether keypoints should be expected as a positional
       argument in the returned function.
+    include_keypoint_scores: Whether keypoint scores should be expected as a
+      positional argument in the returned function.
     include_track_ids: Whether track ids should be expected as a positional
       argument in the returned function.
     **kwargs: Additional kwargs that will be passed to
@@ -359,6 +369,7 @@ def create_visualization_fn(category_index, include_masks=False,
         -- Optional positional arguments --
         instance_masks - a numpy array of shape [N, image_height, image_width].
         keypoints - a numpy array of shape [N, num_keypoints, 2].
+        keypoint_scores - a numpy array of shape [N, num_keypoints].
         track_ids - a numpy array of shape [N] with unique track ids.
 
     Returns:
@@ -369,7 +380,7 @@ def create_visualization_fn(category_index, include_masks=False,
     boxes = args[1]
     classes = args[2]
     scores = args[3]
-    masks = keypoints = track_ids = None
+    masks = keypoints = keypoint_scores = track_ids = None
     pos_arg_ptr = 4  # Positional argument for first optional tensor (masks).
     if include_masks:
       masks = args[pos_arg_ptr]
@@ -377,6 +388,9 @@ def create_visualization_fn(category_index, include_masks=False,
     if include_keypoints:
       keypoints = args[pos_arg_ptr]
       pos_arg_ptr += 1
+    if include_keypoint_scores:
+      keypoint_scores = args[pos_arg_ptr]
+      pos_arg_ptr += 1
     if include_track_ids:
       track_ids = args[pos_arg_ptr]
 
@@ -388,11 +402,106 @@ def create_visualization_fn(category_index, include_masks=False,
         category_index=category_index,
         instance_masks=masks,
         keypoints=keypoints,
+        keypoint_scores=keypoint_scores,
         track_ids=track_ids,
         **kwargs)
   return visualization_py_func_fn
 
 
+def draw_heatmaps_on_image(image, heatmaps):
+  """Draws heatmaps on an image.
+
+  The heatmaps are handled channel by channel and different colors are used to
+  paint different heatmap channels.
+
+  Args:
+    image: a PIL.Image object.
+    heatmaps: a numpy array with shape [image_height, image_width, channel].
+      Note that the image_height and image_width should match the size of input
+      image.
+  """
+  draw = ImageDraw.Draw(image)
+  channel = heatmaps.shape[2]
+  for c in range(channel):
+    heatmap = heatmaps[:, :, c] * 255
+    heatmap = heatmap.astype('uint8')
+    bitmap = Image.fromarray(heatmap, 'L')
+    bitmap.convert('1')
+    draw.bitmap(
+        xy=[(0, 0)],
+        bitmap=bitmap,
+        fill=STANDARD_COLORS[c])
+
+
+def draw_heatmaps_on_image_array(image, heatmaps):
+  """Overlays heatmaps to an image (numpy array).
+
+  The function overlays the heatmaps on top of image. The heatmap values will be
+  painted with different colors depending on the channels. Similar to
+  "draw_heatmaps_on_image_array" function except the inputs are numpy arrays.
+
+  Args:
+    image: a numpy array with shape [height, width, 3].
+    heatmaps: a numpy array with shape [height, width, channel].
+
+  Returns:
+    An uint8 numpy array representing the input image painted with heatmap
+    colors.
+  """
+  if not isinstance(image, np.ndarray):
+    image = image.numpy()
+  if not isinstance(heatmaps, np.ndarray):
+    heatmaps = heatmaps.numpy()
+  image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
+  draw_heatmaps_on_image(image_pil, heatmaps)
+  return np.array(image_pil)
+
+
+def draw_heatmaps_on_image_tensors(images,
+                                   heatmaps,
+                                   apply_sigmoid=False):
+  """Draws heatmaps on batch of image tensors.
+
+  Args:
+    images: A 4D uint8 image tensor of shape [N, H, W, C]. If C > 3, additional
+      channels will be ignored. If C = 1, then we convert the images to RGB
+      images.
+    heatmaps: [N, h, w, channel] float32 tensor of heatmaps. Note that the
+      heatmaps will be resized to match the input image size before overlaying
+      the heatmaps with input images. Theoretically the heatmap height width
+      should have the same aspect ratio as the input image to avoid potential
+      misalignment introduced by the image resize.
+    apply_sigmoid: Whether to apply a sigmoid layer on top of the heatmaps. If
+      the heatmaps come directly from the prediction logits, then we should
+      apply the sigmoid layer to make sure the values are in between [0.0, 1.0].
+
+  Returns:
+    4D image tensor of type uint8, with heatmaps overlaid on top.
+  """
+  # Additional channels are being ignored.
+  if images.shape[3] > 3:
+    images = images[:, :, :, 0:3]
+  elif images.shape[3] == 1:
+    images = tf.image.grayscale_to_rgb(images)
+
+  _, height, width, _ = shape_utils.combined_static_and_dynamic_shape(images)
+  if apply_sigmoid:
+    heatmaps = tf.math.sigmoid(heatmaps)
+  resized_heatmaps = tf.image.resize(heatmaps, size=[height, width])
+
+  elems = [images, resized_heatmaps]
+
+  def draw_heatmaps(image_and_heatmaps):
+    """Draws heatmaps on image."""
+    image_with_heatmaps = tf.py_function(
+        draw_heatmaps_on_image_array,
+        image_and_heatmaps,
+        tf.uint8)
+    return image_with_heatmaps
+  images = tf.map_fn(draw_heatmaps, elems, dtype=tf.uint8, back_prop=False)
+  return images
+
+
 def _resize_original_image(image, image_shape):
   image = tf.expand_dims(image, 0)
   image = tf.image.resize_images(
@@ -412,6 +521,8 @@ def draw_bounding_boxes_on_image_tensors(images,
                                          true_image_shape=None,
                                          instance_masks=None,
                                          keypoints=None,
+                                         keypoint_scores=None,
+                                         keypoint_edges=None,
                                          track_ids=None,
                                          max_boxes_to_draw=20,
                                          min_score_thresh=0.2,
@@ -436,6 +547,11 @@ def draw_bounding_boxes_on_image_tensors(images,
       instance masks.
     keypoints: A 4D float32 tensor of shape [N, max_detection, num_keypoints, 2]
       with keypoints.
+    keypoint_scores: A 3D float32 tensor of shape [N, max_detection,
+      num_keypoints] with keypoint scores.
+    keypoint_edges: A list of tuples with keypoint indices that specify which
+      keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
+      edges from keypoint 0 to 1 and from keypoint 2 to 4.
     track_ids: [N, max_detections] int32 tensor of unique tracks ids (i.e.
       instance ids for each object). If provided, the color-coding of boxes is
       dictated by these ids, and not classes.
@@ -458,7 +574,8 @@ def draw_bounding_boxes_on_image_tensors(images,
       'max_boxes_to_draw': max_boxes_to_draw,
       'min_score_thresh': min_score_thresh,
       'agnostic_mode': False,
-      'line_thickness': 4
+      'line_thickness': 4,
+      'keypoint_edges': keypoint_edges
   }
   if true_image_shape is None:
     true_shapes = tf.constant(-1, shape=[images.shape.as_list()[0], 3])
@@ -473,6 +590,7 @@ def draw_bounding_boxes_on_image_tensors(images,
       category_index,
       include_masks=instance_masks is not None,
       include_keypoints=keypoints is not None,
+      include_keypoint_scores=keypoint_scores is not None,
       include_track_ids=track_ids is not None,
       **visualization_keyword_args)
 
@@ -481,6 +599,8 @@ def draw_bounding_boxes_on_image_tensors(images,
     elems.append(instance_masks)
   if keypoints is not None:
     elems.append(keypoints)
+  if keypoint_scores is not None:
+    elems.append(keypoint_scores)
   if track_ids is not None:
     elems.append(track_ids)
 
@@ -506,7 +626,8 @@ def draw_side_by_side_evaluation_image(eval_dict,
                                        category_index,
                                        max_boxes_to_draw=20,
                                        min_score_thresh=0.2,
-                                       use_normalized_coordinates=True):
+                                       use_normalized_coordinates=True,
+                                       keypoint_edges=None):
   """Creates a side-by-side image with detections and groundtruth.
 
   Bounding boxes (and instance masks, if available) are visualized on both
@@ -519,9 +640,12 @@ def draw_side_by_side_evaluation_image(eval_dict,
     category_index: A category index (dictionary) produced from a labelmap.
     max_boxes_to_draw: The maximum number of boxes to draw for detections.
     min_score_thresh: The minimum score threshold for showing detections.
-    use_normalized_coordinates: Whether to assume boxes and kepoints are in
-      normalized coordinates (as opposed to absolute coordiantes).
+    use_normalized_coordinates: Whether to assume boxes and keypoints are in
+      normalized coordinates (as opposed to absolute coordinates).
       Default is True.
+    keypoint_edges: A list of tuples with keypoint indices that specify which
+      keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
+      edges from keypoint 0 to 1 and from keypoint 2 to 4.
 
   Returns:
     A list of [1, H, 2 * W, C] uint8 tensor. The subimage on the left
@@ -536,7 +660,8 @@ def draw_side_by_side_evaluation_image(eval_dict,
   # Add the batch dimension if the eval_dict is for single example.
   if len(eval_dict[detection_fields.detection_classes].shape) == 1:
     for key in eval_dict:
-      if key != input_data_fields.original_image and key != input_data_fields.image_additional_channels:
+      if (key != input_data_fields.original_image and
+          key != input_data_fields.image_additional_channels):
         eval_dict[key] = tf.expand_dims(eval_dict[key], 0)
 
   for indx in range(eval_dict[input_data_fields.original_image].shape[0]):
@@ -547,15 +672,36 @@ def draw_side_by_side_evaluation_image(eval_dict,
               eval_dict[detection_fields.detection_masks][indx], axis=0),
           tf.uint8)
     keypoints = None
+    keypoint_scores = None
     if detection_fields.detection_keypoints in eval_dict:
       keypoints = tf.expand_dims(
           eval_dict[detection_fields.detection_keypoints][indx], axis=0)
+      if detection_fields.detection_keypoint_scores in eval_dict:
+        keypoint_scores = tf.expand_dims(
+            eval_dict[detection_fields.detection_keypoint_scores][indx], axis=0)
+      else:
+        keypoint_scores = tf.cast(keypoint_ops.set_keypoint_visibilities(
+            keypoints), dtype=tf.float32)
+
     groundtruth_instance_masks = None
     if input_data_fields.groundtruth_instance_masks in eval_dict:
       groundtruth_instance_masks = tf.cast(
           tf.expand_dims(
               eval_dict[input_data_fields.groundtruth_instance_masks][indx],
               axis=0), tf.uint8)
+    groundtruth_keypoints = None
+    groundtruth_keypoint_scores = None
+    gt_kpt_vis_fld = input_data_fields.groundtruth_keypoint_visibilities
+    if input_data_fields.groundtruth_keypoints in eval_dict:
+      groundtruth_keypoints = tf.expand_dims(
+          eval_dict[input_data_fields.groundtruth_keypoints][indx], axis=0)
+      if gt_kpt_vis_fld in eval_dict:
+        groundtruth_keypoint_scores = tf.expand_dims(
+            tf.cast(eval_dict[gt_kpt_vis_fld][indx], dtype=tf.float32), axis=0)
+      else:
+        groundtruth_keypoint_scores = tf.cast(
+            keypoint_ops.set_keypoint_visibilities(
+                groundtruth_keypoints), dtype=tf.float32)
 
     images_with_detections = draw_bounding_boxes_on_image_tensors(
         tf.expand_dims(
@@ -574,6 +720,8 @@ def draw_side_by_side_evaluation_image(eval_dict,
             eval_dict[input_data_fields.true_image_shape][indx], axis=0),
         instance_masks=instance_masks,
         keypoints=keypoints,
+        keypoint_scores=keypoint_scores,
+        keypoint_edges=keypoint_edges,
         max_boxes_to_draw=max_boxes_to_draw,
         min_score_thresh=min_score_thresh,
         use_normalized_coordinates=use_normalized_coordinates)
@@ -596,7 +744,9 @@ def draw_side_by_side_evaluation_image(eval_dict,
         true_image_shape=tf.expand_dims(
             eval_dict[input_data_fields.true_image_shape][indx], axis=0),
         instance_masks=groundtruth_instance_masks,
-        keypoints=None,
+        keypoints=groundtruth_keypoints,
+        keypoint_scores=groundtruth_keypoint_scores,
+        keypoint_edges=keypoint_edges,
         max_boxes_to_draw=None,
         min_score_thresh=0.0,
         use_normalized_coordinates=use_normalized_coordinates)
@@ -628,6 +778,7 @@ def draw_side_by_side_evaluation_image(eval_dict,
                   eval_dict[input_data_fields.true_image_shape][indx], axis=0),
               instance_masks=groundtruth_instance_masks,
               keypoints=None,
+              keypoint_edges=None,
               max_boxes_to_draw=None,
               min_score_thresh=0.0,
               use_normalized_coordinates=use_normalized_coordinates))
@@ -641,51 +792,113 @@ def draw_side_by_side_evaluation_image(eval_dict,
 
 def draw_keypoints_on_image_array(image,
                                   keypoints,
+                                  keypoint_scores=None,
+                                  min_score_thresh=0.5,
                                   color='red',
                                   radius=2,
-                                  use_normalized_coordinates=True):
+                                  use_normalized_coordinates=True,
+                                  keypoint_edges=None,
+                                  keypoint_edge_color='green',
+                                  keypoint_edge_width=2):
   """Draws keypoints on an image (numpy array).
 
   Args:
     image: a numpy array with shape [height, width, 3].
     keypoints: a numpy array with shape [num_keypoints, 2].
+    keypoint_scores: a numpy array with shape [num_keypoints]. If provided, only
+      those keypoints with a score above score_threshold will be visualized.
+    min_score_thresh: A scalar indicating the minimum keypoint score required
+      for a keypoint to be visualized. Note that keypoint_scores must be
+      provided for this threshold to take effect.
     color: color to draw the keypoints with. Default is red.
     radius: keypoint radius. Default value is 2.
     use_normalized_coordinates: if True (default), treat keypoint values as
       relative to the image.  Otherwise treat them as absolute.
+    keypoint_edges: A list of tuples with keypoint indices that specify which
+      keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
+      edges from keypoint 0 to 1 and from keypoint 2 to 4.
+    keypoint_edge_color: color to draw the keypoint edges with. Default is red.
+    keypoint_edge_width: width of the edges drawn between keypoints. Default
+      value is 2.
   """
   image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
-  draw_keypoints_on_image(image_pil, keypoints, color, radius,
-                          use_normalized_coordinates)
+  draw_keypoints_on_image(image_pil,
+                          keypoints,
+                          keypoint_scores=keypoint_scores,
+                          min_score_thresh=min_score_thresh,
+                          color=color,
+                          radius=radius,
+                          use_normalized_coordinates=use_normalized_coordinates,
+                          keypoint_edges=keypoint_edges,
+                          keypoint_edge_color=keypoint_edge_color,
+                          keypoint_edge_width=keypoint_edge_width)
   np.copyto(image, np.array(image_pil))
 
 
 def draw_keypoints_on_image(image,
                             keypoints,
+                            keypoint_scores=None,
+                            min_score_thresh=0.5,
                             color='red',
                             radius=2,
-                            use_normalized_coordinates=True):
+                            use_normalized_coordinates=True,
+                            keypoint_edges=None,
+                            keypoint_edge_color='green',
+                            keypoint_edge_width=2):
   """Draws keypoints on an image.
 
   Args:
     image: a PIL.Image object.
     keypoints: a numpy array with shape [num_keypoints, 2].
+    keypoint_scores: a numpy array with shape [num_keypoints].
+    min_score_thresh: a score threshold for visualizing keypoints. Only used if
+      keypoint_scores is provided.
     color: color to draw the keypoints with. Default is red.
     radius: keypoint radius. Default value is 2.
     use_normalized_coordinates: if True (default), treat keypoint values as
       relative to the image.  Otherwise treat them as absolute.
+    keypoint_edges: A list of tuples with keypoint indices that specify which
+      keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
+      edges from keypoint 0 to 1 and from keypoint 2 to 4.
+    keypoint_edge_color: color to draw the keypoint edges with. Default is red.
+    keypoint_edge_width: width of the edges drawn between keypoints. Default
+      value is 2.
   """
   draw = ImageDraw.Draw(image)
   im_width, im_height = image.size
+  keypoints = np.array(keypoints)
   keypoints_x = [k[1] for k in keypoints]
   keypoints_y = [k[0] for k in keypoints]
   if use_normalized_coordinates:
     keypoints_x = tuple([im_width * x for x in keypoints_x])
     keypoints_y = tuple([im_height * y for y in keypoints_y])
-  for keypoint_x, keypoint_y in zip(keypoints_x, keypoints_y):
-    draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
-                  (keypoint_x + radius, keypoint_y + radius)],
-                 outline=color, fill=color)
+  if keypoint_scores is not None:
+    keypoint_scores = np.array(keypoint_scores)
+    valid_kpt = np.greater(keypoint_scores, min_score_thresh)
+  else:
+    valid_kpt = np.where(np.any(np.isnan(keypoints), axis=1),
+                         np.zeros_like(keypoints[:, 0]),
+                         np.ones_like(keypoints[:, 0]))
+  valid_kpt = [v for v in valid_kpt]
+
+  for keypoint_x, keypoint_y, valid in zip(keypoints_x, keypoints_y, valid_kpt):
+    if valid:
+      draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
+                    (keypoint_x + radius, keypoint_y + radius)],
+                   outline=color, fill=color)
+  if keypoint_edges is not None:
+    for keypoint_start, keypoint_end in keypoint_edges:
+      if (keypoint_start < 0 or keypoint_start >= len(keypoints) or
+          keypoint_end < 0 or keypoint_end >= len(keypoints)):
+        continue
+      if not (valid_kpt[keypoint_start] and valid_kpt[keypoint_end]):
+        continue
+      edge_coordinates = [
+          keypoints_x[keypoint_start], keypoints_y[keypoint_start],
+          keypoints_x[keypoint_end], keypoints_y[keypoint_end]
+      ]
+      draw.line(
+          edge_coordinates, fill=keypoint_edge_color, width=keypoint_edge_width)
 
 
 def draw_mask_on_image_array(image, mask, color='red', alpha=0.4):
@@ -730,6 +943,8 @@ def visualize_boxes_and_labels_on_image_array(
     instance_masks=None,
     instance_boundaries=None,
     keypoints=None,
+    keypoint_scores=None,
+    keypoint_edges=None,
     track_ids=None,
     use_normalized_coordinates=False,
     max_boxes_to_draw=20,
@@ -737,6 +952,7 @@ def visualize_boxes_and_labels_on_image_array(
     agnostic_mode=False,
     line_thickness=4,
     groundtruth_box_visualization_color='black',
+    skip_boxes=False,
     skip_scores=False,
     skip_labels=False,
     skip_track_ids=False):
@@ -762,7 +978,11 @@ def visualize_boxes_and_labels_on_image_array(
     instance_boundaries: a numpy array of shape [N, image_height, image_width]
       with values ranging between 0 and 1, can be None.
     keypoints: a numpy array of shape [N, num_keypoints, 2], can
-      be None
+      be None.
+    keypoint_scores: a numpy array of shape [N, num_keypoints], can be None.
+    keypoint_edges: A list of tuples with keypoint indices that specify which
+      keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
+      edges from keypoint 0 to 1 and from keypoint 2 to 4.
     track_ids: a numpy array of shape [N] with unique track ids. If provided,
       color-coding of boxes will be determined by these ids, and not the class
       indices.
@@ -770,13 +990,15 @@ def visualize_boxes_and_labels_on_image_array(
       normalized coordinates or not.
     max_boxes_to_draw: maximum number of boxes to visualize.  If None, draw
       all boxes.
-    min_score_thresh: minimum score threshold for a box to be visualized
+    min_score_thresh: minimum score threshold for a box or keypoint to be
+      visualized.
     agnostic_mode: boolean (default: False) controlling whether to evaluate in
       class-agnostic mode or not.  This mode will display scores but ignore
       classes.
     line_thickness: integer (default: 4) controlling line width of the boxes.
     groundtruth_box_visualization_color: box color for visualizing groundtruth
       boxes
+    skip_boxes: whether to skip the drawing of bounding boxes.
     skip_scores: whether to skip score when drawing a single detection
     skip_labels: whether to skip label when drawing a single detection
     skip_track_ids: whether to skip track id when drawing a single detection
@@ -791,10 +1013,13 @@ def visualize_boxes_and_labels_on_image_array(
   box_to_instance_masks_map = {}
   box_to_instance_boundaries_map = {}
   box_to_keypoints_map = collections.defaultdict(list)
+  box_to_keypoint_scores_map = collections.defaultdict(list)
   box_to_track_ids_map = {}
   if not max_boxes_to_draw:
     max_boxes_to_draw = boxes.shape[0]
-  for i in range(min(max_boxes_to_draw, boxes.shape[0])):
+  for i in range(boxes.shape[0]):
+    if max_boxes_to_draw == len(box_to_color_map):
+      break
     if scores is None or scores[i] > min_score_thresh:
       box = tuple(boxes[i].tolist())
       if instance_masks is not None:
@@ -803,6 +1028,8 @@ def visualize_boxes_and_labels_on_image_array(
         box_to_instance_boundaries_map[box] = instance_boundaries[i]
       if keypoints is not None:
         box_to_keypoints_map[box].extend(keypoints[i])
+      if keypoint_scores is not None:
+        box_to_keypoint_scores_map[box].extend(keypoint_scores[i])
       if track_ids is not None:
         box_to_track_ids_map[box] = track_ids[i]
       if scores is None:
@@ -860,16 +1087,24 @@ def visualize_boxes_and_labels_on_image_array(
         ymax,
         xmax,
         color=color,
-        thickness=line_thickness,
+        thickness=0 if skip_boxes else line_thickness,
         display_str_list=box_to_display_str_map[box],
         use_normalized_coordinates=use_normalized_coordinates)
     if keypoints is not None:
+      keypoint_scores_for_box = None
+      if box_to_keypoint_scores_map:
+        keypoint_scores_for_box = box_to_keypoint_scores_map[box]
       draw_keypoints_on_image_array(
           image,
           box_to_keypoints_map[box],
+          keypoint_scores_for_box,
+          min_score_thresh=min_score_thresh,
           color=color,
           radius=line_thickness / 2,
-          use_normalized_coordinates=use_normalized_coordinates)
+          use_normalized_coordinates=use_normalized_coordinates,
+          keypoint_edges=keypoint_edges,
+          keypoint_edge_color=color,
+          keypoint_edge_width=line_thickness // 2)
 
   return image
 
@@ -950,7 +1185,8 @@ class EvalMetricOpsVisualization(six.with_metaclass(abc.ABCMeta, object)):
                max_boxes_to_draw=20,
                min_score_thresh=0.2,
                use_normalized_coordinates=True,
-               summary_name_prefix='evaluation_image'):
+               summary_name_prefix='evaluation_image',
+               keypoint_edges=None):
     """Creates an EvalMetricOpsVisualization.
 
     Args:
@@ -958,10 +1194,13 @@ class EvalMetricOpsVisualization(six.with_metaclass(abc.ABCMeta, object)):
       max_examples_to_draw: The maximum number of example summaries to produce.
       max_boxes_to_draw: The maximum number of boxes to draw for detections.
       min_score_thresh: The minimum score threshold for showing detections.
-      use_normalized_coordinates: Whether to assume boxes and kepoints are in
-        normalized coordinates (as opposed to absolute coordiantes).
+      use_normalized_coordinates: Whether to assume boxes and keypoints are in
+        normalized coordinates (as opposed to absolute coordinates).
         Default is True.
       summary_name_prefix: A string prefix for each image summary.
+      keypoint_edges: A list of tuples with keypoint indices that specify which
+        keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
+        edges from keypoint 0 to 1 and from keypoint 2 to 4.
     """
 
     self._category_index = category_index
@@ -970,6 +1209,7 @@ class EvalMetricOpsVisualization(six.with_metaclass(abc.ABCMeta, object)):
     self._min_score_thresh = min_score_thresh
     self._use_normalized_coordinates = use_normalized_coordinates
     self._summary_name_prefix = summary_name_prefix
+    self._keypoint_edges = keypoint_edges
     self._images = []
 
   def clear(self):
@@ -1005,6 +1245,12 @@ class EvalMetricOpsVisualization(six.with_metaclass(abc.ABCMeta, object)):
           int64 tensor with 1-indexed groundtruth classes.
         fields.InputDataFields.groundtruth_instance_masks - (optional)
           [batch_size, num_boxes, H, W] int64 tensor with instance masks.
+        fields.InputDataFields.groundtruth_keypoints - (optional)
+          [batch_size, num_boxes, num_keypoints, 2] float32 tensor with
+          keypoint coordinates in format [y, x].
+        fields.InputDataFields.groundtruth_keypoint_visibilities - (optional)
+          [batch_size, num_boxes, num_keypoints] bool tensor with
+          keypoint visibilities.
         fields.DetectionResultFields.detection_boxes - [batch_size,
           max_num_boxes, 4] float32 tensor with detection boxes in range [0.0,
           1.0].
@@ -1017,6 +1263,9 @@ class EvalMetricOpsVisualization(six.with_metaclass(abc.ABCMeta, object)):
         fields.DetectionResultFields.detection_keypoints - (optional)
           [batch_size, max_num_boxes, num_keypoints, 2] float32 tensor with
           keypoints.
+        fields.DetectionResultFields.detection_keypoint_scores - (optional)
+          [batch_size, max_num_boxes, num_keypoints] float32 tensor with
+          keypoints scores.
 
     Returns:
       A dictionary of image summary names to tuple of (value_op, update_op). The
@@ -1083,16 +1332,20 @@ class VisualizeSingleFrameDetections(EvalMetricOpsVisualization):
                max_boxes_to_draw=20,
                min_score_thresh=0.2,
                use_normalized_coordinates=True,
-               summary_name_prefix='Detections_Left_Groundtruth_Right'):
+               summary_name_prefix='Detections_Left_Groundtruth_Right',
+               keypoint_edges=None):
     super(VisualizeSingleFrameDetections, self).__init__(
         category_index=category_index,
         max_examples_to_draw=max_examples_to_draw,
         max_boxes_to_draw=max_boxes_to_draw,
         min_score_thresh=min_score_thresh,
         use_normalized_coordinates=use_normalized_coordinates,
-        summary_name_prefix=summary_name_prefix)
+        summary_name_prefix=summary_name_prefix,
+        keypoint_edges=keypoint_edges)
 
   def images_from_evaluation_dict(self, eval_dict):
-    return draw_side_by_side_evaluation_image(
-        eval_dict, self._category_index, self._max_boxes_to_draw,
-        self._min_score_thresh, self._use_normalized_coordinates)
+    return draw_side_by_side_evaluation_image(eval_dict, self._category_index,
+                                              self._max_boxes_to_draw,
+                                              self._min_score_thresh,
+                                              self._use_normalized_coordinates,
+                                              self._keypoint_edges)

research/object_detection/utils/visualization_utils_test.py

@@ -45,6 +45,7 @@ class VisualizationUtilsTest(tf.test.TestCase):
 
     # Show that with 34 colors, the closest prime number to 34/10 that
     # satisfies the constraints is 5.
+    default_standard_colors = visualization_utils.STANDARD_COLORS
     visualization_utils.STANDARD_COLORS = [
         'color_{}'.format(str(i)) for i in range(34)
     ]
@@ -59,6 +60,8 @@ class VisualizationUtilsTest(tf.test.TestCase):
     multiplier = visualization_utils._get_multiplier_for_color_randomness()
     self.assertEqual(13, multiplier)
 
+    visualization_utils.STANDARD_COLORS = default_standard_colors
+
   def create_colorful_test_image(self):
     """This function creates an image that can be used to test vis functions.
 
@@ -161,6 +164,8 @@ class VisualizationUtilsTest(tf.test.TestCase):
                            [[0.25, 0.25, 0.75, 0.75], [0.1, 0.3, 0.6, 1.0]]])
       classes = tf.constant([[1, 1], [1, 2]], dtype=tf.int64)
       scores = tf.constant([[0.8, 0.1], [0.6, 0.5]])
+      keypoints = tf.random.uniform((2, 2, 4, 2), maxval=1.0, dtype=tf.float32)
+      keypoint_edges = [(0, 1), (1, 2), (2, 3), (3, 0)]
       images_with_boxes = (
           visualization_utils.draw_bounding_boxes_on_image_tensors(
               images_tensor,
@@ -170,7 +175,9 @@ class VisualizationUtilsTest(tf.test.TestCase):
               category_index,
               original_image_spatial_shape=image_shape,
               true_image_shape=image_shape,
-              min_score_thresh=0.2))
+              keypoints=keypoints,
+              min_score_thresh=0.2,
+              keypoint_edges=keypoint_edges))
 
       with self.test_session() as sess:
         sess.run(tf.global_variables_initializer())
@@ -297,8 +304,34 @@ class VisualizationUtilsTest(tf.test.TestCase):
     test_image = Image.fromarray(test_image)
     width_original, height_original = test_image.size
     keypoints = [[0.25, 0.75], [0.4, 0.6], [0.1, 0.1], [0.9, 0.9]]
+    keypoint_scores = [0.8, 0.2, 0.2, 0.7]
+    keypoint_edges = [(0, 1), (1, 2), (2, 3), (3, 0)]
+
+    visualization_utils.draw_keypoints_on_image(
+        test_image,
+        keypoints,
+        keypoint_scores,
+        keypoint_edges=keypoint_edges,
+        keypoint_edge_width=1,
+        keypoint_edge_color='green')
+    width_final, height_final = test_image.size
+
+    self.assertEqual(width_original, width_final)
+    self.assertEqual(height_original, height_final)
 
-    visualization_utils.draw_keypoints_on_image(test_image, keypoints)
+  def test_draw_keypoints_on_image_with_default_keypoint_scores(self):
+    test_image = self.create_colorful_test_image()
+    test_image = Image.fromarray(test_image)
+    width_original, height_original = test_image.size
+    keypoints = [[0.25, np.nan], [0.4, 0.6], [np.nan, np.nan], [0.9, 0.9]]
+    keypoint_edges = [(0, 1), (1, 2), (2, 3), (3, 0)]
+
+    visualization_utils.draw_keypoints_on_image(
+        test_image,
+        keypoints,
+        keypoint_edges=keypoint_edges,
+        keypoint_edge_width=1,
+        keypoint_edge_color='green')
     width_final, height_final = test_image.size
 
     self.assertEqual(width_original, width_final)
@@ -309,8 +342,14 @@ class VisualizationUtilsTest(tf.test.TestCase):
     width_original = test_image.shape[0]
     height_original = test_image.shape[1]
     keypoints = [[0.25, 0.75], [0.4, 0.6], [0.1, 0.1], [0.9, 0.9]]
-
-    visualization_utils.draw_keypoints_on_image_array(test_image, keypoints)
+    keypoint_edges = [(0, 1), (1, 2), (2, 3), (3, 0)]
+
+    visualization_utils.draw_keypoints_on_image_array(
+        test_image,
+        keypoints,
+        keypoint_edges=keypoint_edges,
+        keypoint_edge_width=1,
+        keypoint_edge_color='green')
     width_final = test_image.shape[0]
     height_final = test_image.shape[1]
 
@@ -328,6 +367,67 @@ class VisualizationUtilsTest(tf.test.TestCase):
                                                  color='Blue', alpha=.5)
     self.assertAllEqual(test_image, expected_result)
 
+  def test_draw_heatmaps_on_image(self):
+    test_image = self.create_colorful_test_image()
+    test_image = Image.fromarray(test_image)
+    width_original, height_original = test_image.size
+    heatmaps = np.ones(shape=[10, 20, 1], dtype=float)
+    visualization_utils.draw_heatmaps_on_image(test_image, heatmaps)
+    width_final, height_final = test_image.size
+    pixels = list(test_image.getdata())
+
+    self.assertEqual(width_original, width_final)
+    self.assertEqual(height_original, height_final)
+    # The pixel shoud be painted as AliceBlue with RGB (240, 248, 255).
+    self.assertAllEqual((240, 248, 255), pixels[10])
+
+  def test_draw_heatmaps_on_image_array(self):
+    test_image = np.asarray([[[0, 0, 0], [0, 0, 0]],
+                             [[0, 0, 0], [0, 0, 0]]], dtype=np.uint8)
+    heatmap1 = np.asarray([[1, 0],
+                           [0, 1]], dtype=np.float)
+    heatmap2 = np.asarray([[0, 1],
+                           [1, 0]], dtype=np.float)
+    heatmaps = np.stack([heatmap1, heatmap2], axis=0)
+    output_image = visualization_utils.draw_heatmaps_on_image_array(
+        test_image, heatmaps)
+    # Output image should be painted as "AliceBlue" at (0, 0), (1, 1)
+    # and "Chartreuse" at (0, 1), (1, 0).
+    self.assertAllEqual(
+        output_image,
+        np.array([[[240, 248, 255], [127, 255, 0]],
+                  [[127, 255, 0], [240, 248, 255]]]))
+
+  def test_draw_heatmaps_on_image_tensors(self):
+    test_image = np.asarray([[[0, 0, 0], [0, 0, 0]],
+                             [[0, 0, 0], [0, 0, 0]]], dtype=np.uint8)
+
+    heatmap1 = np.asarray([[1, 0],
+                           [0, 1]], dtype=np.float)
+    heatmap2 = np.asarray([[0, 1],
+                           [1, 0]], dtype=np.float)
+    heatmaps = np.stack([heatmap1, heatmap2], axis=0)
+    with tf.Graph().as_default():
+      image_tensor = tf.constant(test_image, dtype=tf.uint8)
+      image_tensor = tf.expand_dims(image_tensor, axis=0)
+      heatmaps_tensor = tf.expand_dims(
+          tf.constant(heatmaps, dtype=tf.float32), axis=0)
+      output_image = visualization_utils.draw_heatmaps_on_image_tensors(
+          images=image_tensor,
+          heatmaps=heatmaps_tensor,
+          apply_sigmoid=False)
+
+      with self.test_session() as sess:
+        sess.run(tf.global_variables_initializer())
+
+        output_image_np = sess.run(output_image)
+      self.assertAllEqual(
+          output_image_np,
+          np.expand_dims(
+              np.array([[[240, 248, 255], [127, 255, 0]],
+                        [[127, 255, 0], [240, 248, 255]]]),
+              axis=0))
+
   def test_add_cdf_image_summary(self):
     values = [0.1, 0.2, 0.3, 0.4, 0.42, 0.44, 0.46, 0.48, 0.50]
     visualization_utils.add_cdf_image_summary(values, 'PositiveAnchorLoss')
@@ -371,6 +471,7 @@ class VisualizationUtilsTest(tf.test.TestCase):
                                           minval=0.0,
                                           maxval=1.0,
                                           dtype=tf.float32)
+    num_groundtruth_boxes = tf.constant([3, 8, 0, 2], tf.int32)
     groundtruth_classes = tf.random_uniform([4, 8],
                                             minval=1,
                                             maxval=3,
@@ -390,7 +491,9 @@ class VisualizationUtilsTest(tf.test.TestCase):
         fields.InputDataFields.groundtruth_boxes:
             groundtruth_boxes,
         fields.InputDataFields.groundtruth_classes:
-            groundtruth_classes
+            groundtruth_classes,
+        fields.InputDataFields.num_groundtruth_boxes:
+            num_groundtruth_boxes
     }
     metric_ops = eval_metric_ops.get_estimator_eval_metric_ops(eval_dict)
     _, update_op = metric_ops[next(six.iterkeys(metric_ops))]
@@ -446,6 +549,28 @@ class VisualizationUtilsTest(tf.test.TestCase):
           six.b(''),
           value_ops_out[metric_op_base + '/' + str(max_examples_to_draw - 1)])
 
+  def test_visualize_boxes_and_labels_on_image_array(self):
+    ori_image = np.ones([360, 480, 3], dtype=np.int32) * 255
+    test_image = np.ones([360, 480, 3], dtype=np.int32) * 255
+    detections = np.array([[0.8, 0.1, 0.9, 0.1, 1., 0.1],
+                           [0.1, 0.3, 0.8, 0.7, 1., 0.6]])
+    keypoints = np.array(np.random.rand(2, 5, 2), dtype=np.float32)
+    labelmap = {1: {'id': 1, 'name': 'cat'}, 2: {'id': 2, 'name': 'dog'}}
+    visualization_utils.visualize_boxes_and_labels_on_image_array(
+        test_image,
+        detections[:, :4],
+        detections[:, 4].astype(np.int32),
+        detections[:, 5],
+        labelmap,
+        keypoints=keypoints,
+        track_ids=None,
+        use_normalized_coordinates=True,
+        max_boxes_to_draw=1,
+        min_score_thresh=0.2,
+        agnostic_mode=False,
+        line_thickness=8)
+    self.assertGreater(np.abs(np.sum(test_image - ori_image)), 0)
+
 
 if __name__ == '__main__':
   tf.test.main()