What is Data Augmentation? Use Cases & Examples

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential
import tensorflow_datasets as tfds

Data Loading

TensorFlow Datasets (tfds) makes it easy to load ready-to-use datasets. Here, we split the cats_vs_dogs data into three parts:

• 80% for training
• 10% for validation
• 10% for testing

(train_ds, val_ds, test_ds), metadata = tfds.load(
    'cats_vs_dogs',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
)

You can iterate through a few samples to display them with matplotlib. This helps confirm that the dataset loaded correctly and the labels match what you see.

Data Analysis

The dataset has two classes: cat and dog. It’s always helpful to check the metadata before proceeding with any augmentation or model-building. You can then preview a handful of images to understand their size, orientation, and variety. This insight helps inform which augmentations make sense (e.g., flipping, rotation, color adjustments).

Data Augmentation with Keras Sequential Layers

Keras offers built-in layers specifically for data augmentation. You can define these transformations as part of a Sequential model or apply them in a preprocessing pipeline.

1. Resizing and Rescaling

Often, you’ll resize images to a smaller dimension for faster processing and scale pixel values to a [0, 1] range by dividing by 255.

IMG_SIZE = 180

resize_and_rescale = Sequential([
    layers.Resizing(IMG_SIZE, IMG_SIZE),
    layers.Rescaling(1./255)
])

Applying resize_and_rescale(image) ensures every image is the same resolution and that raw pixel intensities are normalized.

2. Random Flip and Rotate

You can stack augmentation layers in another Sequential model. In the example below, we define random flips and rotations to alter the images geometrically:

data_augmentation = Sequential([
    layers.RandomFlip("horizontal_and_vertical"),
    layers.RandomRotation(0.4),
])

When you pass your images through data_augmentation, it outputs a new, randomly transformed version each time. Plotting them side by side confirms how flipping and rotation can diversify your training data.

3. Incorporating Augmentation Layers in the Model

There are two main approaches to include these preprocessing steps:

1. i) Directly in the Model Definition

model = Sequential([
    resize_and_rescale,
    data_augmentation,
    layers.Conv2D(16, 3, padding='same', activation='relu'),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(64, activation='relu'),
    layers.Dense(1, activation='sigmoid')
])

The augmentation layers only apply during training, not during inference (evaluate or predict).

1. ii) Using a .map Function

aug_ds = train_ds.map(lambda x, y: (data_augmentation(x, training=True), y))

This approach creates a new dataset that includes transformed images for training.

Preprocessing and Model Training

Before feeding data into the model, it’s common to prepare and batch it. Below is a helper function that:

1. Resizes and rescales every sample.
2. Optionally shuffles the data.
3. Batches the samples (e.g., 32 images per batch).
4. (Optionally) applies augmentation.
5. Prefetches to overlap data processing and model execution.

batch_size = 32
AUTOTUNE = tf.data.AUTOTUNE
def prepare(ds, shuffle=False, augment=False):
    ds = ds.map(lambda x, y: (resize_and_rescale(x), y), 
                num_parallel_calls=AUTOTUNE)

    if shuffle:
        ds = ds.shuffle(1000)
    ds = ds.batch(batch_size)
    if augment:
        ds = ds.map(lambda x, y: (data_augmentation(x, training=True), y), 
                    num_parallel_calls=AUTOTUNE)
   return ds.prefetch(buffer_size=AUTOTUNE)
train_ds = prepare(train_ds, shuffle=True, augment=True)
val_ds = prepare(val_ds)
test_ds = prepare(test_ds)

Example Model:

A simple convolutional neural network might look like this:

model = Sequential([
    layers.Conv2D(32, (3, 3), padding='same', activation='relu', input_shape=(180, 180, 3)),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(32, activation='relu'),
    layers.Dense(1, activation='sigmoid')  # Corrected activation function for binary classification
])

Compile the model:

model = Sequential([
    layers.Conv2D(32, (3, 3), padding='same', activation='relu', input_shape=(180, 180, 3)),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(32, activation='relu'),
    layers.Dense(1, activation='sigmoid')  # Corrected activation function for binary classification
])

Train for a given number of epochs:

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=1
)

Even with minimal tuning, you’ll notice a slight improvement in validation accuracy after applying data augmentation compared to training on only raw images.

Data Augmentation with tf.image

For finer-grained control, TensorFlow’s tf.image module provides functions to manipulate tensors directly. Common transformations include flipping, cropping, brightness adjustments, saturation changes, and rotation. Each operation is typically a single function call, such as:

flipped = tf.image.flip_left_right(image)
gray = tf.image.rgb_to_grayscale(image)
rotated = tf.image.rot90(image)

To visualize the impact of each transformation, you can write a small helper function that plots the original image and the transformed version side by side.

Applying Randomness with tf.image.stateless_*

tf.image also provides stateless functions like tf.image.stateless_random_brightness, which use a seed parameter to ensure reproducibility. You can generate multiple variations of the same image, each with a different brightness level, simply by changing the seed.

Dataset-Wide Augmentation with .map

Just like Keras layers, you can apply tf.image transformations across the entire training set by defining an augment function and mapping it to your dataset:

def augment(image, label):
    image = tf.cast(image, tf.float32)
    image = tf.image.resize(image, [IMG_SIZE, IMG_SIZE])
    image = image / 255.0
    image = tf.image.random_crop(image, size=[IMG_SIZE, IMG_SIZE, 3])
    image = tf.image.random_brightness(image, max_delta=0.5)
    return image, label
train_ds = (
    train_ds
    .shuffle(1000)
    .map(augment, num_parallel_calls=AUTOTUNE)
    .batch(batch_size)
    .prefetch(AUTOTUNE)
)

This ensures every image fed to the model has some level of randomized transformation applied.

Data Augmentation with ImageDataGenerator

For users who prefer a simpler, more traditional Keras interface, ImageDataGenerator offers an easy way to apply common augmentations, especially when loading images from a directory or a NumPy array.

datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

You can then fit this generator on your training images, and it will yield batches of randomly augmented images. This is especially convenient for small to medium-scale projects.

Data augmentation can significantly improve a model’s generalization ability by experimenting with these tools, adjusting parameters, layering transformations, and systematically evaluating performance.

Data Augmentation Tools