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