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 Zhichao Lu:

    Makes the freezable_batch_norm_test run w/ v2 behavior.

    The main change is in v2 updates will happen right away when running batchnorm in training mode. So, we need to restore the weights between batchnorm calls to make sure the numerical checks all start from the same place.

--
309425881  by Zhichao Lu:

    Make TF1/TF2 optimizer builder tests explicit.

--
309408646  by Zhichao Lu:

    Make dataset builder tests TF1 and TF2 compatible.

--
309246305  by Zhichao Lu:

    Added the functionality of combining the person keypoints and object detection
    annotations in the binary that converts the COCO raw data to TfRecord.

--
309125076  by Zhichao Lu:

    Convert target_assigner_utils to TF1/TF2.

--
308966359  by huizhongc:

    Support SSD training with partially labeled groundtruth.

--
308937159  by rathodv:

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

--
308774302  by Zhichao Lu:

    Internal

--
308732860  by rathodv:

    Make core/prefetcher.py  compatible with TF1 only.

--
308726984  by rathodv:

    Update core/multiclass_nms_test.py to be TF1/TF2 compatible.

--
308714718  by rathodv:

    Update core/region_similarity_calculator_test.py to be TF1/TF2 compatible.

--
308707960  by rathodv:

    Update core/minibatch_sampler_test.py to be TF1/TF2 compatible.

--
308700595  by rathodv:

    Update core/losses_test.py to be TF1/TF2 compatible and remove losses_test_v2.py

--
308361472  by rathodv:

    Update core/matcher_test.py to be TF1/TF2 compatible.

--
308335846  by Zhichao Lu:

    Updated the COCO evaluation logics and populated the groundturth area
    information through. This change matches the groundtruth format expected by the
    COCO keypoint evaluation.

--
308256924  by rathodv:

    Update core/keypoints_ops_test.py to be TF1/TF2 compatible.

--
308256826  by rathodv:

    Update class_agnostic_nms_test.py to be TF1/TF2 compatible.

--
308256112  by rathodv:

    Update box_list_ops_test.py to be TF1/TF2 compatible.

--
308159360  by Zhichao Lu:

    Internal change

308145008  by Zhichao Lu:

    Added 'image/class/confidence' field in the TFExample decoder.

--
307651875  by rathodv:

    Refactor core/box_list.py to support TF1/TF2.

--
307651798  by rathodv:

    Modify box_coder.py base class to work with with TF1/TF2

--
307651652  by rathodv:

    Refactor core/balanced_positive_negative_sampler.py to support TF1/TF2.

--
307651571  by rathodv:

    Modify BoxCoders tests to use test_case:execute method to allow testing with TF1.X and TF2.X

--
307651480  by rathodv:

    Modify Matcher tests to use test_case:execute method to allow testing with TF1.X and TF2.X

--
307651409  by rathodv:

    Modify AnchorGenerator tests to use test_case:execute method to allow testing with TF1.X and TF2.X

--
307651314  by rathodv:

    Refactor model_builder to support TF1 or TF2 models based on TensorFlow version.

--
307092053  by Zhichao Lu:

    Use manager to save checkpoint.

--
307071352  by ronnyvotel:

    Fixing keypoint visibilities. Now by default, the visibility is marked True if the keypoint is labeled (regardless of whether it is visible or not).
    Also, if visibilities are not present in the dataset, they will be created based on whether the keypoint coordinates are finite (vis = True) or NaN (vis = False).

--
307069557  by Zhichao Lu:

    Internal change to add few fields related to postprocessing parameters in
    center_net.proto and populate those parameters to the keypoint postprocessing
    functions.

--
307012091  by Zhichao Lu:

    Make Adam Optimizer's epsilon proto configurable.

    Potential issue: tf.compat.v1's AdamOptimizer has a default epsilon on 1e-08 ([doc-link](https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/AdamOptimizer))  whereas tf.keras's AdamOptimizer has default epsilon 1e-07 ([doc-link](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam))

--
306858598  by Zhichao Lu:

    Internal changes to update the CenterNet model:
    1) Modified eval job loss computation to avoid averaging over batches with zero loss.
    2) Updated CenterNet keypoint heatmap target assigner to apply box size to heatmap Guassian standard deviation.
    3) Updated the CenterNet meta arch keypoint losses computation to apply weights outside of loss function.

--
306731223  by jonathanhuang:

    Internal change.

--
306549183  by rathodv:

    Internal Update.

--
306542930  by rathodv:

    Internal Update

--
306322697  by rathodv:

    Internal.

--
305345036  by Zhichao Lu:

    Adding COCO Camera Traps Json to tf.Example beam code

--
304104869  by lzc:

    Internal changes.

--
304068971  by jonathanhuang:

    Internal change.

--
304050469  by Zhichao Lu:

    Internal change.

--
303880642  by huizhongc:

    Support parsing partially labeled groundtruth.

--
303841743  by Zhichao Lu:

    Deprecate nms_on_host in SSDMetaArch.

--
303803204  by rathodv:

    Internal change.

--
303793895  by jonathanhuang:

    Internal change.

--
303467631  by rathodv:

    Py3 update for detection inference test.

--
303444542  by rathodv:

    Py3 update to metrics module

--
303421960  by rathodv:

    Update json_utils to python3.

--
302787583  by ronnyvotel:

    Coco results generator for submission to the coco test server.

--
302719091  by Zhichao Lu:

    Internal change to add the ResNet50 image feature extractor for CenterNet model.

--
302116230  by Zhichao Lu:

    Added the functions to overlay the heatmaps with images in visualization util
    library.

--
301888316  by Zhichao Lu:

    Fix checkpoint_filepath not defined error.

--
301840312  by ronnyvotel:

    Adding keypoint_scores to visualizations.

--
301683475  by ronnyvotel:

    Introducing the ability to preprocess `keypoint_visibilities`.

    Some data augmentation ops such as random crop can filter instances and keypoints. It's important to also filter keypoint visibilities, so that the groundtruth tensors are always in alignment.

--
301532344  by Zhichao Lu:

    Don't use tf.divide since "Quantization not yet supported for op: DIV"

--
301480348  by ronnyvotel:

    Introducing keypoint evaluation into model lib v2.
    Also, making some fixes to coco keypoint evaluation.

--
301454018  by Zhichao Lu:

    Added the image summary to visualize the train/eval input images and eval's
    prediction/groundtruth side-by-side image.

--
301317527  by Zhichao Lu:

    Updated the random_absolute_pad_image function in the preprocessor library to
    support the keypoints argument.

--
301300324  by Zhichao Lu:

    Apply name change(experimental_run_v2 -> run) for all callers in Tensorflow.

--
301297115  by ronnyvotel:

    Utility function for setting keypoint visibilities based on keypoint coordinates.

--
301248885  by Zhichao Lu:

    Allow MultiworkerMirroredStrategy(MWMS) use by adding checkpoint handling with temporary directories in model_lib_v2. Added missing WeakKeyDictionary cfer_fn_cache field in CollectiveAllReduceStrategyExtended.

--
301224559  by Zhichao Lu:

    ...1) Fixes model_lib to also use keypoints while preparing model groundtruth.
    ...2) Tests model_lib with newly added keypoint metrics config.

--
300836556  by Zhichao Lu:

    Internal changes to add keypoint estimation parameters in CenterNet proto.

--
300795208  by Zhichao Lu:

    Updated the eval_util library to populate the keypoint groundtruth to
    eval_dict.

--
299474766  by Zhichao Lu:

    ...Modifies eval_util to create Keypoint Evaluator objects when configured in eval config.

--
299453920  by Zhichao Lu:

    Add swish activation as a hyperperams option.

--
299240093  by ronnyvotel:

    Keypoint postprocessing for CenterNetMetaArch.

--
299176395  by Zhichao Lu:

    Internal change.

--
299135608  by Zhichao Lu:

    Internal changes to refactor the CenterNet model in preparation for keypoint estimation tasks.

--
298915482  by Zhichao Lu:

    Make dataset_builder aware of input_context for distributed training.

--
298713595  by Zhichao Lu:

    Handling data with negative size boxes.

--
298695964  by Zhichao Lu:

    Expose change_coordinate_frame as a config parameter; fix multiclass_scores optional field.

--
298492150  by Zhichao Lu:

    Rename optimizer_builder_test_v2.py -> optimizer_builder_v2_test.py

--
298476471  by Zhichao Lu:

    Internal changes to support CenterNet keypoint estimation.

--
298365851  by ronnyvotel:

    Fixing a bug where groundtruth_keypoint_weights were being padded with a dynamic dimension.

--
297843700  by Zhichao Lu:

    Internal change.

--
297706988  by lzc:

    Internal change.

--
297705287  by ronnyvotel:

    Creating the "snapping" behavior in CenterNet, where regressed keypoints are refined with updated candidate keypoints from a heatmap.

--
297700447  by Zhichao Lu:

    Improve checkpoint checking logic with TF2 loop.

--
297686094  by Zhichao Lu:

    Convert "import tensorflow as tf" to "import tensorflow.compat.v1".

--
297670468  by lzc:

    Internal change.

--
297241327  by Zhichao Lu:

    Convert "import tensorflow as tf" to "import tensorflow.compat.v1".

--
297205959  by Zhichao Lu:

    Internal changes to support refactored the centernet object detection target assigner into a separate library.

--
297143806  by Zhichao Lu:

    Convert "import tensorflow as tf" to "import tensorflow.compat.v1".

--
297129625  by Zhichao Lu:

    Explicitly replace "import tensorflow" with "tensorflow.compat.v1" for TF2.x migration

--
297117070  by Zhichao Lu:

    Explicitly replace "import tensorflow" with "tensorflow.compat.v1" for TF2.x migration

--
297030190  by Zhichao Lu:

    Add configuration options for visualizing keypoint edges

--
296359649  by Zhichao Lu:

    Support DepthwiseConv2dNative (of separable conv) in weight equalization loss.

--
296290582  by Zhichao Lu:

    Internal change.

--
296093857  by Zhichao Lu:

    Internal changes to add general target assigner utilities.

--
295975116  by Zhichao Lu:

    Fix visualize_boxes_and_labels_on_image_array to show max_boxes_to_draw correctly.

--
295819711  by Zhichao Lu:

    Adds a flag to visualize_boxes_and_labels_on_image_array to skip the drawing of axis aligned bounding boxes.

--
295811929  by Zhichao Lu:

    Keypoint support in random_square_crop_by_scale.

--
295788458  by rathodv:

    Remove unused checkpoint to reduce repo size on github

--
295787184  by Zhichao Lu:

    Enable visualization of edges between keypoints

--
295763508  by Zhichao Lu:

    [Context RCNN] Add an option to enable / disable cropping feature in the post
    process step in the meta archtecture.

--
295605344  by Zhichao Lu:

    internal change.

--
294926050  by ronnyvotel:

    Adding per-keypoint groundtruth weights. These weights are intended to be used as multipliers in a keypoint loss function.

    Groundtruth keypoint weights are constructed as follows:
    - Initialize the weight for each keypoint type based on user-specified weights in the input_reader proto
    - Mask out (i.e. make zero) all keypoint weights that are not visible.

--
294829061  by lzc:

    Internal change.

--
294566503  by Zhichao Lu:

    Changed internal CenterNet Model configuration.

--
294346662  by ronnyvotel:

    Using NaN values in keypoint coordinates that are not visible.

--
294333339  by Zhichao Lu:

    Change experimetna_distribute_dataset -> experimental_distribute_dataset_from_function

--
293928752  by Zhichao Lu:

    Internal change

--
293909384  by Zhichao Lu:

    Add capabilities to train 1024x1024 CenterNet models.

--
293637554  by ronnyvotel:

    Adding keypoint visibilities to TfExampleDecoder.

--
293501558  by lzc:

    Internal change.

--
293252851  by Zhichao Lu:

    Change tf.gfile.GFile to tf.io.gfile.GFile.

--
292730217  by Zhichao Lu:

    Internal change.

--
292456563  by lzc:

    Internal changes.

--
292355612  by Zhichao Lu:

    Use tf.gather and tf.scatter_nd instead of matrix ops.

--
292245265  by rathodv:

    Internal

--
291989323  by richardmunoz:

    Refactor out building a DataDecoder from building a tf.data.Dataset.

--
291950147  by Zhichao Lu:

    Flip bounding boxes in arbitrary shaped tensors.

--
291401052  by huizhongc:

    Fix multiscale grid anchor generator to allow fully convolutional inference. When exporting model with identity_resizer as image_resizer, there is an incorrect box offset on the detection results. We add the anchor offset to address this problem.

--
291298871  by Zhichao Lu:

    Py3 compatibility changes.

--
290957957  by Zhichao Lu:

    Hourglass feature extractor for CenterNet.

--
290564372  by Zhichao Lu:

    Internal change.

--
290155278  by rathodv:

    Remove Dataset Explorer.

--
290155153  by Zhichao Lu:

    Internal change

--
290122054  by Zhichao Lu:

    Unify the format in the faster_rcnn.proto

--
290116084  by Zhichao Lu:

    Deprecate tensorflow.contrib.

--
290100672  by Zhichao Lu:

    Update MobilenetV3 SSD candidates

--
289926392  by Zhichao Lu:

    Internal change

--
289553440  by Zhichao Lu:

    [Object Detection API] Fix the comments about the dimension of the rpn_box_encodings from 4-D to 3-D.

--
288994128  by lzc:

    Internal changes.

--
288942194  by lzc:

    Internal change.

--
288746124  by Zhichao Lu:

    Configurable channel mean/std. dev in CenterNet feature extractors.

--
288552509  by rathodv:

    Internal.

--
288541285  by rathodv:

    Internal update.

--
288396396  by Zhichao Lu:

    Make object detection import contrib explicitly

--
288255791  by rathodv:

    Internal

--
288078600  by Zhichao Lu:

    Fix model_lib_v2 test

--
287952244  by rathodv:

    Internal

--
287921774  by Zhichao Lu:

    internal change

--
287906173  by Zhichao Lu:

    internal change

--
287889407  by jonathanhuang:

    PY3 compatibility

--
287889042  by rathodv:

    Internal

--
287876178  by Zhichao Lu:

    Internal change.

--
287770490  by Zhichao Lu:

    Add CenterNet proto and builder

--
287694213  by Zhichao Lu:

    Support for running multiple steps per tf.function call.

--
287377183  by jonathanhuang:

    PY3 compatibility

--
287371344  by rathodv:

    Support loading keypoint labels and ids.

--
287368213  by rathodv:

    Add protos supporting keypoint evaluation.

--
286673200  by rathodv:

    dataset_tools PY3 migration

--
286635106  by Zhichao Lu:

    Update code for upcoming tf.contrib removal

--
286479439  by Zhichao Lu:

    Internal change

--
286311711  by Zhichao Lu:

    Skeleton of context model within TFODAPI

--
286005546  by Zhichao Lu:

    Fix Faster-RCNN training when using keep_aspect_ratio_resizer with pad_to_max_dimension

--
285906400  by derekjchow:

    Internal change

--
285822795  by Zhichao Lu:

    Add CenterNet meta arch target assigners.

--
285447238  by Zhichao Lu:

    Internal changes.

--
285016927  by Zhichao Lu:

    Make _dummy_computation a tf.function. This fixes breakage caused by
    cl/284256438

--
284827274  by Zhichao Lu:

    Convert to python 3.

--
284645593  by rathodv:

    Internal change

--
284639893  by rathodv:

    Add missing documentation for keypoints in eval_util.py.

--
284323712  by Zhichao Lu:

    Internal changes.

--
284295290  by Zhichao Lu:

    Updating input config proto and dataset builder to include context fields

    Updating standard_fields and tf_example_decoder to include context features

--
284226821  by derekjchow:

    Update exporter.

--
284211030  by Zhichao Lu:

    API changes in CenterNet informed by the experiments with hourlgass network.

--
284190451  by Zhichao Lu:

    Add support for CenterNet losses in protos and builders.

--
284093961  by lzc:

    Internal changes.

--
284028174  by Zhichao Lu:

    Internal change

--
284014719  by derekjchow:

    Do not pad top_down feature maps unnecessarily.

--
284005765  by Zhichao Lu:

    Add new pad_to_multiple_resizer

--
283858233  by Zhichao Lu:

    Make target assigner work when under tf.function.

--
283836611  by Zhichao Lu:

    Make config getters more general.

--
283808990  by Zhichao Lu:

    Internal change

--
283754588  by Zhichao Lu:

    Internal changes.

--
282460301  by Zhichao Lu:

    Add ability to restore v2 style checkpoints.

--
281605842  by lzc:

    Add option to disable loss computation in OD API eval job.

--
280298212  by Zhichao Lu:

    Add backwards compatible change

--
280237857  by Zhichao Lu:

    internal change

--

PiperOrigin-RevId: 310447280

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()