Train an ML model for custom object Image Classification

(Using Google Colab)

In this article, I will train a machine learning model for custom object image classification using TensorFlow 2.x in Google Colab. Following is the roadmap for it.

Roadmap

Collect the dataset of images.
Setup the environment, mount drive, and create a folder for backup on drive.
Pre-process the model, compile the model and finally train the model.
Evaluate the model.
Export the model

Here we will train a model for mask image classification. This is done in 12 steps mentioned in the section below:

Import dependencies and libraries
Mount drive and link your folder
Create custom_ic folder in your google drive
Upload your dataset to your drive and unzip it in the Content directory
Select Variables
Preprocess data
Compile model
Load Tensorboard
Train the model
Check the predictions
Export your model
Test the saved model by reloading

HOW TO BEGIN?

Open my Colab notebook on your browser.
Click on File in the menu bar and click on Save a copy in drive. This will open a copy of my Colab notebook on your browser which you can now use.
Next, once you have opened the copy of my notebook and are connected to the Google Colab VM, click on Runtime in the menu bar and click on Change runtime type. Select GPU and click on save.

LET’S BEGIN !!

1) Setup (Import dependencies and libraries)

from __future__ import absolute_import, division, print_function, unicode_literals

import matplotlib.pylab as plt
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras import layers
import numpy as np

import keras as ker

print(ker.__version__)

2) Mount drive and link your folder

# mount drive

from google.colab import drive
drive.mount('/content/gdrive')

# this creates a symbolic link so that now the path /content/gdrive/My Drive/ is equal to /mydrive

!ln -s /content/gdrive/My Drive/ /mydrive
!ls /mydrive

**3) Create custom_ic folder in your google drive**

Create a folder named custom_ic for your custom image classification files in your google drive.

4) Upload your dataset to your drive and unzip it in the Content directory

# Upload dataset to the custom_ic folder in your drive and copy it to the content dir
!cp /mydrive/custom_ic/obj.zip /content

# Unzip the dataset
!unzip -qq "obj.zip"

Simple transfer learning

5) Select Variables

Select Model

#@title Enter Variables

Type_of_model = 'efficient' #@param ["efficient", "mobilenetv2", "inceptionv3", "inceptionResnetv2"]

How to Choose Model

# https://paperswithcode.com/sota/image-classification-on-imagenet

# Mobilenet : less number of classes and needed low size model in less time.
# Inception : More accuracy than Mobilenet inceptionResnet > inception v3 but v3 has less number so small size.
# if need small size more accuracy go for v3, if less size issue then inceptionResnet

# Efficient : More Size, More Accuracy, More time to train. 


# Advice : Efficient is STATE OF THE ART,  USE IT 

# There are other State of art models but their packages are not available now


# Accuracy : 
# EfficientNet > inceptionResnetv2 > InceptionNet >= Mobilenet  

# Size : 
# Mobilenet < inceptionNet < inceptionResnetv2 < EfficientNet

#https://tfhub.dev/google/efficientnet/b3/classification/1
model_link = {"efficient":"https://tfhub.dev/google/efficientnet/b3/classification/1", 
              "mobilenetv2" : "https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4", 
              "inceptionv3" : "https://tfhub.dev/google/tf2-preview/inception_v3/classification/4", 
              "inceptionResnetv2" : "https://tfhub.dev/google/imagenet/inception_resnet_v2/classification/4"}

model_shape = {"efficient":(300,300), 
              "mobilenetv2" : (224,224), 
              "inceptionv3" : (300,300), 
              "inceptionResnetv2" : (300,300)
}

6) Preprocess data

IMAGE_SHAPE = model_shape[Type_of_model]
image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1/255, validation_split=0.3)

train_gen = image_generator.flow_from_directory('obj', target_size=IMAGE_SHAPE,  subset='training')
val_gen = image_generator.flow_from_directory('obj', target_size=IMAGE_SHAPE, subset='validation')


for image_batch, label_batch in train_gen:
  print("Image batch shape: ", image_batch.shape)
  print("Label batch shape: ", label_batch.shape)
  break
for image_batch, label_batch in val_gen:
  print("Image batch shape: ", image_batch.shape)
  print("Label batch shape: ", label_batch.shape)
  break


feature_extractor_url= model_link[Type_of_model] 

feature_extractor_layer = hub.KerasLayer(feature_extractor_url,
                                         input_shape=(IMAGE_SHAPE[0],IMAGE_SHAPE[1],3))

feature_batch = feature_extractor_layer(image_batch)
print(feature_batch.shape)

feature_extractor_layer.trainable = False

7) Compile model

model = tf.keras.Sequential([
  feature_extractor_layer,
  layers.Dense(2, activation='softmax')
])

model.summary()

predictions = model(image_batch)
predictions.shape


#Use compile to configure the training process:

model.compile(
  optimizer=tf.keras.optimizers.Adam(),
  loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
  metrics=['acc'])

Set model directory and checkpoint

import os

model_dir = "/mydrive/custom_ic/training/"
checkpoint_dir = "weights"

os.makedirs(model_dir, exist_ok=True)
# os.makedirs(checkpoint_dir, exist_ok=True)



class CollectBatchStats(tf.keras.callbacks.Callback):
  def __init__(self):
    self.batch_losses = []
    self.batch_acc = []

  def on_train_batch_end(self, batch, logs=None):
    self.batch_losses.append(logs['loss'])
    self.batch_acc.append(logs['acc'])
    self.model.reset_metrics()


from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint

checkpoint = ModelCheckpoint(checkpoint_dir, monitor='val_loss', verbose=0)
tensorboard = TensorBoard(log_dir=model_dir)

8) Load tensorboard

%load_ext tensorboard
%tensorboard --logdir "training/"

9) Train the model

%steps_per_epoch = np.ceil(train_gen.samples/train_gen.batch_size)

batch_stats_callback = CollectBatchStats()

# If you want to save every checkpoint change callbacks = [batch_stats_callback, tensorboard, checkpoint]

history = model.fit_generator(train_gen, validation_data=val_gen, epochs=11,
                              steps_per_epoch=steps_per_epoch,
                              callbacks = [batch_stats_callback, tensorboard])

#AUTO-CLICK TO AVOID BEING KICKED OFF COLAB
#Press (Ctrl + Shift + i) . Go to console. Paste the following code and press Enter.

function ClickConnect(){
console.log("Working"); 
document
  .querySelector('#top-toolbar > colab-connect-button')
  .shadowRoot.querySelector('#connect')
  .click() 
}
setInterval(ClickConnect,60000)

Check training progress

Now after, even just a few training iterations, we can already see that the model is making progress on the task.

Plot loss graph

plt.figure()
plt.ylabel("Loss")
plt.xlabel("Training Steps")
plt.ylim([0,2])
plt.plot(batch_stats_callback.batch_losses)

Plot accuracy graph

plt.figure()
plt.ylabel("Accuracy")
plt.xlabel("Training Steps")
plt.ylim([0,1])
plt.plot(batch_stats_callback.batch_acc)

10) Check the predictions

To redo the plot from before, first get the ordered list of class names:

import numpy as np
class_names = sorted(val_gen.class_indices.items(), key=lambda pair:pair[1])
class_names = np.array([key.title() for key, value in class_names])
class_names

Run the image batch through the model and convert the indices to class names.

predicted_batch = model.predict(image_batch)
predicted_id = np.argmax(predicted_batch, axis=-1)
predicted_label_batch = class_names[predicted_id]

Plot the result

label_id = np.argmax(label_batch, axis=-1)

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  color = "green" if predicted_id[n] == label_id[n] else "red"
  plt.title(predicted_label_batch[n].title(), color=color)
  plt.axis('off')
_ = plt.suptitle("Model predictions (green: correct, red: incorrect)")

11) Export your model

Now that you’ve trained the model, export it as a saved model:

# Saving as tensorflow file

export_path = "saved_models"
model.save(export_path, save_format='tf')

export_path

model.save('/mydrive/custom_ic/image_classification_model.h5')

import tensorflow as tf
sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(log_device_placement=True))

my_reloaded_model = tf.keras.models.load_model(('/mydrive/custom_ic/image_classification_model.h5'),custom_objects={'KerasLayer':hub.KerasLayer})

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

12) Test the saved model by reloading

Train a model for Image Classification

Reload the model

export_path = "saved_models"
reloaded = tf.keras.models.load_model(export_path)

Run inference to test the model

import cv2
from google.colab.patches import cv2_imshow


# Enter your labels here

label=['with_mask','without_mask']

# Read an image
image = cv2.imread('/mydrive/mask_test_images/image1.jpg')

crop_img = cv2.resize(image, IMAGE_SHAPE)
crop_img = crop_img/255
crop_img = crop_img.reshape(1,IMAGE_SHAPE[0],IMAGE_SHAPE[1],3)

result = reloaded.predict(crop_img)
print(' result',result)
predicted_id = np.argmax(result, axis=-1)


print(predicted_id)
id=predicted_id[0]
print(label[id])
cv2_imshow(image)

To learn how to use this trained model to build an app read the article Build Android app for custom object image classification using google colab.

Mask dataset

I have used the following dataset for test purposes. You can use other datasets. See more datasets sources mentioned below in credits section.

Prajnasb Github

My Colab notebook for this

Train image classification

Check out my Youtube Video on this

Coming Soon!

CREDITS

Documentation / References

Dataset Sources

You can download datasets for many objects from the sites mentioned below. These sites also contain images of many classes of objects along with their annotations/labels in multiple formats such as the YOLO_DARKNET txt files and the PASCAL_VOC xml files.

More Mask Datasets

Prasoonkottarathil Kaggle (20000 images)
Ashishjangra27 Kaggle (12000 images )