import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model

# Image dimensions required by your model
img_height, img_width = 180, 180

# 1. Load the model
# We use compile=False to avoid errors with optimizers since we are only running predictions.
try:
    model_flower = load_model('model_flower.h5', compile=False)
except Exception as e:
    print(f"Error loading model: {e}")
    raise

class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']

def predict_image(img):
    if img is None:
        return None
        
    # 2. Resize the image to match the training size
    img_resized = tf.image.resize(img, (img_height, img_width))
    
    # 3. Add the batch dimension (1, 180, 180, 3)
    img_array = tf.expand_dims(img_resized, 0)
    
    # 4. Predict
    prediction = model_flower.predict(img_array)[0]
    
    # 5. Apply Softmax
    # Since your model was trained with from_logits=True, the output is raw scores.
    # We apply softmax to convert them into percentages (0.0 to 1.0).
    score = tf.nn.softmax(prediction)
    
    return {class_names[i]: float(score[i]) for i in range(len(class_names))}

# 6. Define the Interface
# Note: 'image_mode' is removed in Gradio 4.0+. It defaults to RGB automatically.
image = gr.Image(label="Upload Image")
label = gr.Label(num_top_classes=5)

gr.Interface(
    fn=predict_image, 
    inputs=image, 
    outputs=label,
    title="Flower Classification",
    description="Upload an image to classify it as a daisy, dandelion, rose, sunflower, or tulip."
).launch()