import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
import torch
import cv2
import os
from segment_anything import sam_model_registry, SamPredictor

# Avoid matplotlib permission errors
os.environ["MPLCONFIGDIR"] = "/tmp"

# Load classification model
model = tf.keras.models.load_model("3.keras")
CLASS_NAMES = ['Fungi', 'Healthy', 'Nematode', 'Pest', 'Phytopthora', 'Virus']

# Load SAM
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SAM_PATH = "./sam_vit_b.pth"
sam = sam_model_registry["vit_b"](checkpoint=SAM_PATH).to(DEVICE)
predictor = SamPredictor(sam)
print("✅ Models loaded")

# Inference function (no division by 255)
def classify_leaf(image: Image.Image):
    image_np = np.array(image.convert("RGB"))
    image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    # Segment with SAM
    predictor.set_image(image_bgr)
    H, W, _ = image_bgr.shape
    point = np.array([[W // 2, H // 2]])
    label = np.array([1])

    masks, scores, _ = predictor.predict(
        point_coords=point,
        point_labels=label,
        multimask_output=True
    )

    if len(masks) == 0:
        return "No leaf detected", {}, image

    # Choose the largest mask
    largest_mask = max(masks, key=lambda m: m.sum())
    segmented = image_np * largest_mask[:, :, None]

    seg_pil = Image.fromarray(segmented.astype("uint8")).convert("RGB").resize((224, 224))
    arr = np.expand_dims(np.array(seg_pil).astype("float32"), axis=0)  # ⚠️ No division by 255

    prediction = model.predict(arr)[0]
    predicted_class = CLASS_NAMES[int(np.argmax(prediction))]
    probs = {CLASS_NAMES[i]: float(round(prediction[i], 4)) for i in range(len(CLASS_NAMES))}

    return predicted_class, probs, seg_pil

# Gradio Interface
demo = gr.Interface(
    fn=classify_leaf,
    inputs=gr.Image(type="pil", label="Upload Potato Leaf"),
    outputs=[
        gr.Label(label="Predicted Disease"),
        gr.JSON(label="Class Probabilities"),
        gr.Image(label="Segmented Leaf"),
    ],
    title="🥔 Potato Leaf Disease Detection",
    description="Upload a potato leaf image. SAM will detect the largest leaf, and the model will classify its disease without normalizing pixel values.",
)

demo.launch()