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


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


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

    def __call__(self, image, config: Optional[Dict[str, Any]] = None) -> 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
        """
        if config is None:
            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
        result = self.model.predict(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