|
from ultralyticsplus import YOLO |
|
from typing import Dict, Any, List |
|
|
|
|
|
DEFAULT_CONFIG = {'conf': 0.25, 'iou': 0.45, 'agnostic_nms': False, 'max_det': 1000} |
|
BOX_KEYS = ['xmin', 'ymin', 'xmax', 'ymax'] |
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
self.model = YOLO('ultralyticsplus/yolov8s') |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
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: |
|
A :obj: `dict` | `dict`: {scores, labels, boxes} |
|
""" |
|
inputs = data.pop("inputs", data) |
|
input_config = inputs.pop("config", DEFAULT_CONFIG) |
|
config = {**DEFAULT_CONFIG, **input_config} |
|
|
|
if config is None: |
|
config = DEFAULT_CONFIG |
|
|
|
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') |
|
|
|
|
|
names = self.model.model.names |
|
|
|
|
|
result = self.model.predict(inputs['image'])[0] |
|
prediction = [] |
|
for score, label, box in zip(result.boxes.conf, result.boxes.cls, result.boxes.xyxy): |
|
item_score = score.item() |
|
item_label = names[int(label)] |
|
item_box = box.to(dtype=int).tolist() |
|
|
|
item_prediction = { |
|
'score': item_score, |
|
'label': item_label, |
|
'box': dict(zip(BOX_KEYS, item_box)) |
|
} |
|
|
|
prediction.append(item_prediction) |
|
|
|
return prediction |
