Custom Data Augmentation using Keras ImageDataGenerator

We often run into trouble when training a neural network because we don’t have enough data to train it. And as we all know, most machine learning models, and deep learning models, in particular, are only as good as the training data they are given. Now we can try to collect more data, but collecting data is a time-consuming and tedious task. A solution to this problem is data augmentation.

The contents of this blog are as follows:

1. What is Data Augmentation?
2. Keras ImageDataGenerator
3. Steps to Create Custom Transformations

What is Data Augmentation?

Augmentation is a technique used to increase the size of your dataset by applying different transformations to the existing data and creating new samples using them. As seen in Figure 1, a single image of a dog can be augmented using transformations like rotation and cropping. The aim of augmentation is not just to increase the number of samples in the dataset but also to create variations in the dataset, which will prevent overfitting and help the model generalize better on unseen images.

Figure 1. Augmented Images of a Dog

Keras ImageDataGenerator

An easy way of augmenting data without creating a large overhead is by using the Keras ImageDataGenerator. The ImageDataGenerator generates batches of tensor image data with real-time data augmentation. It ensures that our model will receive variation in data at each epoch, which helps prevent overfitting. At the same time, the ImageDataGenerator uses less memory as it loads all the images in batches rather than all at once.

ImageDataGenerator also provides a wide range of transformations:

1. Rotation
2. Brightness
3. Shear
4. Zoom
5. Horizontal and Vertical Flip

But what if you want to add more transformations? How can you keep all the awesome features of the ImageDataGenerator and add your custom transformations to it? So, in this blog, I will be explaining to you how you can add your custom transformations to ImageDataGenerator.

Steps to Create Custom Transformations

The most efficient way of creating your custom transformations is by creating a Custom Image Data Generator class that inherits from the original ImageDataGenerator class. This class can be easily extended to add new augmentations in the future. Let us get into the code now.

We will be implementing 2 custom transformations: vertical and horizontal motion blur. This will help us understand how multiple transformations can be implemented.

Step 0: Importing Libraries

import cv2
import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt

Step 1: Declare CustomDataGenerator class

Figure 2. Skeleton Code for CustomDataGenerator

As you can see, we extend ImageDataGenerator so that we can specify our parameters in the __init__() method. We declare the processing_function as actions where we will call our transformations horizontal and vertical motion blur. Now let's define our parameters in the __init__() function.

Step 2: __init__() function

Figure 3. Declaring parameters in __init__()

We have declared two parameters here, v_kernel_size and h_kernel_size, for the vertical motion blur and horizontal motion blur, respectively.

Step 3: Writing the Transformation Logic

Figure 4. Logic for Transformations

This is not the most efficient way of implementing these transformations, but for the sake of implementing multiple transformations, we have created two different functions for motion blur. Note that this transformation can be done using a single function as well. You can check out this implementation of motion blur.

Step 4: Action Function Logic

Figure 5. Action Function

Here, we call these our transformations. We need to randomly select one of the transformations for which we create a list of our transformations and generate random values in the range of 0 to 1, corresponding to each transformation. We then select the transformation with the highest value.

Step 5: Putting all the functions together

Figure 6. Complete CustomDataGenrator Script

This is what our CustomDataGenerator will look like. Now it's time to execute this data generator and see what the output looks like. Note that we can still use the transformations present in the ImageDataGenerator as we are inheriting them into our custom data generator.

Figure 7. Executing the CustomDataGenerator

The output of this script can be seen in Figure 8. The point to note here is that we had also included a default parameter of the ImageDataGenerator class horizontal_flip which can be seen in the subplot 1, 2, and 9. Our custom transformations are applied to subplots 3, 4, 5, 6, 7, and 8.

Figure 8. Output of CustomDataGenerator

Conclusion

We have finally built a CustomDataGenerator with our custom transformations and that too, with only a few lines of code. You can try to implement other transformations, like blocking random patches and adding Gaussian noise. I hope this blog will help you create custom augmentations using Keras ImageDataGenerator without any problems. If you have enjoyed reading this blog, do consider following me on Medium for similar content and connecting with me on LinkedIn!

Check out my other articles here!

References:

Image Augmentation Keras | Keras ImageDataGenerator

Tutorial: Custom Data Augmentation in Jeras

tf.keras.preprocessing.image.ImageDataGenerator | TensorFlow v2.10.0