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

# Download model from your model repo
model_path = hf_hub_download(
    repo_id="ProConceptTech/jambazisight",
    filename="best.pt"
)

model = YOLO(model_path)

def detect(image):

    image = Image.fromarray(image)
    results = model(image, conf=0.15)

    detections = []
    draw = ImageDraw.Draw(image)

    for r in results:
        for box in r.boxes:

            x1, y1, x2, y2 = box.xyxy[0].tolist()
            conf = float(box.conf)
            cls = int(box.cls)
            label = model.names[cls]

            detections.append({
                "class": label,
                "confidence": conf,
                "bbox": [x1, y1, x2, y2]
            })

            # draw bounding box
            draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
            draw.text((x1, y1), f"{label} {conf:.2f}", fill="red")

    return image, detections


demo = gr.Interface(
    fn=detect,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(label="Detected Image"),
        gr.JSON(label="Detections")
    ],
    title="JambaziSight Object Detection API",
    description="Upload an image to detect objects using the JambaziSight YOLO model."
)

demo.launch()