The easiest and most obvious way is to modify the Keras code (you can find full implementation of ImageDataGenerator in here). However, Keras does provide an elegant API to deal with this, although it is not very well documented.
Writing a generator
We would need to write a new Keras generator, inheriting from Iterator class. The Iterator class itself is just a convenient child class from Sequence, a detailed tutorial of which can be found here.
from keras.preprocessing.image import Iterator, ImageDataGenerator
class MyIterator(Iterator):
"""This is a toy example of a wrapper around ImageDataGenerator"""
def __init__(self, n, batch_size, shuffle, seed, **kwargs):
super().__init__(n, batch_size, shuffle, seed)
# Load any data you need here (CSV, HDF5, raw stuffs). The code
# below is just a pseudo-code for demonstration purpose.
input_images = ...
ground_truth = ...
# Here is our beloved image augmentator <3
self.generator = ImageDataGenerator(**kwargs)
def _get_batches_of_transformed_samples(self, index_array):
"""Gets a batch of transformed samples from array of indices"""
# Get a batch of image data
batch_x = input_images[index_array].copy()
batch_y = ground_truth[index_array].copy()
# Transform the inputs and correct the outputs accordingly
for i, (x, y) in enumerate(zip(batch_x, batch_y)):
transform_params = self.generator.get_random_transform(x.shape)
batch_x[i] = self.generator.apply_transform(x, transform_params)
batch_y[i] = process_outputs_accordingly(y, transform_params)
return batch_x, batch_y
Why we inherit from Keras Iterator?
It is recommended to have your generators inherited from keras.utils.Sequence (also the base class of other classes). It allows the data be loaded in parallel between several threads.
I want the same training API!
You can write a custom generator with methods flow, flow_from_directory, and flow_from_dataframe --- core functions of Keras API.
