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

# Load the YOLOv9 model
model = YOLO('best.pt')  # Make sure best.pt is in the same folder

def detect(image):
    results = model(image)
    annotated_frame = results[0].plot()  # Draw predictions on the image
    
    # Get detection results
    detections = []
    for r in results:
        boxes = r.boxes
        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0].tolist()  # Get box coordinates
            conf = float(box.conf[0])  # Get confidence
            cls = int(box.cls[0])  # Get class
            class_name = model.names[cls]  # Get class name
            detections.append({
                'class': cls,
                'class_name': class_name,
                'confidence': conf,
                'box': [x1, y1, x2, y2]
            })
    
    return Image.fromarray(annotated_frame), detections

# Launch the Gradio interface
gr.Interface(
    fn=detect,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Detected Image"),
        gr.JSON(label="Detection Results")
    ],
    title="YOLOv9 (Ultralytics) Object Detection",
    description="Upload an image to run object detection using your custom YOLOv9 model. The results will show both the annotated image and the detection details including class names.",
).launch()