import gradio as gr
from ultralytics import YOLO
import json

model = YOLO("best.pt")

def predict(img):
    results = model(img)
    annotated_img = results[0].plot() 
    
    detections = []
    # Safely extract boxes if they exist
    if len(results) > 0 and len(results[0].boxes) > 0:
        for box in results[0].boxes:
            # Use .item() to safely convert PyTorch tensors to standard Python numbers
            class_id = int(box.cls.item())
            label = model.names[class_id]
            conf = float(box.conf.item())
            
            detections.append({
                "label": label, 
                "confidence": conf
            })
        
    return annotated_img, json.dumps(detections)

demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=[
        gr.Image(type="pil"), 
        gr.Textbox(label="Detections")
    ]
)

demo.launch()