Creating Inference.py file

inference.py

@@ -0,0 +1,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()