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()