import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications.mobilenet_v3 import preprocess_input

MODEL_PATH = "best_model_normal.keras"
CLASS_NAMES = [
    "Leaf_Algal",
    "Leaf_Blight",
    "Leaf_Colletotrichum",
    "Leaf_Healthy",
    "Leaf_Phomopsis",
    "Leaf_Rhizoctonia",
]
IMG_SIZE = (224, 224)

model = tf.keras.models.load_model(MODEL_PATH)

def predict(img):
    # img: numpy array HxWx3 (uint8)
    x = tf.image.resize(img, IMG_SIZE)
    x = tf.cast(x, tf.float32)
    x = preprocess_input(x)          # pipeline train
    x = tf.expand_dims(x, 0)         # (1,224,224,3)

    prob = model.predict(x, verbose=0)[0]
    return {CLASS_NAMES[i]: float(prob[i]) for i in range(len(CLASS_NAMES))}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload A Durian Leaf"),
    outputs=gr.Label(num_top_classes=3, label="Result"),
    title="Durian Leaf Disease Classification",
)
demo.launch()