| from ultralyticsplus import YOLO | |
| from typing import List, Dict, Any | |
| from sahi import ObjectPrediction | |
| import torch, torchvision | |
| 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 | |
| return torch.__version__, torchvision.__version__ | |