|
import gradio as gr |
|
import tensorflow as tf |
|
import numpy as np |
|
from tensorflow.keras.preprocessing import image |
|
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input |
|
from PIL import Image |
|
|
|
|
|
MODEL_PATH = "setosys_dogs_model.h5" |
|
model = tf.keras.models.load_model(MODEL_PATH) |
|
|
|
|
|
class_labels = list(model.class_indices.keys()) |
|
|
|
|
|
def preprocess_image(img: Image.Image) -> np.ndarray: |
|
"""Preprocess the image to match the model's input requirements.""" |
|
img = img.resize((224, 224)) |
|
img_array = np.array(img) |
|
img_array = preprocess_input(img_array) |
|
img_array = np.expand_dims(img_array, axis=0) |
|
return img_array |
|
|
|
|
|
def predict_dog_breed(img: Image.Image) -> dict: |
|
"""Predict the breed of the dog in the uploaded image.""" |
|
img_array = preprocess_image(img) |
|
predictions = model.predict(img_array) |
|
|
|
|
|
print("Predictions Shape:", predictions.shape) |
|
|
|
class_idx = np.argmax(predictions) |
|
confidence = float(np.max(predictions)) |
|
|
|
|
|
predicted_breed = class_labels[class_idx] if class_idx < len(class_labels) else "Unknown" |
|
|
|
return {predicted_breed: confidence} |
|
|
|
|
|
interface = gr.Interface( |
|
fn=predict_dog_breed, |
|
inputs=gr.Image(type="pil"), |
|
outputs=gr.Label(), |
|
title="Dog Breed Classifier", |
|
description="Upload an image of a dog to predict its breed." |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
interface.launch() |
|
|