Supercharge Your Image Models with On-The-Fly Augmentation Using Keras ImageDataGenerator

As a deep learning practitioner, you‘ve likely encountered situations where you have a limited number of training images to work with. Maybe you‘re trying to classify some rare plant species, detect manufacturing defects, or build a facial recognition model for a niche application. The challenge is, deep learning models tend to be data-hungry, often requiring thousands or even millions of examples to learn robust, generalizable features. So what do you do when you simply don‘t have a huge dataset available?

Enter image augmentation – a technique that can significantly expand your training data and improve model performance, without requiring you to collect any new images. The idea is simple: take your existing images and programmatically generate new "augmented" versions by applying various transformations like rotations, flips, zooms, shifts, and brightness adjustments. Your model then learns from this expanded dataset, seeing a wider variety of examples that improve its ability to handle variations in real-world data.

While image augmentation is a powerful approach, the traditional way of applying it – generating a bunch of augmented images upfront and saving them to disk – has some downsides. It‘s time consuming, eats up storage space, and still only exposes your model to a fixed set of predetermined augmentations. Wouldn‘t it be great if you could generate augmented images on-the-fly while training your model, exposing it to a virtually infinite stream of random variations? Well, that‘s exactly what Keras‘ ImageDataGenerator allows you to do! Let‘s dive in and see how it works.

The Magic of Keras ImageDataGenerator

The ImageDataGenerator class in Keras is a fantastic utility that makes it super easy to implement a wide variety of image augmentations. The key advantage is that it can dynamically generate batches of augmented images in real-time during model training. This means you get a huge diversity of training examples, without having to create and store the augmented images beforehand.

Here are some of the augmentation techniques ImageDataGenerator supports:

Random Rotations

By specifying a ‘rotation_range‘, you can have ImageDataGenerator randomly rotate images clockwise or counterclockwise by any angle within that range. This is useful for training your model to recognize objects at different orientations. Keep in mind that pixels rotated outside the image bounds will be filled in, either with a constant value, the nearest pixel, or a wrap-around effect, based on the ‘fill_mode‘ parameter.

Width and Height Shifts

The ‘width_shift_range‘ and ‘height_shift_range‘ parameters let you randomly translate the image horizontally or vertically by a fraction of the image‘s width or height. For example, setting width_shift_range to 0.2 would randomly shift the image left or right by up to 20%. This helps the model handle variations in object position.

Flips

Setting ‘horizontal_flip‘ or ‘vertical_flip‘ to True will randomly flip the image along that axis with 50% probability. Flipping can be very useful augmentation for many types of images, but should be used carefully depending on the problem domain. Horizontal flips make sense for natural scenery or objects, but vertical flips are less common and may not be appropriate for things like faces or text.

Brightness

By providing a ‘brightness_range‘, you can randomly adjust the image brightness within that range. This is a nice way to simulate different lighting conditions, making your model more robust. Values less than 1 will darken the image, while values greater than 1 will brighten it.

Zoom

The ‘zoom_range‘ parameter lets you randomly zoom into or out from the image. A value less than 1 zooms in, while a value greater than 1 zooms out. Zooming can help the model learn features at different scales and handle variations in object size.

In addition to these basic transforms, there are many other augmentations available that can be composed in creative ways depending on your specific problem. The Keras documentation provides a great overview of all the options.

Demonstrating ImageDataGenerator

To see ImageDataGenerator in action, let‘s work through an example using a dataset of emergency vs. non-emergency vehicle images. With just ~1600 total images, this is a great case where augmentation can really help boost performance.

The first step is to set up our augmentations. Here‘s an example configuration that performs a variety of transforms:

datagen = ImageDataGenerator(
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        horizontal_flip=True,
        brightness_range=(0.8, 1.2),
        zoom_range=0.2)

Next, we need to actually flow our data through this pipeline. ImageDataGenerator provides a couple convenient methods for this:

flow_from_directory

This method automatically infers the class labels from the directory structure, assuming images for each class are in separate subfolders. We just need to provide the path to the parent directory:

train_generator = datagen.flow_from_directory(
        ‘data/train‘,
        target_size=(150, 150),
        batch_size=32,
        class_mode=‘binary‘,
        seed=42)

flow_from_dataframe

If we have a DataFrame that already contains the image paths and labels, we can use this method instead. Just provide the DataFrame, the path to the directory containing the images, and column names for the image paths and class labels:

train_df = pd.read_csv(‘train.csv‘)
train_generator = datagen.flow_from_dataframe(
        dataframe=train_df,
        directory=‘data/images‘,
        x_col=‘filename‘,
        y_col=‘label‘,
        target_size=(150, 150),
        batch_size=32,
        class_mode=‘binary‘, 
        seed=42)

In both cases, we get a generator object that will yield batches of augmented images and their corresponding labels, ready to be fed into our model during training.

Speaking of models – let‘s build a simple CNN using Keras‘ Sequential API:

model = Sequential([
    Conv2D(32, (3, 3), activation=‘relu‘, input_shape=(150, 150, 3)),
    MaxPooling2D((2, 2)),
    Conv2D(64, (3, 3), activation=‘relu‘),
    MaxPooling2D((2, 2)),
    Conv2D(64, (3, 3), activation=‘relu‘),
    Flatten(),
    Dense(64, activation=‘relu‘),
    Dense(1, activation=‘sigmoid‘)
])

model.compile(optimizer=‘adam‘,
              loss=‘binary_crossentropy‘,
              metrics=[‘accuracy‘])

Now, to train this model using our ImageDataGenerator pipeline, we can use the fit method (which has replaced the fit_generator method used in earlier versions of Keras):

model.fit(
      train_generator,
      steps_per_epoch=train_generator.samples // batch_size,
      epochs=30,
      validation_data=val_generator,
      validation_steps=val_generator.samples // batch_size)

The key thing to note here is that we‘re passing our generator objects (train_generator and val_generator) directly to the fit method. Behind the scenes, Keras will pull batches of augmented images from these generators to train and validate the model.

After training this model for 30 epochs, we get a validation accuracy of around 92% – not bad for such a small dataset! But the real proof is in comparing to a baseline model trained without augmentation. Using the same architecture and hyperparameters, the non-augmented model only achieves around 75% validation accuracy and shows signs of overfitting, with the validation loss diverging from the train loss.

This really highlights the power of augmentation, especially for smaller datasets. By exposing the model to a wide variety of generated examples, we‘re able to learn more robust features and significantly improve generalization performance.

Taking It Further

We‘ve covered the basics of using Keras ImageDataGenerator for on-the-fly augmentation, but there‘s always more to explore and experiment with. Here are a few ideas to take your augmentation pipeline to the next level:

• Try out more advanced augmentation techniques like elastic deformations, random erasing, or mixing images
• Dynamically adjust your augmentation parameters over the course of training, e.g. reducing the intensity towards the end
• Combine ImageDataGenerator with transfer learning to leverage pre-trained models while still benefiting from augmentation
• Implement custom augmentation functions for specific problem domains, like defect simulation for manufacturing applications

The possibilities are endless, and the right augmentation strategy will depend on your specific data and problem. But one thing is clear – ImageDataGenerator is an incredibly powerful tool to have in your deep learning toolkit, and mastering it will help you get the most out of your image models.

So go forth and augment! Your models will thank you.