import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from PIL import Image
import numpy as np
import os

# Optional: Set YOLO config dir for write access (for HF Spaces)
os.environ["YOLO_CONFIG_DIR"] = "/tmp"

# Load YOLO model from Hugging Face
yolo_path = hf_hub_download(repo_id="RippedKek/leafmed-models", filename="yolo11x_leaf.pt")
yolo_model = YOLO(yolo_path)

def detect_and_crop_leaf(image: Image.Image):
    image_rgb = image.convert("RGB")
    image_rgb.save("temp_input.jpg")

    results = yolo_model.predict("temp_input.jpg", task="detect", conf=0.55)
    boxes = results[0].boxes
    class_ids = boxes.cls.cpu().numpy()
    xyxy = boxes.xyxy.cpu().numpy()
    class_names = [yolo_model.names[int(i)] for i in class_ids]

    for i, name in enumerate(class_names):
        if name.lower() == "leaf":
            x1, y1, x2, y2 = map(int, xyxy[i])
            x1_adj = max(int(x1 - x1 * 0.1), 0)
            y1_adj = max(int(y1 - y1 * 0.1), 0)
            x2_adj = int(x2 + x2 * 0.1)
            y2_adj = int(y2 + y2 * 0.1)

            cropped = image_rgb.crop((x1_adj, y1_adj, x2_adj, y2_adj))
            return cropped, f"Leaf detected with bounding box: ({x1_adj}, {y1_adj}, {x2_adj}, {y2_adj})"
    
    return None, "No leaf detected in the image."

demo = gr.Interface(
    fn=detect_and_crop_leaf,
    inputs=gr.Image(type="pil", label="Upload Leaf Image"),
    outputs=[
        gr.Image(type="pil", label="Cropped Leaf"),
        gr.Textbox(label="Detection Info")
    ],
    title="Leaf Detector",
    description="Upload an image containing a leaf. The model will detect and crop the leaf area."
)

demo.launch()