autoaugment_tensorflow.py

import tensorflow as tf
import tensorflow_datasets as tfds

from aug.autoaugment import distort_image_with_autoaugment


def aug_fn(sample, policy="v0"):
    """Apply an AutoAugment policy to a dataset sample. Make sure to check
    the format of your dataset to call the distort function with the appropriate
    arguments. bboxes should have format [y_min, x_min, y_max, x_max] with values 
    normalized to [0,1]."""
    aug_img, aug_bboxes = distort_image_with_autoaugment(image=sample["image"],
                                                         bboxes=sample["objects"]["bbox"],
                                                         augmentation_name=policy)
    sample["image"], sample["objects"]["bbox"] = aug_img, aug_bboxes


# download pascal voc dataset
train_ds = tfds.load(name="voc/2007", split='train')

# apply augmentation
aug_train_ds = train_ds.map(aug_fn)