Survey of Keras's image related generators

In François Chollet's Deep learning with Python, he teaches how to build image classification models using ImageDataGenerators. Let's look into how these are made and how to use them.

tf.keras.preprocessing.image.ImageDataGenerator
Generate batches of tensor image data with real-time data augmentation.

Basic usage

Use images in a folder

A standard snippet from his book to prepare the datasets would look like this:

train_datagen = ImageDataGenerator(rescale=1./255, rotation_range=40, width_shift_range=.2, height_shift_range=.2, shear_range=.2, zoom_range=.2, horizontal_flip=True, fill_mode='nearest')
test_datagen = ImageDataGenerator(rescale=1./255)

batch_size=30
target_size=224
train_generator = train_datagen.flow_from_directory(train_dir, target_size=(target_size,target_size), batch_size=batch_size)
validation_generator = test_datagen.flow_from_directory(validation_dir, target_size=(target_size,target_size), batch_size=batch_size)

> Found 3512 images belonging to 6 classes.
> Found 633 images belonging to 6 classes.

ImageDataGenerator generates an iterator of batches of images and labels using a folder's directory structure. It also takes a few parameters that make the iterator apply random transformations. This helps prevent overfitting.

By default it assumes a categorical classification model. If you want binary, set class_mode to "binary".

flow_from_directory creates the DirectoryIterator.

Train a model with a pre-trained model

Then to train the model, you could use a pre-trained image classification model, freeze it, put a Dense network on top, and train it.

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.applications import Xception
conv_base = Xception(include_top=False, input_shape=(target_size, target_size, 3))
# Freeze the pre-trained model's base
conv_base.trainable = False

model = models.Sequential()
model.add(conv_base)
model.add(layers.Flatten())
model.add(layers.Dense(512, activation='relu'))
model.add(layers.Dense(train_generator.num_classes, activation='softmax'))
model.compile(loss='categorical_crossentropy', metrics=['accuracy'])
model.summary()
history = model.fit(train_generator, steps_per_epoch=len(train_generator),
    epochs=20, validation_data=validation_generator, validation_steps=len(validation_generator),
    callbacks=[tf.keras.callbacks.EarlyStopping(patience=1)])

DirectoryIterator

tf.keras.preprocessing.image.DirectoryIterator
Iterator capable of reading images from a directory on disk.

Printing out the first batch

If you want to iterate through the Iterator manually to inspect its data, you could use a loop:

for x_batch, y_batch in directoryIterator:
    print(x_batch)
    print(y_batch)
    break;

This will print out the first batch. If you didn't put the break, it would loop infinitely.

Getting steps per epoch

If you want to know how many steps you'd need to cover one epoch during training, you'd need to know how many images are being provided and divide that by the batch size. The DirectoryIterator actually does that for us in __len__.

Looping through the data once

On top of __len__, DirectoryIterator provides __getitem__. So one simple way to loop through it is:

for i in range(len(directoryIterator)):
    x_batch, y_batch = directoryIterator[i]
    ...

Another way, which is more verbose, is to do:

i = 0
steps_per_epoch = len(directoryIterator)
for (x_batch, y_batch) in directoryIterator:
    ...
    i += 1
    if (i == steps_per_epoch):
        break

Finding the number of classes

I wasn't able to find where in the source code it was computed, but this works:

directoryIterator.num_classes

Finding the mapping of class index to class names

As per the doc of flow_from_directory:

directoryIterator.class_indices

Determining if it's binary or categorical

One could look at num_classes, but that doesn't cover the case of a categorical model with only two classes. Instead check its class_mode.

directoryIterator.class_mode