Create handler.py

handler.py

@@ -0,0 +1,39 @@
+from typing import Dict, List, Any
+from ultralytics import YOLO
+from PIL import Image
+import torch
+import io
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Load the YOLO model from the current directory
+        # The 'path' argument is provided by Hugging Face automatically
+        self.model = YOLO(f"{path}/best.pt")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs (:obj: `bytes`) : The raw image bytes sent from Supabase
+        Return:
+            A :obj:`list` | `dict`: The detection results
+        """
+        # Get inputs
+        inputs = data.pop("inputs", data)
+
+        # Convert bytes to PIL Image
+        img = Image.open(io.BytesIO(inputs))
+
+        # Run inference
+        results = self.model(img, conf=0.25)
+
+        # Format results for your dashboard
+        payload = []
+        for r in results:
+            for box in r.boxes:
+                payload.append({
+                    "label": self.model.names[int(box.cls)],
+                    "confidence": float(box.conf),
+                    "points": box.xyxy[0].tolist(), # [x1, y1, x2, y2]
+                })
+
+        return payload