From 182d941fadc7bb1001336dbca08b89c51b281bbb Mon Sep 17 00:00:00 2001 From: wartek69 Date: Mon, 17 Jan 2022 20:56:58 +0100 Subject: [PATCH] Added YOLOv5 serverless function for auto labeling (#4178) * added yolov5 serverless function for auto labeling * updated CHANGELOG.md Co-authored-by: Alex --- CHANGELOG.md | 2 +- README.md | 1 + .../ultralytics/yolov5/nuclio/function.yaml | 121 ++++++++++++++++++ .../pytorch/ultralytics/yolov5/nuclio/main.py | 40 ++++++ 4 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 serverless/pytorch/ultralytics/yolov5/nuclio/function.yaml create mode 100644 serverless/pytorch/ultralytics/yolov5/nuclio/main.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d7c688ec703..c90b3611ec0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,7 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added information about OpenVINO toolkit to login page () - Support for working with ellipses () - Add several flags to task creation CLI () - +- Add YOLOv5 serverless function for automatic annotation () ### Changed - Users don't have access to a task object anymore if they are assigneed only on some jobs of the task () - Different resources (tasks, projects) are not visible anymore for all CVAT instance users by default () diff --git a/README.md b/README.md index 3a9e092ecf16..ea369cae13ee 100644 --- a/README.md +++ b/README.md @@ -90,6 +90,7 @@ For more information about supported formats look at the | [Object reidentification](/serverless/openvino/omz/intel/person-reidentification-retail-300/nuclio) | reid | OpenVINO | X | | | [Semantic segmentation for ADAS](/serverless/openvino/omz/intel/semantic-segmentation-adas-0001/nuclio) | detector | OpenVINO | X | | | [Text detection v4](/serverless/openvino/omz/intel/text-detection-0004/nuclio) | detector | OpenVINO | X | | +| [YOLO v5](/serverless/pytorch/ultralytics/yolov5/nuclio) | detector | PyTorch | X | | | [SiamMask](/serverless/pytorch/foolwood/siammask/nuclio) | tracker | PyTorch | X | X | | [f-BRS](/serverless/pytorch/saic-vul/fbrs/nuclio) | interactor | PyTorch | X | | | [HRNet](/serverless/pytorch/saic-vul/hrnet/nuclio) | interactor | PyTorch | | X | diff --git a/serverless/pytorch/ultralytics/yolov5/nuclio/function.yaml b/serverless/pytorch/ultralytics/yolov5/nuclio/function.yaml new file mode 100644 index 000000000000..8fcb743598ec --- /dev/null +++ b/serverless/pytorch/ultralytics/yolov5/nuclio/function.yaml @@ -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 diff --git a/serverless/pytorch/ultralytics/yolov5/nuclio/main.py b/serverless/pytorch/ultralytics/yolov5/nuclio/main.py new file mode 100644 index 000000000000..92bcf2e02a32 --- /dev/null +++ b/serverless/pytorch/ultralytics/yolov5/nuclio/main.py @@ -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)