from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from ultralytics import YOLO

# Class mapping and colors
class_names = {0: "weed", 1: "wheat"}
class_colors = {
    "weed": "red",
    "wheat": "green"
}

# Load ONNX model
model = YOLO("best.onnx")

def detect(image):
    original_image = image.copy()
    results = model.predict(image, imgsz=640)[0]
    preds = results.boxes.data

    draw = ImageDraw.Draw(original_image)
    try:
        font = ImageFont.truetype("arial.ttf", 16)
    except:
        font = ImageFont.load_default()

    output = []
    for box in preds.cpu():
        x1, y1, x2, y2, conf, cls = box.tolist()
        if conf >= 0.5:
            class_id = int(cls)
            class_name = class_names.get(class_id, "unknown")
            color = class_colors.get(class_name, "blue")
            label = f"{class_name} ({conf:.2f})"
            draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
            draw.text((x1, y1 - 10), label, fill=color, font=font)
            output.append({
                "x1": int(x1), "y1": int(y1),
                "x2": int(x2), "y2": int(y2),
                "confidence": round(conf, 2),
                "class": class_name
            })

    return output  # Return only JSON for Android API usage

# Launch Gradio as API
gr.Interface(
    fn=detect,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(label="Detection Results"),
    title="YOLOv11M Weed Detection API",
    description="Send an image to detect weed (🔴) and wheat (🟢).",
    allow_flagging="never"
).launch()