from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import gradio as gr
from PIL import Image, ImageDraw
import requests
from io import BytesIO

# Download custom pre-trained model from Hugging Face Hub
model_path = hf_hub_download(
    repo_id="piky/yolo11",
    filename="yolo11m.onnx"
)

# Load ONNX model
model = YOLO(model_path, task="detect")

def load_image(uploaded_image, image_url):
    if image_url and image_url.strip():
        response = requests.get(image_url.strip(), timeout=10)
        response.raise_for_status()
        return Image.open(BytesIO(response.content)).convert("RGB")
    elif uploaded_image is not None:
        return uploaded_image
    else:
        raise ValueError("Please upload an image or provide an image URL.")

def detect_pill(uploaded_image, image_url):
    image = load_image(uploaded_image, image_url)

    results = model.predict(image, conf=0.25)
    result = results[0]

    img = image.copy()
    draw = ImageDraw.Draw(img)

    detections = []

    for box in result.boxes:
        x1, y1, x2, y2 = box.xyxy[0].tolist()
        conf = float(box.conf[0])

        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
        draw.text((x1, max(0, y1 - 10)), f"pill {conf:.2f}", fill="red")

        detections.append(conf)

    if detections:
        summary = f"Detected {len(detections)} pill(s), max confidence: {max(detections):.2f}"
    else:
        summary = "No pill detected"

    return img, summary

demo = gr.Interface(
    fn=detect_pill,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(label="Or Enter Image URL")
    ],
    outputs=[
        gr.Image(type="pil", label="Detection Result"),
        gr.Textbox(label="Summary")
    ],
    title="YOLO11 Pill Detector",
    description="Upload an image or provide an image URL to detect pills using a custom YOLO11 ONNX model."
)

demo.launch()