You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
import json
|
|
import base64
|
|
from PIL import Image
|
|
import io
|
|
import torch
|
|
|
|
def init_context(context):
|
|
context.logger.info("Init context... 0%")
|
|
|
|
# Read the DL model
|
|
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # or yolov5m, yolov5l, yolov5x, custom
|
|
context.user_data.model = model
|
|
|
|
context.logger.info("Init context...100%")
|
|
|
|
def handler(context, event):
|
|
context.logger.info("Run yolo-v5 model")
|
|
data = event.body
|
|
buf = io.BytesIO(base64.b64decode(data["image"]))
|
|
threshold = float(data.get("threshold", 0.5))
|
|
context.user_data.model.conf = threshold
|
|
image = Image.open(buf)
|
|
yolo_results_json = context.user_data.model(image).pandas().xyxy[0].to_dict(orient='records')
|
|
|
|
encoded_results = []
|
|
for result in yolo_results_json:
|
|
encoded_results.append({
|
|
'confidence': result['confidence'],
|
|
'label': result['name'],
|
|
'points': [
|
|
result['xmin'],
|
|
result['ymin'],
|
|
result['xmax'],
|
|
result['ymax']
|
|
],
|
|
'type': 'rectangle'
|
|
})
|
|
|
|
return context.Response(body=json.dumps(encoded_results), headers={},
|
|
content_type='application/json', status_code=200)
|