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

# Constants
IMAGE_SIZE = (300, 300)

# Class index to label mapping
CLASS_NAMES = {
    0: "A healthy tomato leaf",
    1: "A tomato leaf with Leaf Mold",
    2: "A tomato leaf with Target Spot",
    3: "A tomato leaf with Late Blight",
    4: "A tomato leaf with Early Blight",
    5: "A tomato leaf with Bacterial Spot",
    6: "A tomato leaf with Septoria Leaf Spot",
    7: "A tomato leaf with Tomato Mosaic Virus",
    8: "A tomato leaf with Tomato Yellow Leaf Curl Virus",
    9: "A tomato leaf with Spider Mites Two-spotted Spider Mite"
}

# Load the model from Hugging Face Hub
model_url = "https://huggingface.co/chimithecat/penyakit_tomat/resolve/main/Tomato_Models.h5"
model_path = tf.keras.utils.get_file("Tomato_Models.h5", model_url)
model = tf.keras.models.load_model(model_path)

# Prediction function
def classify(image):
    if image is None:
        return "Tidak ada gambar.", ""

    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    img = image.resize(IMAGE_SIZE)
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    predictions = model.predict(img_array)[0]
    predicted_index = int(np.argmax(predictions))
    confidence = float(predictions[predicted_index])
    label = CLASS_NAMES[predicted_index]

    return f"{label}", f"Confidence: {confidence*100:.2f}%"

# UI with Gradio Blocks
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        # 🍅 Tomato Leaf Disease Classifier  
        Upload a photo of a tomato leaf to detect its potential disease.
        """
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Leaf Image")
            submit_button = gr.Button("Analyze", variant="primary")

        with gr.Column():
            result_output = gr.Text(label="Prediction Result")
            confidence_output = gr.Text(label="Confidence")

    submit_button.click(
        fn=classify,
        inputs=image_input,
        outputs=[result_output, confidence_output],
        api_name="predict"
    )

app.launch(show_api=True)