Move visualization method for Retinanet to utils/visualization_utils.py.

PiperOrigin-RevId: 296077315

official/vision/detection/executor/detection_executor.py

@@ -21,11 +21,9 @@ from __future__ import print_function
 
 from absl import logging
 
-import os
-import json
 import tensorflow.compat.v2 as tf
 from official.modeling.training import distributed_executor as executor
-from official.vision.detection.utils import box_utils
+from official.vision.detection.utils.object_detection import visualization_utils
 
 
 class DetectionDistributedExecutor(executor.DistributedExecutor):
@@ -115,7 +113,7 @@ class DetectionDistributedExecutor(executor.DistributedExecutor):
           # TODO(hongjunchoi): Once dynamic slicing is supported on TPU, only
           # write correct slice of outputs to summary file.
           if num_remaining_visualizations > 0:
-            box_utils.visualize_bounding_boxes(
+            visualization_utils.visualize_images_with_bounding_boxes(
                 inputs, prediction_outputs['detection_boxes'],
                 self.global_train_step, self.eval_summary_writer)
 

official/vision/detection/utils/object_detection/visualization_utils.py

@@ -32,6 +32,7 @@ import PIL.ImageFont as ImageFont
 import six
 import tensorflow.compat.v2 as tf
 
+from official.vision.detection.utils import box_utils
 from official.vision.detection.utils.object_detection import shape_utils
 
 
@@ -93,6 +94,23 @@ def encode_image_array_as_png_str(image):
   return png_string
 
 
+def visualize_images_with_bounding_boxes(images, box_outputs, step,
+                                         summary_writer):
+  """Records subset of evaluation images with bounding boxes."""
+  image_shape = tf.shape(images[0])
+  image_height = tf.cast(image_shape[0], tf.float32)
+  image_width = tf.cast(image_shape[1], tf.float32)
+  normalized_boxes = box_utils.normalize_boxes(box_outputs,
+                                               [image_height, image_width])
+
+  bounding_box_color = tf.constant([[1.0, 1.0, 0.0, 1.0]])
+  image_summary = tf.image.draw_bounding_boxes(images, normalized_boxes,
+                                               bounding_box_color)
+  with summary_writer.as_default():
+    tf.summary.image('bounding_box_summary', image_summary, step=step)
+    summary_writer.flush()
+
+
 def draw_bounding_box_on_image_array(image,
                                      ymin,
                                      xmin,