Image Quality Classification — From Training to Production (Part 1)

4 min readJun 21, 2019

In this tutorial, we are going to train a neural network classifier to predict if a provided image is of “Good Quality” or “Bad Quality” using Transfer Learning technique with Keras and Python.

Introduction

Recently i was tasked to solve a problem for a client. At the registration stage of the client’s product, the user is required to upload an image of their valid ID Card (international passport, voters card, driver’s license, national identity card), during processing, employees realized that many of the images uploaded were of bad quality. After identifying the bad imagess, the employee will have to deny access to the user and ask the user to re-upload the ID. This is not a very good customer experience. So we decided to build a system that can quickly access the quality of an image at the moment of upload during registration.

The first part of this article will provide a straight forward guide to training image classification model by fine-tuning a pre-trained Neural Network (Transfer Learning) using Keras and Tensorflow. The second part will show you how to serve your model into production using Tensorflow-serving and Flask.

Data Overview

The training data used in this project consists of two classes, “good” and “bad” with 250 images each. The test data contains the same classes with 32 images each.

Transfer Learning

Due to low amount of training images, transfer learning technique was introduced. Transfer Learning is a technique widely used by computer vision practitioners, which allows you to build accurate models from already trained models instead of building from scratch. This already trained models are usually called Pre-Trained Models. Pre-trained models are models trained on a large benchmark on a task similar to our own. It often requires large data, time and energy to build a model from scratch. Initializing a model from a pre-trained model give faster convergence and also saves time and energy.

Fine tuning involves loading a pre-trained model and retraining only a few layers. This approach works best when the problem you are trying to solve is different from the images the model was trained on. This is a big advantage when dealing with smaller data.

Step 1: Generate Bottleneck Features

Bottleneck features are features generated from complex pre-trained models. The images are passed forward through the model and the pre final layer is stored.

def save_bottlebeck_features():
    datagen = ImageDataGenerator(rescale=1. / 255)# build the VGG16 network
    model = tf.keras.applications.VGG16(include_top=False, weights='imagenet')generator = datagen.flow_from_directory(
        train_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode=None,
        shuffle=False)
    
    bottleneck_features_train = model.predict_generator(
        generator, 32)
    
    np.save(open('bottleneck_features_train.npy', 'wb'),
            bottleneck_features_train)generator = datagen.flow_from_directory(
        validation_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode=None,
        shuffle=False)
    
    bottleneck_features_validation = model.predict_generator(
        generator, 8)
    
    np.save(open('bottleneck_features_validation.npy', 'wb'),
    bottleneck_features_validation)

Step 2: Train Top Model

Next, we’ll train a classifier on the bottleneck features and save the weights. This will act as top model for the fine tuned model.

def train_top_model():
    import keras
    train_data = np.load(open('bottleneck_features_train.npy', 'rb'))
    train_labels = np.array([0] * int(no_train / 2) + [1] * int(no_train / 2))validation_data = np.load(open('bottleneck_features_validation.npy', 'rb'))
    validation_labels = np.array([0] * int(no_test / 2) + [1] * int(no_test / 2))
    
    tbCallBack = tf.keras.callbacks.TensorBoard(log_dir='./Graph', histogram_freq=0,  
          write_graph=True, write_images=True)model = Sequential()
    model.add(Flatten(input_shape=train_data.shape[1:]))
    model.add(Dense(8, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(1, activation='sigmoid'))model.compile(optimizer=tf.keras.optimizers.SGD(lr=0.001),
                  loss='binary_crossentropy', metrics=['accuracy'])model.fit(train_data, train_labels,
              epochs=epochs,
              batch_size=batch_size,
              validation_data=(validation_data, validation_labels),
              callbacks=[tbCallBack])
    
    model.save_weights(top_model_weights_path)

Step 3: Design Model Architecture

load the VGG pre-trained model

model = tf.keras.applications.VGG16(include_top=False, input_shape=input_shape)

2. Define Fine-Tune Model

model_fine_tune = tf.keras.models.Sequential()
model_fine_tune.add(tf.keras.layers.Flatten(input_shape=model.output_shape[1:]))
model_fine_tune.add(tf.keras.layers.Dense(8, activation='relu'))
model_fine_tune.add(tf.keras.layers.Dropout(0.5))
model_fine_tune.add(tf.keras.layers.Dense(1, activation='sigmoid'))

3. Load top model weights and top model is then added to convolutional base.

from tensorflow.keras.models import Model
model_fine_tune.load_weights(top_model_weights_path)
model = Model(inputs= model.input, outputs= model_fine_tune(model.output))

4. We can set the top 18 layers of the model to be non trainable so we do not update the weights. Only the remaining layers will be updated.

for vgg_layer in model.layers[:18]:
    vgg_layer.trainable = False

5. Compile the model

model.compile(loss='binary_crossentropy',
              optimizer=tf.keras.optimizers.Adam(lr=1e-5),
metrics=['accuracy'])

6. Train the model

import datetime
from tensorflow.keras.callbacks import ModelCheckpoint# filepath to save checkpoint
filepath = "ckpt/weights.best.hdf5"# Monitor the validation loss and save the best 
checkpoint = ModelCheckpoint(filepath, monitor='val_loss', verbose=1, save_best_only=True, mode='min')# Visualize model training using tensorboard
log_dir="Graph/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=True, write_images=True)# Fit model
model.fit_generator(
    generator_train,
    epochs=epochs,
    validation_data=generator_test, 
    callbacks=[tensorboard_callback, checkpoint])

Results

Conclusion

We have successfully trained an image classifier to detect the quality of an ID Card image using popular computer vision technique called transfer learning. We also looked at how we can fine tune a pre-trained model and the benefits of taking such approach. In the next part we would serve our trained model as an API (Application Programming Interface) using the convenient tensorflow-serving and Flask.

I hope this tutorial provided a guide to developing your own unique computer vision solutions. I’m glad to be of help. If you have any questions or suggestions, please let me know.