import gradio as gr
from ultralytics import YOLO
import numpy as np

# Load YOLOv8 model
model = YOLO("best.pt")  # Make sure best.pt is in the repo root

def predict(image):
    # image is numpy array from Gradio
    results = model(image)[0]  # YOLO prediction

    annotated_image = results.plot()  # returns numpy array
    detections = []

    if results.boxes is not None:
        for box, cls, conf in zip(results.boxes.xyxy, results.boxes.cls, results.boxes.conf):
            detections.append({
                "label": model.names[int(cls)],
                "confidence": float(conf),
                "box": [float(coord) for coord in box]
            })

    return annotated_image, detections

# Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=[gr.Image(type="numpy"), gr.JSON()],
    title="Helmet Detection YOLOv8",
    description="Upload an image and detect helmet / head using YOLOv8"
)

demo.launch()