Create inference.py

inference.py

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