Commit b0cf164b authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TF Object Detection Team
Browse files

Optimize data preprocessing for 3-channel images by bypassing concat in _augment_only_rgb_channels.

PiperOrigin-RevId: 471008666
parent 51078b5d
...@@ -1032,6 +1032,9 @@ def random_image_scale(image, ...@@ -1032,6 +1032,9 @@ def random_image_scale(image,
def _augment_only_rgb_channels(image, augment_function): def _augment_only_rgb_channels(image, augment_function):
"""Augments only the RGB slice of an image with additional channels.""" """Augments only the RGB slice of an image with additional channels."""
# Skipping the concat if possible reduces latency.
if image.shape[2] == 3:
return augment_function(image)
rgb_slice = image[:, :, :3] rgb_slice = image[:, :, :3]
augmented_rgb_slice = augment_function(rgb_slice) augmented_rgb_slice = augment_function(rgb_slice)
image = tf.concat([augmented_rgb_slice, image[:, :, 3:]], -1) image = tf.concat([augmented_rgb_slice, image[:, :, 3:]], -1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment