predict: add segmentation models to onnx/coreml and refactor openvino

plugins/coreml/src/coreml/__init__.py

@@ -16,6 +16,7 @@ from common import yolo
 from coreml.face_recognition import CoreMLFaceRecognition
 from coreml.custom_detection import CoreMLCustomDetection
 from coreml.clip_embedding import CoreMLClipEmbedding
+from coreml.segment import CoreMLSegmentation
 
 try:
     from coreml.text_recognition import CoreMLTextRecognition
@@ -105,6 +106,7 @@ class CoreMLPlugin(
         self.faceDevice = None
         self.textDevice = None
         self.clipDevice = None
+        self.segmentDevice = None
 
         if not self.forked:
             asyncio.ensure_future(self.prepareRecognitionModels(), loop=self.loop)
@@ -149,6 +151,18 @@ class CoreMLPlugin(
                     "name": "CoreML CLIP Embedding",
                 }
             )
+
+            await scrypted_sdk.deviceManager.onDeviceDiscovered(
+                {
+                    "nativeId": "segment",
+                    "type": scrypted_sdk.ScryptedDeviceType.Builtin.value,
+                    "interfaces": [
+                        scrypted_sdk.ScryptedInterface.ClusterForkInterface.value,
+                        scrypted_sdk.ScryptedInterface.ObjectDetection.value,
+                    ],
+                    "name": "CoreML Segmentation",
+                }
+            )
         except:
             pass
 
@@ -162,6 +176,9 @@ class CoreMLPlugin(
         elif nativeId == "clipembedding":
             self.clipDevice = self.clipDevice or CoreMLClipEmbedding(self, nativeId)
             return self.clipDevice
+        elif nativeId == "segment":
+            self.segmentDevice = self.segmentDevice or CoreMLSegmentation(self, nativeId)
+            return self.segmentDevice
         custom_model = self.custom_models.get(nativeId, None)
         if custom_model:
             return custom_model

plugins/coreml/src/coreml/segment.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import traceback
+
+import numpy as np
+
+import coremltools as ct
+from common import async_infer
+from common import yolov9_seg
+from predict.segment import Segmentation
+
+prepareExecutor, predictExecutor = async_infer.create_executors("Segment")
+
+
+class CoreMLSegmentation(Segmentation):
+    def __init__(self, plugin, nativeId: str):
+        super().__init__(plugin=plugin, nativeId=nativeId)
+
+    def loadModel(self, name):
+        model_path = self.plugin.downloadHuggingFaceModelLocalFallback(name)
+        modelFile = os.path.join(model_path, f"{name}.mlpackage")
+        model = ct.models.MLModel(modelFile)
+        return model
+
+    async def detect_once(self, input, settings, src_size, cvss):
+        def predict():
+            input_name = self.model.get_spec().description.input[0].name
+            out_dict = self.model.predict({input_name: input})
+
+            outputs = list(out_dict.values())
+            pred = outputs[0]
+            proto = outputs[1]
+            pred = yolov9_seg.non_max_suppression(pred, nm=32)
+
+            return self.process_segmentation_output(pred, proto)
+
+        try:
+            objs = await asyncio.get_event_loop().run_in_executor(
+                predictExecutor, lambda: predict()
+            )
+        except:
+            traceback.print_exc()
+            raise
+
+        ret = self.create_detection_result(objs, src_size, cvss)
+        return ret