No public description

PiperOrigin-RevId: 595901724

Author: tensorflower-gardener

official/vision/utils/object_detection/visualization_utils.py

@@ -969,6 +969,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(
@@ -981,10 +986,6 @@ def _denormalize_images(images: tf.Tensor) -> tf.Tensor:
       ),
   )
 
-  if images.shape[3] > 3:
-    images = images[:, :, :, 0:3]
-  elif images.shape[3] == 1:
-    images = tf.image.grayscale_to_rgb(images)
   if true_image_shape is None:
     true_shapes = tf.constant(-1, shape=[images.shape.as_list()[0], 3])
   else: