|
|
|
import io |
|
import base64 |
|
import json |
|
|
|
import cv2 |
|
import numpy as np |
|
from ultralytics import YOLO |
|
|
|
|
|
def init_context(context): |
|
context.logger.info('Init context... 0%') |
|
model = YOLO('best.pt') |
|
context.user_data.model_handler = model |
|
context.logger.info('Init context...100%') |
|
|
|
|
|
def handler(context, event): |
|
context.logger.info('Run custom yolov8 model') |
|
data = event.body |
|
image_buffer = io.BytesIO(base64.b64decode(data['image'])) |
|
image = cv2.imdecode(np.frombuffer(image_buffer.getvalue(), np.uint8), cv2.IMREAD_COLOR) |
|
|
|
results = context.user_data.model_handler(image) |
|
result = results[0] |
|
|
|
boxes = result.boxes.data[:,:4] |
|
confs = result.boxes.conf |
|
clss = result.boxes.cls |
|
class_name = result.names |
|
|
|
detections = [] |
|
threshold = 0.1 |
|
for box, conf, cls in zip(boxes, confs, clss): |
|
label = class_name[int(cls)] |
|
if conf >= threshold: |
|
|
|
detections.append({ |
|
'confidence': str(float(conf)), |
|
'label': label, |
|
'points': box.tolist(), |
|
'type': 'rectangle', |
|
}) |
|
|
|
return context.Response(body=json.dumps(detections), headers={}, |
|
content_type='application/json', status_code=200) |
|
|
|
|
