Upload 2 files

function.yaml

@@ -0,0 +1,64 @@
+metadata:
+  name: custom-model-yolov8
+  namespace: cvat
+  annotations:
+    name: custom-model-yolov8
+    type: detector
+    framework: pytorch
+    # change this accordingly to your model output/classes
+    spec: |
+      [
+        {"id": 0, "name": "tench"},
+        {"id": 1, "name": "goldfish"},
+        .
+        .
+        .
+        .
+        {"id": 998, "name": "ear"},
+        {"id": 999, "name": "toilet paper"}
+      ]
+spec:
+  description: custom-model-yolov8
+  runtime: 'python:3.9'
+  handler: main:handler
+  eventTimeout: 30s
+
+  build:
+    image: custom-model-yolov8
+    baseImage: ubuntu:22.04
+
+    directives:
+      preCopy:
+        - kind: ENV
+          value: DEBIAN_FRONTEND=noninteractive
+        - kind: RUN
+          value: apt-get update && apt-get -y install curl git python3 python3-pip 
+        - kind: RUN
+          value: apt-get -y install libgl1-mesa-glx libglib2.0-dev
+        - kind: WORKDIR
+          value: /opt/nuclio
+          #
+          # make sure that for the next step (at least) the ultralytics package version
+          # is compatible to that of the the ultralytics package used to train the custom model
+        - kind: RUN
+          value: pip3 install ultralytics==8.0.114 opencv-python==4.7.0.72 numpy==1.24.3
+          #
+        - kind: RUN
+          value: ln -s /usr/bin/pip3 /usr/local/bin/pip
+        - kind: RUN
+          value: ln -s /usr/bin/python3 /usr/local/bin/python
+
+  triggers:
+    myHttpTrigger:
+      maxWorkers: 1
+      kind: 'http'
+      workerAvailabilityTimeoutMilliseconds: 10000
+      attributes:
+        maxRequestBodySize: 33554432 # 32MB
+
+  platform:
+    attributes:
+      restartPolicy:
+        name: always
+        maximumRetryCount: 3
+      mountMode: volume

main.py

@@ -0,0 +1,47 @@
+
+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('custom-yolov8n.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)
+