coreml: yolo-nas

plugins/coreml/package.json

@@ -42,5 +42,5 @@
    "devDependencies": {
       "@scrypted/sdk": "file:../../sdk"
    },
-   "version": "0.1.49"
+   "version": "0.1.50"
 }

plugins/coreml/src/coreml/__init__.py

@@ -6,6 +6,7 @@
 import os
 import re
 from typing import Any, List, Tuple
+import numpy as np
 
 import coremltools as ct
 import scrypted_sdk
@@ -26,6 +27,7 @@
 
 availableModels = [
     "Default",
+    "scrypted_yolo_nas_s_320",
     "scrypted_yolov9c_320",
     "scrypted_yolov9c",
     "scrypted_yolov6n_320",
@@ -77,6 +79,7 @@ def __init__(self, nativeId: str | None = None):
                 self.storage.setItem("model", "Default")
             model = "scrypted_yolov9c_320"
         self.yolo = "yolo" in model
+        self.scrypted_yolo_nas = "scrypted_yolo_nas" in model
         self.scrypted_yolo = "scrypted_yolo" in model
         self.scrypted_model = "scrypted" in model
         model_version = "v7"
@@ -211,7 +214,12 @@ async def detect_once(self, input: Image.Image, settings: Any, src_size, cvss):
         if self.yolo:
             out_dict = await self.queue_batch({self.input_name: input})
 
-            if self.scrypted_yolo:
+            if self.scrypted_yolo_nas:
+                predictions  = list(out_dict.values())
+                objs = yolo.parse_yolo_nas(predictions)
+                ret = self.create_detection_result(objs, src_size, cvss)
+                return ret
+            elif self.scrypted_yolo:
                 results = list(out_dict.values())[0][0]
                 objs = yolo.parse_yolov9(results)
                 ret = self.create_detection_result(objs, src_size, cvss)

plugins/openvino/src/common/yolo.py

@@ -6,6 +6,20 @@
 
 defaultThreshold = .2
 
+def parse_yolo_nas(predictions):
+    objs = []
+    for pred_scores, pred_bboxes in zip(*predictions):
+        i, j = np.nonzero(pred_scores > .5)
+        pred_bboxes = pred_bboxes[i]
+        pred_cls_conf = pred_scores[i, j]
+        pred_cls_label = j[:]
+        for box, conf, label in zip(pred_bboxes, pred_cls_conf, pred_cls_label):
+            obj = Prediction(
+                int(label), conf.astype(float), Rectangle(box[0].astype(float), box[1].astype(float), box[2].astype(float), box[3].astype(float))
+            )
+            objs.append(obj)
+    return objs
+
 def parse_yolov9(results, threshold = defaultThreshold, scale = None, confidence_scale  = None):
     objs = []
     keep = np.argwhere(results[4:] > threshold)

plugins/tensorflow-lite/src/detect/__init__.py

@@ -73,7 +73,6 @@ async def detectObjects(self, mediaObject: MediaObject, session: ObjectDetection
         if mediaObject.mimeType == ScryptedMimeTypes.Image.value:
             image = await scrypted_sdk.sdk.connectRPCObject(mediaObject)
         else:
-            print('non image provided')
             image = await scrypted_sdk.mediaManager.convertMediaObjectToBuffer(mediaObject, ScryptedMimeTypes.Image.value)
 
         return await self.run_detection_image(image, session)