tensorflow_hub_mobilenet_image_classification_with_imagenet_labels.py

1import tensorflow as tf
2import tensorflow_hub as hub
3import numpy as np
4import PIL.Image as Image
5
6# 1. Load a pre-trained MobileNetV2 classification model from TF Hub
7model_url = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4"
8model = tf.keras.Sequential([
9    hub.KerasLayer(model_url, input_shape=(224, 224, 3))
10])
11
12# 2. Download and preprocess a sample image (a Grace Hopper image)
13image_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg"
14image_path = tf.keras.utils.get_file('image.jpg', image_url)
15img = Image.open(image_path).resize((224, 224))
16img_array = np.array(img) / 255.0  # Normalize to [0, 1]
17img_array = img_array[np.newaxis, ...]  # Add batch dimension
18
19# 3. Run prediction
20predictions = model.predict(img_array)
21predicted_class = np.argmax(predictions[0], axis=-1)
22
23# 4. Fetch the ImageNet labels to map the ID to a name
24labels_path = tf.keras.utils.get_file('ImageNetLabels.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
25imagenet_labels = np.array(open(labels_path).read().splitlines())
26
27print(f"Predicted class ID: {predicted_class}")
28print(f"Predicted label: {imagenet_labels[predicted_class]}")