import gradio as gr
import cv2
import numpy as np
import tensorflow as tf

from PIL import Image

# Assuming you have already defined img_height, img_width, and class_names
# img_height, img_width = 180, 180
class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']

# Load the fine-tuned model (from local)
resnet_model = tf.keras.models.load_model('./flower_image_classification_ResNet50_v1.0.h5')

def preprocess_image(image):
    # Convert the PIL image to an array
    image = np.array(image)
    img_height = 180
    img_width = 180
    # Read and resize the image
    image_resized = cv2.resize(image, (img_height, img_width))
    
    # Preprocess the image
    image = np.expand_dims(image_resized, axis=0)
    
    # Predict with the model
    pred = resnet_model.predict(image)
    
    # Get the predicted class label
    predicted_class = np.argmax(pred)
    output_class = class_names[predicted_class]
    
    # Get the confidence level (probability)
    confidence_level = pred[0][predicted_class]
    
    return image_resized, output_class, confidence_level

def predict(image):
    image_resized, output_class, confidence_level = preprocess_image(image)
    return Image.fromarray(image_resized), output_class, str(confidence_level)

# Define the Gradio interface
inputs = gr.Image(type="pil", label="Upload Image")
outputs = [
    gr.Image(type="pil", label="Resized Image"),
    gr.Textbox(label="Predicted Class"),
    gr.Textbox(label="Confidence Level")
]

# Create the Gradio Interface
gr.Interface(
    fn=predict,
    inputs=inputs,
    outputs=outputs,
    title="Flower Classification with ResNet50",
    description="Upload an image of a flower to classify it into one of the five categories (Roses / Dandelion / Tulips / Sunflower / Daisy).",
    live=True
).launch()