import torch
from ultralytics import YOLO
from PIL import Image
import io

# Load YOLOv8 model
model = YOLO("model.pt")

def predict(image_bytes):
    img = Image.open(io.BytesIO(image_bytes))
    results = model.predict(img)
    output = []
    for result in results:
        for i in range(len(result.boxes)):
            output.append({
                "bbox": result.boxes.xyxy[i].tolist(),
                "class": int(result.boxes.cls[i].item()),
                "confidence": float(result.boxes.conf[i].item())
            })
    return output