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

# Load the trained model
MODEL_PATH = "setosys_dogs_model.h5"
model = tf.keras.models.load_model(MODEL_PATH)

# Get class labels from the model (assuming the model has a 'class_indices' attribute)
class_labels = list(model.class_indices.keys())  # Fetch class labels from the model

# Image preprocessing function using EfficientNetV2S
def preprocess_image(img: Image.Image) -> np.ndarray:
    """Preprocess the image to match the model's input requirements."""
    img = img.resize((224, 224))  # Resize image to model input size
    img_array = np.array(img)
    img_array = preprocess_input(img_array)  # EfficientNetV2 preprocessing
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    return img_array

# Prediction function
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)
    
    # Check the shape of the predictions to make sure the output is correct
    print("Predictions Shape:", predictions.shape)
    
    class_idx = np.argmax(predictions)  # Index of the highest prediction probability
    confidence = float(np.max(predictions))  # Confidence score
    
    # Get predicted breed and its confidence score
    predicted_breed = class_labels[class_idx] if class_idx < len(class_labels) else "Unknown"
    
    return {predicted_breed: confidence}

# Create Gradio interface
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."
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()