from typing import Dict, List, Any
from ultralytics import YOLO

class EndpointHandler():
    def __init__(self, path=""):
        # Preload all the elements you are going to need at inference.
        # pseudo:
        # self.model= load_model(path)
        # yolov8_model_name = 'yolov8_2023-07-19_yolov8m.pt'
        self.model = YOLO(path)    

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        data args:
            inputs (:obj: `str` | `PIL.Image` | `np.array`)
            kwargs
        Return:
            A :obj:`list` | `dict`: will be serialized and returned
        """
        # Get the prediction
        result = self.model(data['inputs'])
        # Get the original image with channel shifted
        img = result[0].orig_img[:,:,::-1]
        H, W, _ = img.shape
        annotated = img.copy()
        # Modify crop so that it is square
        try:
            x1, y1, x2, y2 = result[0].boxes.xyxy.numpy().astype('int')[0]
            if result[0].boxes.conf[0].item() < 0.75: # if low in confidence
                x1, y1, x2, y2 = 0, 0, W, H
            else:
                annotated = result[0].plot(labels=False, conf=False)[:,:,::-1]
        except: # in case there is no detection
            x1, y1, x2, y2 = 0, 0, W, H

        h, w = y2-y1, x2-x1
        offset = abs(h-w) // 2
        if h > w:
            x1 = max(x1 - offset, 0)
            x2 = min(x2 + offset, W)
        else:
            y1 = max(y1 - offset, 0)
            y2 = min(y2 + offset, H)
        new_image = img[y1:y2, x1:x2]
        # Return the annotated original image with the square cropped
        return annotated, new_image