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

# Load TFLite model
interpreter = tf.lite.Interpreter(model_path="plant_disease_model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Load labels
with open("class_labels.json") as f:
    labels = json.load(f)

def predict(img):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)
    img = img.resize((224, 224))  # adjust size to your model
    img = np.array(img, dtype=np.float32) / 255.0
    img = np.expand_dims(img, axis=0)

    interpreter.set_tensor(input_details[0]['index'], img)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])

    pred_class = int(np.argmax(output_data[0]))
    print("Input shape:", img.shape)
    print("Input dtype:", img.dtype)
    print("Output:", output_data)
    print("Predicted class:", pred_class)
    print("Label:", labels[str(pred_class)])
    return labels[str(pred_class)]
    

demo = gr.Interface(fn=predict, inputs="image", outputs="label")
demo.launch()