from ultralyticsplus import YOLO
from typing import List, Dict, Any
from sahi import ObjectPrediction


DEFAULT_CONFIG = {'conf': 0.25, 'iou': 0.45, 'agnostic_nms': False, 'max_det': 1000}


class EndpointHandler():
    def __init__(self, path=""):
        self.model = YOLO('ultralyticsplus/yolov8s')

    def __call__(self, data: str) -> List[ObjectPrediction]:
        """
       data args:
            image: image path to segment
            config: (conf - NMS confidence threshold,
                     iou - NMS IoU threshold,
                     agnostic_nms - NMS class-agnostic: True / False,
                     max_det - maximum number of detections per image)
      Return:
            object_predictions
        """
        config = DEFAULT_CONFIG
        # Set model parameters
        self.model.overrides['conf'] = config.get('conf')
        self.model.overrides['iou'] = config.get('iou')
        self.model.overrides['agnostic_nms'] = config.get('agnostic_nms')
        self.model.overrides['max_det'] = config.get('max_det')   
        
        # perform inference
        inputs = data.pop("inputs", data)
        result = self.model.predict(inputs['image'])[0]

        names = self.model.model.names
        boxes = result.boxes

        object_predictions = []
        if boxes is not None:
            det_ind = 0
            for xyxy, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls):
                object_prediction = ObjectPrediction(
                    bbox=xyxy.tolist(),
                    category_name=names[int(cls)],
                    category_id=int(cls),
                    score=conf,
                )
                object_predictions.append(object_prediction)
                det_ind += 1
        return object_predictions