Added YOLOv5 serverless function for auto labeling (#4178)
* added yolov5 serverless function for auto labeling * updated CHANGELOG.md Co-authored-by: Alex <aleksander.chernyavskiy@seafar.eu>main
wartek69
4 years ago
committed by
GitHub
parent
f97db1c1ed
commit
182d941fad
@ -0,0 +1,121 @@
|
||||
metadata:
|
||||
name: ultralytics-yolov5
|
||||
namespace: cvat
|
||||
annotations:
|
||||
name: YOLO v5
|
||||
type: detector
|
||||
framework: pytorch
|
||||
spec: |
|
||||
[
|
||||
{ "id": 0, "name": "person" },
|
||||
{ "id": 1, "name": "bicycle" },
|
||||
{ "id": 2, "name": "car" },
|
||||
{ "id": 3, "name": "motorbike" },
|
||||
{ "id": 4, "name": "aeroplane" },
|
||||
{ "id": 5, "name": "bus" },
|
||||
{ "id": 6, "name": "train" },
|
||||
{ "id": 7, "name": "truck" },
|
||||
{ "id": 8, "name": "boat" },
|
||||
{ "id": 9, "name": "traffic light" },
|
||||
{ "id": 10, "name": "fire hydrant" },
|
||||
{ "id": 11, "name": "stop sign" },
|
||||
{ "id": 12, "name": "parking meter" },
|
||||
{ "id": 13, "name": "bench" },
|
||||
{ "id": 14, "name": "bird" },
|
||||
{ "id": 15, "name": "cat" },
|
||||
{ "id": 16, "name": "dog" },
|
||||
{ "id": 17, "name": "horse" },
|
||||
{ "id": 18, "name": "sheep" },
|
||||
{ "id": 19, "name": "cow" },
|
||||
{ "id": 20, "name": "elephant" },
|
||||
{ "id": 21, "name": "bear" },
|
||||
{ "id": 22, "name": "zebra" },
|
||||
{ "id": 23, "name": "giraffe" },
|
||||
{ "id": 24, "name": "backpack" },
|
||||
{ "id": 25, "name": "umbrella" },
|
||||
{ "id": 26, "name": "handbag" },
|
||||
{ "id": 27, "name": "tie" },
|
||||
{ "id": 28, "name": "suitcase" },
|
||||
{ "id": 29, "name": "frisbee" },
|
||||
{ "id": 30, "name": "skis" },
|
||||
{ "id": 31, "name": "snowboard" },
|
||||
{ "id": 32, "name": "sports ball" },
|
||||
{ "id": 33, "name": "kite" },
|
||||
{ "id": 34, "name": "baseball bat" },
|
||||
{ "id": 35, "name": "baseball glove" },
|
||||
{ "id": 36, "name": "skateboard" },
|
||||
{ "id": 37, "name": "surfboard" },
|
||||
{ "id": 38, "name": "tennis racket" },
|
||||
{ "id": 39, "name": "bottle" },
|
||||
{ "id": 40, "name": "wine glass" },
|
||||
{ "id": 41, "name": "cup" },
|
||||
{ "id": 42, "name": "fork" },
|
||||
{ "id": 43, "name": "knife" },
|
||||
{ "id": 44, "name": "spoon" },
|
||||
{ "id": 45, "name": "bowl" },
|
||||
{ "id": 46, "name": "banana" },
|
||||
{ "id": 47, "name": "apple" },
|
||||
{ "id": 48, "name": "sandwich" },
|
||||
{ "id": 49, "name": "orange" },
|
||||
{ "id": 50, "name": "broccoli" },
|
||||
{ "id": 51, "name": "carrot" },
|
||||
{ "id": 52, "name": "hot dog" },
|
||||
{ "id": 53, "name": "pizza" },
|
||||
{ "id": 54, "name": "donut" },
|
||||
{ "id": 55, "name": "cake" },
|
||||
{ "id": 56, "name": "chair" },
|
||||
{ "id": 57, "name": "sofa" },
|
||||
{ "id": 58, "name": "pottedplant" },
|
||||
{ "id": 59, "name": "bed" },
|
||||
{ "id": 60, "name": "diningtable" },
|
||||
{ "id": 61, "name": "toilet" },
|
||||
{ "id": 62, "name": "tvmonitor" },
|
||||
{ "id": 63, "name": "laptop" },
|
||||
{ "id": 64, "name": "mouse" },
|
||||
{ "id": 65, "name": "remote" },
|
||||
{ "id": 66, "name": "keyboard" },
|
||||
{ "id": 67, "name": "cell phone" },
|
||||
{ "id": 68, "name": "microwave" },
|
||||
{ "id": 69, "name": "oven" },
|
||||
{ "id": 70, "name": "toaster" },
|
||||
{ "id": 71, "name": "sink" },
|
||||
{ "id": 72, "name": "refrigerator" },
|
||||
{ "id": 73, "name": "book" },
|
||||
{ "id": 74, "name": "clock" },
|
||||
{ "id": 75, "name": "vase" },
|
||||
{ "id": 76, "name": "scissors" },
|
||||
{ "id": 77, "name": "teddy bear" },
|
||||
{ "id": 78, "name": "hair drier" },
|
||||
{ "id": 79, "name": "toothbrush" }
|
||||
]
|
||||
|
||||
spec:
|
||||
description: YOLO v5 via pytorch hub
|
||||
runtime: 'python:3.6'
|
||||
handler: main:handler
|
||||
eventTimeout: 30s
|
||||
build:
|
||||
image: cvat/ultralytics-yolov5
|
||||
baseImage: ultralytics/yolov5:latest
|
||||
|
||||
directives:
|
||||
preCopy:
|
||||
- kind: USER
|
||||
value: root
|
||||
- kind: WORKDIR
|
||||
value: /opt/nuclio
|
||||
|
||||
triggers:
|
||||
myHttpTrigger:
|
||||
maxWorkers: 2
|
||||
kind: 'http'
|
||||
workerAvailabilityTimeoutMilliseconds: 10000
|
||||
attributes:
|
||||
maxRequestBodySize: 33554432 # 32MB
|
||||
|
||||
platform:
|
||||
attributes:
|
||||
restartPolicy:
|
||||
name: always
|
||||
maximumRetryCount: 3
|
||||
mountMode: volume
|
||||
@ -0,0 +1,40 @@
|
||||
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)
|
||||
Loading…
Reference in New Issue