Save the date! Google I/O returns May 18-20 Register now

Image Classification with TensorFlow Hub

In this colab, you'll try multiple image classification models from TensorFlow Hub and decide which one is best for your use case.

Because TF Hub encourages a consistent input convention for models that operate on images, it's easy to experiment with different architectures to find the one that best fits your needs.

View on TensorFlow.org View on GitHub Download notebook See TF Hub models
import tensorflow as tf
import tensorflow_hub as hub

import requests
from PIL import Image
from io import BytesIO

import matplotlib.pyplot as plt
import numpy as np

Helper functions for loading image (hidden)

Select an Image Classification Model. After that, some internal variables are set and the labels file is downloaded and prepared for use.

There are some technical differences between the models, like different input size, model size, accuracy, and inference time. Here you can change the model you are using until you find the one most suitable for your use case.

The handle (url) of the model is printed for your convenience. More documentation about each model is available there.

Select an Image Classification model

Selected model: mobilenet_v3_small_100_224 : https://hub.tensorflow.google.cn/google/imagenet/mobilenet_v3_small_100_224/classification/5
Images will be converted to 224x224
Downloading data from https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt
16384/10484 [==============================================] - 0s 0us/step

You can select one of the images below, or use your own image. Just remember that the input sizes for the models vary and some of them use a dynamic input size (enabling inference on the unscaled image). Given that, the method load_image will already rescale the image to the expected format.

Select an Input Image

png

Now that the model was chosen, loading it with TensorFlow Hub is simple.

This also calls the model with a random input as a "warmup" run. Subsequent calls are often much faster, and you can compare this with the latency below.

classifier = hub.load(model_handle)

input_shape = image.shape
warmup_input = tf.random.uniform(input_shape, 0, 1.0)
%time warmup_logits = classifier(warmup_input).numpy()
CPU times: user 1.66 s, sys: 8.52 s, total: 10.2 s
Wall time: 10.5 s

Everything is ready for inference. Here you can see the top 5 results from the model for the selected image.

# Run model on image
%time probabilities = tf.nn.softmax(classifier(image)).numpy()

top_5 = tf.argsort(probabilities, axis=-1, direction="DESCENDING")[0][:5].numpy()
np_classes = np.array(classes)

# Some models include an additional 'background' class in the predictions, so
# we must account for this when reading the class labels.
includes_background_class = probabilities.shape[1] == 1001

for i, item in enumerate(top_5):
  class_index = item if not includes_background_class else item - 1
  line = f'({i+1}) {class_index:4} - {classes[class_index]}: {probabilities[0][top_5][i]}'
  print(line)

show_image(image, '')
CPU times: user 5.51 ms, sys: 375 µs, total: 5.88 ms
Wall time: 5.39 ms
(1)   34 - leatherback turtle: 0.775492787361145
(2)  980 - volcano: 0.07103672623634338
(3)   33 - loggerhead: 0.02250688150525093
(4)  106 - wombat: 0.007518391124904156
(5)  276 - hyena: 0.0055127376690506935

png

Learn More

If you want to learn more and try how to do Transfer Learning with these models you can try this tutorial: Transfer Learning for Image classification

If you want to check on more image models you can check them out on hub.tensorflow.google.cn