# app.py

import gradio as gr
from transformers import pipeline
from PIL import Image

# -------------------------------
# 1. Load model (only once)
# -------------------------------
# Using a free Hugging Face model (no API key needed)
classifier = pipeline(
    "image-classification",
    model="nateraw/vit-base-beans"
)

# -------------------------------
# 2. Disease → Treatment mapping
# -------------------------------
treatment_suggestions = {
    "angular_leaf_spot": "Use fungicide sprays and avoid overhead watering.",
    "bean_rust": "Apply sulfur-based fungicides and remove infected leaves.",
    "healthy": "Your plant looks healthy. Maintain proper watering and sunlight.",
}

# -------------------------------
# 3. Prediction function
# -------------------------------
def predict_disease(image):
    if image is None:
        return "Please upload an image.", "", ""

    # Ensure image is in correct format
    image = Image.fromarray(image)

    # Run prediction
    results = classifier(image)

    # Get top prediction
    top_result = results[0]
    label = top_result["label"]
    confidence = round(top_result["score"] * 100, 2)

    # Get treatment suggestion
    treatment = treatment_suggestions.get(
        label.lower(),
        "No specific treatment found. Consult an agricultural expert."
    )

    return label, f"{confidence}%", treatment

# -------------------------------
# 4. Gradio UI
# -------------------------------
with gr.Blocks(title="🌿 Crop Disease Detection") as app:

    gr.Markdown("# 🌿 Crop Disease Detection from Leaf Images")
    gr.Markdown("Upload a leaf image to detect disease and get treatment suggestions.")

    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload Leaf Image")

    with gr.Row():
        predict_button = gr.Button("🔍 Predict")

    with gr.Row():
        label_output = gr.Textbox(label="Predicted Disease")
        confidence_output = gr.Textbox(label="Confidence Score")
        treatment_output = gr.Textbox(label="Suggested Treatment")

    # Button click action
    predict_button.click(
        fn=predict_disease,
        inputs=image_input,
        outputs=[label_output, confidence_output, treatment_output]
    )

# -------------------------------
# 5. Launch app
# -------------------------------
if _name_ == "_main_":
    app.launch()