Writing Custom Keras Generators

The idea behind using a Keras generator is to get batches of input and corresponding output on the fly during training process, e.g. reading in 100 images, getting corresponding 100 label vectors and then feeding this set to the gpu for training step.

The problem I faced was memory requirement for the standard Keras generator. The in memory generator creates copies of the orginal data as well as has to convert the dtype from uint8 to float64. On the other hand, the keras generator to read from directory expects images in each class to be in an independent directory (Not possible in multilabel problems, segmentation problems etc.)

So, I like to split the batch genrator into 4 steps:

1. Get input : input_path -> image
2. Get output : input_path -> label
3. Preprocess input : image -> pre-processing step -> image
4. Get generator output : (batch_input, batch_labels )

Step 1 : Define a function to get input (can be subsetting an array, reading in from disk etc.) :

from skimage.io import imread
def get_input(path):

img = imread(path)


Step 2 : Define a function to get output :

import numpy as np
import pandas as pd
def get_output(path,label_file=None):

img_id = path.split('/')[-1].split('.')[0]
img_id = np.int64(img_id)
labels = label_file.loc[img_id].values


Step 3 : Define a function to preprocess input :

def preprocess_input(image):

--- Rescale Image
--- Rotate Image
--- Resize Image
--- Flip Image
--- PCA etc.

Step 4 : Bring everything together to define your generator :

def image_generator(files,label_file, batch_size = 64):

while True:
# Select files (paths/indices) for the batch
batch_paths = np.random.choice(a = files, size = batch_size)
batch_input = []
batch_output = []
# Read in each input, perform preprocessing and get labels
for input_path in batch_paths:
input = get_input(input_path )
output = get_output(input_path,label_file=label_file )

input = preprocess_input(image=input)
batch_input += [ input ]
batch_output += [ output ]
# Return a tuple of (input,output) to feed the network
batch_x = np.array( batch_input )
batch_y = np.array( batch_output )

yield( batch_x, batch_y )

And there you have it, you can add to preprocessing function as defined by your specific dataset , have the output as a image mask ( segmentation problems, localization problems etc. ),

After this, using this custom generator is the same as using a predefined keras ImageDataGenerator and simply pass the generator object to model.fit_generator().

Source: Deep Learning on Medium