How to augment data in a Tensorflow Dataset?

Question

For a set of images, I was confused if the term Data Augmentation meant to transform the current dataset (e.g. crop/flip/rotate/...) or if it meant to increase the amount of data by adding the cropped/flipped/rotated images to the initial dataset. As far as I understand, from this question and this one, it means both. Please correct me if I'm wrong.

So, using Tensorflow Dataset, I want to achieve the second one: augmenting the amount of data.

I'm using the ImageNet data from TFDS (trainning set is not available):

import tensorflow_datasets as tfds
ds = tfds.load('imagenet_a', split='test', as_supervised=True)

And I want to flip the images:

def transform(image, label):
    image = tf.image.flip_left_right(image)
    return image, label

It works well if I apply the transformation directly to the dataset. But it doesn't increase the amount of data:

ds = ds.map(transform)

So, I tried to create a second dataset and concatenate both:

ds0 = ds.map(transform)
ds = ds.concatenate(ds0)

But I get the following error:

TypeError: Two datasets to concatenate have different types (tf.uint8, tf.int64) and (tf.float32, tf.int64)

Is it the way to do to concatenate two datasets to increase a training set? Or how to do it correctly? (or how to fix my error)

I'm aware of ImageDataGenerator, but it doesn't contain the transformation I want

Answer 1

As the error clearly says, the two datasets should be having the same data types, you can achieve this using tf.cast but this is a bit of a hectic process for a large dataset.

You can also merge datasets using tf.data.Dataset.sample_from_datasets.

Below is some example code to illustrate this:

    import tensorflow_datasets as tfds
    import matplotlib.pyplot as plt
    import tensorflow as tf
    from tensorflow.keras.preprocessing.image import ImageDataGenerator
    from keras.preprocessing.image import img_to_array, array_to_img

    ds, info = tfds.load('imagenet_a', split='test', as_supervised=True, with_info=True)

Original sample images:

    vis = tfds.visualization.show_examples(ds, info)

We take 10 images for testing and flip those 10 images randomly with map() function to create a new dataset.

Note (2024-06-21): In the original answer the non-stateless random_flip_left_right was used. Whilst not deprecated, this is strongly discouraged by the TensorFlow documentation in favour of stateless_random_flip_left_right.

    #Take 10 images from original dataset and apply map of random flip
    ds1 = ds.take(10)
    ds2 = ds1.map(lambda image, label: (tf.image.stateless_random_flip_left_right(image), label))

    #Merge original and randomly flipped datasets
    # Returns 20 images, 10 original plus 10 randomly flipped
    new_ds = tf.data.Dataset.sample_from_datasets([ds1, ds2])
    print(len(list(new_ds))) 
    
    #Plot resulting merged dataset to show our augmentation
    f, axarr = plt.subplots(5, 4, figsize=(15, 15))
    
    ix = 0
    i = 0
    count = 0
    k = 0
    
    for images, labels in new_ds:
      crop_img = array_to_img(images)
      axarr[i, ix].imshow(crop_img)
      ix = ix + 1
      count = count + 1
      if count == 4:
         i = i + 1
         count = 0
         ix = 0

Merged Dataset:

You can see the merged data with original images and randomly flipped images.