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

# Load the saved model (ensure the model file is in the same directory)
model = tf.keras.models.load_model('sugarcane_disease_model.h5')
# Define class names in the same order as used during training
class_names = ['BacterialBlights', 'Healthy', 'Mosaic', 'RedRot', 'Rust', 'Yellow']

def predict_image(img):
    """
    Preprocess the input image, run model inference,
    and return the predicted class label.
    """
    # Resize image to (256,256) - adjust if necessary
    img_resized = cv2.resize(img, (256, 256))
    # Normalize the image to [0, 1]
    img_normalized = img_resized.astype("float32") / 255.0
    # Expand dimensions to match model's input shape (batch size of 1)
    img_input = np.expand_dims(img_normalized, axis=0)
    
    # Run inference
    preds = model.predict(img_input)
    predicted_index = np.argmax(preds, axis=1)[0]
    predicted_class = class_names[predicted_index]
    
    return predicted_class

# Create Gradio Interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy"),
    outputs="text",
    title="Sugarcane Disease Classification",
    description="Upload a sugarcane leaf image to classify its disease or check if it's healthy."
)

if __name__ == "__main__":
    iface.launch()