Use the ImageDataGenerator!

Source: Deep Learning on Medium

Use the ImageDataGenerator!

Do you want to do basic image classification? And only specify the path to your data? Let’s do it with the ImageDataGenerator.

import os
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

Now let’s assume the training data is in “path/to/my/data/train” and the validation data is in “path/to/my/data/validation”.

dir = 'path/to/my/data'
train_dir = os.path.join(dir, 'train')
validation_dir = os.path.join(dir, 'validation')

We initialize the ImageDataGenerator and normalize the values to 0–1.

train_gen = ImageDataGenerator(rescale=1.0/255.)
test_gen = ImageDataGenerator(rescale=1.0/255.)

Now we initialize our train and validation generator. We can choose a target size, use the binary mode, and a batch size from 32. If we have more than two classes, we can use the categorical mode.

train_generator = train_gen.flow_from_directory(
train_dir,
batch_size=32,
class_mode='binary',
target_size=(150, 150))
validation_generator = test_gen.flow_from_directory(
validation_dir,
batch_size=32,
class_mode = 'binary',
target_size = (150, 150))

For the training, we use a convolutional dummy model.

model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(16, (3, 3),
activation='relu',
input_shape=(150, 150, 3)),
tf.keras.layers.MaxPooling2D(2, 2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid')
])

We compile our model and use the Adam optimizer for stochastic gradient descent.

model.compile(optimizer='Adam',
loss='binary_crossentropy',
metrics = ['acc'])

Let’s train our model 🙂

model.fit_generator(train_generator,
validation_data=validation_generator,
steps_per_epoch=75,
epochs=10,
validation_steps=50,
verbose=2)

Done!

But our ImageDataGenerator can do more! We can do data augmentation in our memory on the fly. We can rotate, shift, shear, zoom, flip our images without modifying our original dataset.

ImageDataGenerator(
rescale = 1./255,
rotation_range=40,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=TRUE,
fill_mode='nearest')

The ImageDataGenerator is an easy input pipeline for beginners who want to get started with image classification.