File size: 1,987 Bytes
d349524 5bb5008 d349524 5bb5008 d349524 3a43c10 d349524 5bb5008 d349524 5bb5008 d349524 5bb5008 04f67ba d349524 5bb5008 5007afa 5bb5008 04f67ba 5bb5008 40901b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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
# 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')
# Get label idx-to-name
names = self.model.model.names
# perform inference
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 |