import gradio as gr
from ultralytics import YOLO
from PIL import Image
import torch

# Load a YOLO model (e.g., YOLOv8)
model = YOLO('yolov8n.pt')  # You can choose other YOLO models as well, like 'yolov8s.pt' or 'yolov8m.pt'

def yolo_inference(image):
    # Convert PIL image to a format suitable for the YOLO model
    results = model(image)
    
    # Draw bounding boxes on the image
    result_image = results[0].plot()
    
    # Convert result image from OpenCV format to PIL format
    result_image = Image.fromarray(result_image)

    # Get detection results
    labels, boxes = results[0].names, results[0].boxes
    
    return result_image, {labels[int(box.cls)]: float(box.conf) for box in boxes}


# Set up Gradio interface using the updated method
inputs = gr.Image(type="pil", label="Upload an Image")
outputs = [
    gr.Image(type="pil", label="Detection Output"),
    gr.JSON(label="Detection Results"),
]

gr.Interface(
    fn=yolo_inference,
    inputs=inputs,
    outputs=outputs,
    title="YOLO Object Detection",
    description="Upload an image and detect objects using YOLO pre-trained models.",
    allow_flagging="never"
).launch()