| 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() |