from ultralytics import YOLO
import io, base64
from PIL import Image
import numpy as np

MODEL = YOLO("best.pt")

def inference(image_bytes: bytes) -> bytes:
    """Takes raw image bytes, returns annotated PNG bytes."""
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    results = MODEL(img)
    annotated = results[0].plot()  # returns a NumPy array
    buf = io.BytesIO()
    Image.fromarray(annotated).save(buf, format="PNG")
    return buf.getvalue()