File size: 1,180 Bytes
2774c5b e916ed1 2774c5b |
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 |
import io
import base64
import json
import cv2
import numpy as np
from ultralytics import YOLO
# Initialize your model
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%')
# Inference endpoint
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:
# must be in this format
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)
|