import gradio as gr
import numpy as np
import tensorflow as tf
from PIL import Image

# =========================
# LOAD TFLITE MODEL
# =========================
interpreter = tf.lite.Interpreter(model_path="potato_model.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# CHANGE THIS TO YOUR CLASSES
class_names = ['Potato_Early_Blight', 'Potato_Healthy', 'Potato_Late_Blight']

IMG_SIZE = 224

# =========================
# PREDICT FUNCTION
# =========================
def predict(image):
    image = image.convert("RGB")
    image = image.resize((IMG_SIZE, IMG_SIZE))

    img = np.array(image, dtype=np.float32) / 255.0
    img = np.expand_dims(img, axis=0)

    interpreter.set_tensor(input_details[0]['index'], img)
    interpreter.invoke()

    output = interpreter.get_tensor(output_details[0]['index'])[0]

    return {class_names[i]: float(output[i]) for i in range(len(class_names))}

# =========================
# GRADIO UI
# =========================
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Potato Disease Detection (TFLite)",
    description="Fast lightweight inference using TensorFlow Lite"
)

demo.launch()