Image Classification Model MobileNet V2 from TensorFlow Hub

In this blog, we will use models from TensorFlow Hub and classify a image with pre-trained model MobileNet V2.

What is TensorFlow Hub?

TensorFlow Hub is a repository of pre-trained TensorFlow models.

An ImageNet classifier?

An Imagenet classifier is pre-trained model on the ImageNet benchmark dataset. No initial training required for this classifier. In this example we will use MobileNetV2 pre-trained classifier from TensorFlow Hub.

MobileNet V2

MobileNet V2 is a family of neural network architectures for efficient on-device image classification and related tasks, originally published by

Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen: "Inverted Residuals and Linear Bottlenecks: Mobile Networks for Classification, Detection and Segmentation", 2018.

Mobilenets come in various sizes controlled by a multiplier for the depth (number of features) in the convolutional layers. They can also be trained for various sizes of input images to control inference speed.

This TF-Hub module uses the TF-Slim implementation of mobilenet_v2 with a depth multiplier of 1.0 and an input size of 224x224 pixels.

The module contains a trained instance of the network, packaged to get feature vectors from images. If you want the full model including the classification it was originally trained for, use module google/tf2-preview/mobilenet_v2/classification/4 instead.

Download classifier

mobilenet_v2 ="https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4"

Create an instance of MobileNetV2 model

classifier_model = mobilenet_v2

Import Libraries

import tensorflow as tf
import tensorflow_hub as hub

import PIL.Image as Image
import numpy as np
import matplotlib.pylab as plt

Wrap classifier model as Keras layer with hub.KerasLayer

Keras is TensorFlow's high-level API for building deep learning models by composing Keras Layer objects. The tensorflow_hub library provides the class hub.KerasLayer that gets initialized with the URL (or filesystem path) of a SavedModel and then provides the computation from the SavedModel, including its pre-trained weights.

IMAGE_SHAPE = (224, 224)

classifier = tf.keras.Sequential([
    hub.KerasLayer(classifier_model, input_shape=IMAGE_SHAPE+(3,))
])

Open a image and resize

airbus = 'input/airbus.jpg'
airbus = Image.open(airbus).resize(IMAGE_SHAPE)
airbus

Rescale image

airbus

airbus = np.array(airbus)/255.0
airbus.shape

(224, 224, 3)

airbus

array([[[1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        ...,
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [0.92156863, 0.92156863, 0.92156863]],

       [[1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        ...,
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [0.92156863, 0.92156863, 0.92156863]],

       [[1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        ...,
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [0.92156863, 0.92156863, 0.92156863]],

       ...,

       [[1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        ...,
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [0.92156863, 0.92156863, 0.92156863]],

       [[1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        ...,
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [0.92156863, 0.92156863, 0.92156863]],

       [[1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        ...,
        [1.        , 1.        , 1.        ],
        [1.        , 1.        , 1.        ],
        [0.92156863, 0.92156863, 0.92156863]]])

Add a batch dimension with np.newaxis

airbus_expended_dim = airbus[np.newaxis, ...]
airbus_expended_dim.shape

(1, 224, 224, 3)

Predict

result = classifier.predict(airbus_expended_dim)
result.shape

(1, 1001)

The result is a 1001-element vector of logits, rating the probability of each class for the image.

Get the top class ID

predicted_class = tf.math.argmax(result[0], axis=-1)
predicted_class

<tf.Tensor: shape=(), dtype=int64, numpy=405>

Decode the predictions

Download the ImageNet dataset labels

labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
labels_path

'/home/jupyter-thakur/.keras/datasets/ImageNetLabels.txt'

Read the ImageNet dataset labels

imagenet_labels = np.array(open(labels_path).read().splitlines())
imagenet_labels

array(['background', 'tench', 'goldfish', ..., 'bolete', 'ear',
       'toilet tissue'], dtype='<U30')

Plot the image and predicted class

plt.imshow(airbus)
predicted_class_name = imagenet_labels[predicted_class]
plt.title("Prediction: " + predicted_class_name.title())
plt.show()