Spaces:
Sleeping
Sleeping
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
from PIL import Image | |
model_path = "DogClassifier2.5.keras" | |
model = tf.keras.models.load_model(model_path) | |
# Define the core prediction function | |
def predict_breed(image): | |
# Preprocess image | |
image = Image.fromarray(image.astype('uint8')) # Convert numpy array to PIL image | |
image = image.convert("RGB") # Ensure the image is in RGB format | |
image = image.resize((150, 150)) # Resize the image to 150x150 | |
image = np.array(image) | |
image = np.expand_dims(image, axis=0) # Add batch dimension | |
# Predict | |
prediction = model.predict(image) | |
# Apply softmax to get probabilities for each class | |
prediction = tf.nn.softmax(prediction) | |
# Print prediction probabilities to the console | |
print("Prediction probabilities:", prediction) | |
# Define class names | |
class_names = ['Afghan', 'African Wild Dog', 'Airedale', 'American Hairless', 'American Spaniel', 'Basenji', 'Basset', 'Beagle', | |
'Bearded Collie', 'Bermaise', 'Bichon Frise', 'Blenheim', 'Bloodhound', 'Bluetick', 'Border Collie', 'Borzoi', | |
'Boston Terrier', 'Boxer', 'Bull Mastiff', 'Bull Terrier', 'Bulldog', 'Cairn', 'Chihuahua', 'Chinese Crested', | |
'Chow', 'Clumber','Cockapoo', 'Cocker', 'Collie', 'Corgi', 'Coyote', 'Dalmation', 'Dhole', 'Dingo', 'Doberman', | |
'Elk Hound', 'French Bulldog', 'German Sheperd', 'Golden Retriever', 'Great Dane', 'Great Perenees', 'Greyhound', | |
'Groenendael', 'Irish Spaniel', 'Irish Wolfhound', 'Japanese Spaniel', 'Komondor', 'Labradoodle', 'Labrador', | |
'Lhasa', 'Malinois', 'Maltese', 'Mex Hairless', 'Newfoundland', 'Pekinese', 'Pit Bull', 'Pomeranian', | |
'Poodle', 'Pug', 'Rhodesian', 'Rottweiler', 'Saint Bernard', 'Schnauzer', 'Scotch Terrier', 'Shar_Pei', | |
'Shiba Inu', 'Shih-Tzu', 'Siberian Husky', 'Vizsla', 'Yorkie'] | |
# Apply threshold and set probabilities lower than 0.015 to 0.0 | |
threshold = 0.01395 | |
prediction = np.array(prediction) | |
prediction[prediction < threshold] = 0.0 | |
# Recalculate the probabilities | |
total_probability = np.sum(prediction) | |
if total_probability > 0: | |
prediction = prediction / total_probability | |
# Create a dictionary with the probabilities for each dog breed | |
prediction_dict = {class_names[i]: np.round(float(prediction[0][i]), 2) for i in range(len(class_names))} | |
sorted_predictions = dict(sorted(prediction_dict.items(), key=lambda item: item[1], reverse=True)) | |
return sorted_predictions | |
input_image = gr.Image() | |
iface = gr.Interface( | |
fn=predict_breed, | |
inputs=input_image, | |
outputs=gr.Label(), | |
description="A simple classification model for determining a dog breed.") | |
iface.launch(share=True) | |