import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
import numpy as np

# Load the trained model (replace with correct path to your model)
model = tf.keras.models.load_model("setosys_dogs_model.h5")

# Get class labels dynamically from model (from class_indices)
class_labels = {v: k for k, v in model.class_indices.items()}

# Preprocessing function for EfficientNetV2 model
def preprocess_image(image_path):
    img = load_img(image_path, target_size=(224, 224))  # Resize image to 224x224
    img_array = img_to_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(image):
    img_array = preprocess_image(image)
    predictions = model.predict(img_array)
    class_idx = np.argmax(predictions)  # Get index of the class with the highest probability
    breed = class_labels[class_idx]  # Get label using the index
    confidence = predictions[0][class_idx]
    return breed, confidence

# Gradio interface definition
iface = gr.Interface(fn=predict_dog_breed, 
                     inputs=gr.inputs.Image(type="filepath"), 
                     outputs=["text", "number"], 
                     live=True)

# Launch the Gradio app
iface.launch(share=True)  # `share=True` gives you a public link