File size: 468 Bytes
7dd7b93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | 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() |