import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

model_path = "DogClassifier2.2.keras"
model = tf.keras.models.load_model(model_path)

# Define the core prediction function
def predict_bmwX(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)
    
    # Debug statement to check the shape of the prediction
    print(f"Prediction shape: {prediction.shape}")
    
    # 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']
    
    # Check if the number of predictions matches the number of class names
    if len(prediction[0]) != len(class_names):
        return f"Error: Number of model outputs ({len(prediction[0])}) does not match number of class names ({len(class_names)})."
    
    # 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))}
    
    # Sort the dictionary by value in descending order and get the top 3 classes
    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_bmwX,
    inputs=input_image, 
    outputs=gr.Label(),
    description="A simple MLP classification model for image classification using the MNIST dataset.")
iface.launch(share=True)