def _augment_image(
image: tf.Tensor,
is_training: bool,
augment_name: Optional[str] = None,
) -> tf.Tensor:
"""Applies AA/RA to an image."""
if is_training and augment_name:
if 'autoaugment' in augment_name or 'randaugment' in augment_name:
input_image_type = image.dtype
image = tf.clip_by_value(image, 0.0, 255.0)
# Autoaugment requires a uint8 image; we cast here and then cast back
image = tf.cast(image, dtype=tf.uint8)
if 'autoaugment' in augment_name:
logging.info(f'Applying AutoAugment policy {augment_name}')
image = autoaugment.distort_image_with_autoaugment(image, 'v0')
elif 'randaugment' in augment_name:
magnitude = int(augment_name.split('_')[-1]) # pytype: disable=attribute-error
# Allow passing in num_layers as a magnitude > 100
if magnitude > 100:
num_layers = magnitude // 100
magnitude = magnitude - int(num_layers * 100)
else:
num_layers = 2
logging.info(f'Applying RA {num_layers} x {magnitude}')
image = autoaugment.distort_image_with_randaugment(
image, num_layers=num_layers, magnitude=magnitude)
image = tf.cast(image, dtype=input_image_type)
return image
def preprocess_for_train(image_bytes,
use_bfloat16,
image_size=IMAGE_SIZE,
autoaugment_name=None):
"""Preprocesses the given image for evaluation.
Args:
image_bytes: `Tensor` representing an image binary of arbitrary size.
use_bfloat16: `bool` for whether to use bfloat16.
image_size: image size.
autoaugment_name: `string` that is the name of the autoaugment policy
to apply to the image. If the value is `None` autoaugment will not be
applied.
Returns:
A preprocessed image `Tensor`.
"""
image = _decode_and_random_crop(image_bytes, image_size)
image = _flip(image)
image = tf.reshape(image, [image_size, image_size, 3])
image = tf.image.convert_image_dtype(
image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
if autoaugment_name:
tf.logging.info('Apply AutoAugment policy {}'.format(autoaugment_name))
image = tf.clip_by_value(image, 0.0, 255.0)
image = tf.cast(image, dtype=tf.uint8)
image = autoaugment.distort_image_with_autoaugment(
image, autoaugment_name)
image = tf.cast(image,
dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
return image
def preprocess_for_train(image_bytes,
use_bfloat16,
image_size=IMAGE_SIZE,
augment_name=None,
randaug_num_layers=None,
randaug_magnitude=None,
resize_method=None):
"""Preprocesses the given image for evaluation.
Args:
image_bytes: `Tensor` representing an image binary of arbitrary size.
use_bfloat16: `bool` for whether to use bfloat16.
image_size: image size.
augment_name: `string` that is the name of the augmentation method
to apply to the image. `autoaugment` if AutoAugment is to be used or
`randaugment` if RandAugment is to be used. If the value is `None` no
augmentation method will be applied applied. See autoaugment.py for more
details.
randaug_num_layers: 'int', if RandAug is used, what should the number of
layers be. See autoaugment.py for detailed description.
randaug_magnitude: 'int', if RandAug is used, what should the magnitude
be. See autoaugment.py for detailed description.
resize_method: resize method. If none, use bicubic.
Returns:
A preprocessed image `Tensor`.
"""
image = _decode_and_random_crop(image_bytes, image_size, resize_method)
image = _flip(image)
image = tf.reshape(image, [image_size, image_size, 3])
if augment_name:
try:
import autoaugment # pylint: disable=g-import-not-at-top
except ImportError as e:
logging.exception('Autoaugment is not supported in TF 2.x.')
raise e
logging.info('Apply AutoAugment policy %s', augment_name)
input_image_type = image.dtype
image = tf.clip_by_value(image, 0.0, 255.0)
image = tf.cast(image, dtype=tf.uint8)
if augment_name == 'autoaugment':
logging.info('Apply AutoAugment policy %s', augment_name)
image = autoaugment.distort_image_with_autoaugment(image, 'v0')
elif augment_name == 'randaugment':
image = autoaugment.distort_image_with_randaugment(
image, randaug_num_layers, randaug_magnitude)
else:
raise ValueError('Invalid value for augment_name: %s' %
(augment_name))
image = tf.cast(image, dtype=input_image_type)
image = tf.image.convert_image_dtype(
image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
return image
def test_autoaugment(self):
"""Smoke test to be sure no syntax errors."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
aug_image = autoaugment.distort_image_with_autoaugment(image, 'v0')
self.assertEqual((224, 224, 3), aug_image.shape)
def test_autoaugment_policy(self):
# A very simple test to verify no syntax error.
image = tf.placeholder(tf.uint8, shape=[640, 640, 3])
bboxes = tf.placeholder(tf.float32, shape=[4, 4])
autoaugment.distort_image_with_autoaugment(image, bboxes, 'v0')
def augment_image(image_filename, labels, image_size, format):
""" Load and resize image and perform augmentation """
image, labels = load_images(image_filename, labels, image_size, format)
image = autoaugment.distort_image_with_autoaugment(image, 'v0')
return image, labels
