No public description

PiperOrigin-RevId: 676475897
tensorflow · Sep 19, 2024 · 951c619 · 951c619
1 parent 8a8cb09
commit 951c619
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/official/vision/utils/object_detection/visualization_utils.py b/official/vision/utils/object_detection/visualization_utils.py
@@ -438,6 +438,11 @@ def _denormalize_images(images: tf.Tensor) -> tf.Tensor:
       )
     return tf.cast(images, dtype=tf.uint8)
 
+  if images.shape[3] > 3:
+    images = images[:, :, :, 0:3]
+  elif images.shape[3] == 1:
+    images = tf.image.grayscale_to_rgb(images)
+
   images = tf.nest.map_structure(
       tf.identity,
       tf.map_fn(