File size: 2,945 Bytes
6cb8839 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import gradio as gr
import os
import tensorflow
from typing import Tuple, Dict
import tensorflow as tf
import numpy as np
from PIL import Image
from timeit import default_timer as timer
# Setup class names
with open('class_names.txt') as f:
class_names = [breed.strip() for breed in f.readlines()]
# Load the pre-trained EfficientNetV2 model
effnet = tensorflow.keras.models.load_model('demo/dog_breed_classifier/dog_breed_effnet_augmentation.h5')
# Get the preprocessing function for EfficientNetV2
effnet_preprocess_input = tensorflow.keras.applications.efficientnet_v2.preprocess_input
# Create examples list from "examples/" directory
example_list = [["examples/" + example] for example in os.listdir("examples")]
## `predict` function
def predict(img) -> Tuple[Dict, float]:
"""
Transforms and performs a prediction on an image and returns prediction and time taken.
Args:
image_path (str): Path to the input image.
Returns:
Tuple[Dict, float]: A tuple containing a dictionary of class labels and prediction probabilities
and the prediction time.
"""
# Start the timer
start_time = timer()
# Open the image using PIL
img = img.resize((224, 224)) # Resize the image to the model's expected input size
# Convert the image to a NumPy array
x = np.array(img)
x = effnet_preprocess_input(x)
# Add a batch dimension
x = tf.expand_dims(x, axis=0)
# Pass the image through the model
predictions = effnet(x)
top_classes_indices = np.argsort(predictions[0])[::-1][:3] # Get the indices of top 3 classes
top_classes = [class_names[i] for i in top_classes_indices] # Get the class names of top 3 classes
top_probabilities = [predictions[0][index] for index in top_classes_indices] * 100 # Get the probabilities of top 3 classes
# Create a dictionary of class labels and prediction probabilities
pred_labels_and_probs = {top_classes[i]: float(top_probabilities[i]) for i in range(len(top_classes_indices))}
# Calculate the prediction time
pred_time = round(timer() - start_time, 5)
# Return the prediction dictionary and prediction time
return pred_labels_and_probs, pred_time
# Create title, description, and article strings
title = "🐶 Dog Breeds Classifier 🐾"
description = "🚀 An EfficientNetV2S feature extractor computer vision model to classify images of 120 different breeds. 📸"
article = "🌟 Created at [GitHub](https://github.com/adinmg/dog_breed_classifier)."
# Create the Gradio demo
demo = gr.Interface(
fn=predict, # mapping function from input to output
inputs=gr.Image(type="pil"),
outputs=[
gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?
gr.Number(label="Prediction time (s)"),
],
examples=example_list,
title=title,
description=description,
article=article,
)
# Launch the demo!
demo.launch()
|