File size: 1,205 Bytes
599338d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | 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
|