Theingh02's picture
Update app.py
d5ce954 verified
import gradio as gr
import pandas as pd
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image
from sklearn.preprocessing import LabelEncoder
# Load the trained model
model = load_model("best_dog_breed_model.keras") # Ensure this file exists in the same directory
# Load and encode labels
df = pd.read_csv("labels.csv")
le = LabelEncoder()
le.fit(df['breed'])
breed_list = list(le.classes_)
# Define prediction function
def predict(image):
image = image.resize((224, 224))
img_array = np.expand_dims(np.array(image) / 255.0, axis=0)
preds = model.predict(img_array)[0]
top5_idx = preds.argsort()[-5:][::-1]
top5_preds = {breed_list[i]: float(preds[i]) for i in top5_idx}
# Threshold check
top_pred_label = breed_list[top5_idx[0]]
top_pred_confidence = preds[top5_idx[0]]
# Define unknown labels
unknown_labels = {'toy', 'human'}
if top_pred_label in unknown_labels or top_pred_confidence < 0.80:
final_result = "Unknown"
else:
final_result = f"{top_pred_label} ({top_pred_confidence * 100:.2f}%)"
return top5_preds, final_result
# Gradio Interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=[
gr.Label(num_top_classes=5, label="Top 5 Predictions"),
gr.Textbox(label="Final Result")
],
title="Dog Breed Classifier",
description="Upload a dog image. If it's not a known breed or the confidence is too low, it will return 'Unknown'."
)
iface.launch()