predict plugins: refactor recog, add onnx, fix spurious model leaks

plugins/coreml/package.json

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

plugins/coreml/src/coreml/__init__.py

@@ -134,6 +134,8 @@ def __init__(self, nativeId: str | None = None):
         self.loop = asyncio.get_event_loop()
         self.minThreshold = 0.2
 
+        self.faceDevice = None
+        self.textDevice = None
         asyncio.ensure_future(self.prepareRecognitionModels(), loop=self.loop)
 
     async def prepareRecognitionModels(self):
@@ -171,9 +173,11 @@ async def prepareRecognitionModels(self):
 
     async def getDevice(self, nativeId: str) -> Any:
         if nativeId == "facerecognition":
-            return CoreMLFaceRecognition(nativeId)
+            self.faceDevice = self.faceDevice or CoreMLFaceRecognition(nativeId)
+            return self.faceDevice
         if nativeId == "textrecognition":
-            return CoreMLTextRecognition(nativeId)
+            self.textDevice = self.textDevice or CoreMLTextRecognition(nativeId)
+            return self.textDevice
         raise Exception("unknown device")
 
     async def getSettings(self) -> list[Setting]:

plugins/coreml/src/coreml/face_recognition.py

@@ -3,13 +3,15 @@
 import concurrent.futures
 import os
 
+import asyncio
 import coremltools as ct
 import numpy as np
 # import Quartz
 # from Foundation import NSData, NSMakeSize
 
 # import Vision
 from predict.face_recognize import FaceRecognizeDetection
+from PIL import Image
 
 
 def euclidean_distance(arr1, arr2):
@@ -29,6 +31,8 @@ def cosine_similarity(vector_a, vector_b):
 class CoreMLFaceRecognition(FaceRecognizeDetection):
     def __init__(self, nativeId: str | None = None):
         super().__init__(nativeId=nativeId)
+        self.detectExecutor = concurrent.futures.ThreadPoolExecutor(1, "detect-face")
+        self.recogExecutor = concurrent.futures.ThreadPoolExecutor(1, "recog-face")
 
     def downloadModel(self, model: str):
         model_version = "v7"
@@ -51,23 +55,29 @@ def downloadModel(self, model: str):
         inputName = model.get_spec().description.input[0].name
         return model, inputName
 
-    def predictDetectModel(self, input):
-        model, inputName = self.detectModel
-        out_dict = model.predict({inputName: input})
-        results = list(out_dict.values())[0][0]
+    async def predictDetectModel(self, input: Image.Image):
+        def predict():
+            model, inputName = self.detectModel
+            out_dict = model.predict({inputName: input})
+            results = list(out_dict.values())[0][0]
+            return results
+
+        results = await asyncio.get_event_loop().run_in_executor(
+            self.detectExecutor, lambda: predict()
+        )
         return results
 
-    def predictFaceModel(self, input):
-        model, inputName = self.faceModel
-        out_dict = model.predict({inputName: input})
-        return out_dict["var_2167"][0]
+    async def predictFaceModel(self, input: np.ndarray):
+        def predict():
+            model, inputName = self.faceModel
+            out_dict = model.predict({inputName: input})
+            results = out_dict["var_2167"][0]
+            return results
+        results = await asyncio.get_event_loop().run_in_executor(
+            self.recogExecutor, lambda: predict()
+        )
+        return results
 
-    def predictTextModel(self, input):
-        model, inputName = self.textModel
-        out_dict = model.predict({inputName: input})
-        preds = out_dict["linear_2"]
-        return preds
-
     # def predictVision(self, input: Image.Image) -> asyncio.Future[list[Prediction]]:
     #     buffer = input.tobytes()
     #     myData = NSData.alloc().initWithBytes_length_(buffer, len(buffer))

plugins/coreml/src/coreml/text_recognition.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
 
+import concurrent.futures
 import os
 
+import asyncio
+
 import coremltools as ct
+import numpy as np
+from PIL import Image
 
 from predict.text_recognize import TextRecognition
 
@@ -11,6 +16,9 @@ class CoreMLTextRecognition(TextRecognition):
     def __init__(self, nativeId: str | None = None):
         super().__init__(nativeId=nativeId)
 
+        self.detectExecutor = concurrent.futures.ThreadPoolExecutor(1, "detect-text")
+        self.recogExecutor = concurrent.futures.ThreadPoolExecutor(1, "recog-text")
+
     def downloadModel(self, model: str):
         model_version = "v7"
         mlmodel = "model"
@@ -32,14 +40,24 @@ def downloadModel(self, model: str):
         inputName = model.get_spec().description.input[0].name
         return model, inputName
 
-    def predictDetectModel(self, input):
-        model, inputName = self.detectModel
-        out_dict = model.predict({inputName: input})
-        results = list(out_dict.values())[0]
+    async def predictDetectModel(self, input: Image.Image):
+        def predict():
+            model, inputName = self.detectModel
+            out_dict = model.predict({inputName: input})
+            results = list(out_dict.values())[0]
+            return results
+        results = await asyncio.get_event_loop().run_in_executor(
+            self.detectExecutor, lambda: predict()
+        )
         return results
 
-    def predictTextModel(self, input):
-        model, inputName = self.textModel
-        out_dict = model.predict({inputName: input})
-        preds = out_dict["linear_2"]
+    async def predictTextModel(self, input: np.ndarray):
+        def predict():
+            model, inputName = self.textModel
+            out_dict = model.predict({inputName: input})
+            preds = out_dict["linear_2"]
+            return preds
+        preds = await asyncio.get_event_loop().run_in_executor(
+            self.recogExecutor, lambda: predict()
+        )
         return preds

plugins/onnx/package.json

@@ -33,6 +33,7 @@
       "runtime": "python",
       "type": "API",
       "interfaces": [
+         "DeviceProvider",
          "Settings",
          "ObjectDetection",
          "ObjectDetectionPreview"
@@ -41,5 +42,5 @@
    "devDependencies": {
       "@scrypted/sdk": "file:../../sdk"
    },
-   "version": "0.1.87"
+   "version": "0.1.88"
 }

plugins/onnx/src/ort/__init__.py

@@ -20,6 +20,13 @@
 import common.yolo as yolo
 from predict import PredictPlugin
 
+from .face_recognition import ONNXFaceRecognition
+
+try:
+    from .text_recognition import ONNXTextRecognition
+except:
+    ONNXTextRecognition = None
+
 availableModels = [
     "Default",
     "scrypted_yolo_nas_s_320",
@@ -72,6 +79,7 @@ def __init__(self, nativeId: str | None = None):
         deviceIds = json.loads(deviceIds)
         if not len(deviceIds):
             deviceIds = ["0"]
+        self.deviceIds = deviceIds
 
         compiled_models = []
         self.compiled_models = {}
@@ -124,6 +132,52 @@ def executor_initializer():
             thread_name_prefix="onnx-prepare",
         )
 
+        self.faceDevice = None
+        self.textDevice = None
+        asyncio.ensure_future(self.prepareRecognitionModels(), loop=self.loop)
+
+    async def prepareRecognitionModels(self):
+        try:
+            devices = [
+                {
+                    "nativeId": "facerecognition",
+                    "type": scrypted_sdk.ScryptedDeviceType.Builtin.value,
+                    "interfaces": [
+                        scrypted_sdk.ScryptedInterface.ObjectDetection.value,
+                    ],
+                    "name": "ONNX Face Recognition",
+                },
+            ]
+
+            if ONNXTextRecognition:
+                devices.append(
+                    {
+                        "nativeId": "textrecognition",
+                        "type": scrypted_sdk.ScryptedDeviceType.Builtin.value,
+                        "interfaces": [
+                            scrypted_sdk.ScryptedInterface.ObjectDetection.value,
+                        ],
+                        "name": "ONNX Text Recognition",
+                    },
+                )
+
+            await scrypted_sdk.deviceManager.onDevicesChanged(
+                {
+                    "devices": devices,
+                }
+            )
+        except:
+            pass
+
+    async def getDevice(self, nativeId: str) -> Any:
+        if nativeId == "facerecognition":
+            self.faceDevice = self.faceDevice or ONNXFaceRecognition(self, nativeId)
+            return self.faceDevice
+        elif nativeId == "textrecognition":
+            self.textDevice = self.textDevice or ONNXTextRecognition(self, nativeId)
+            return self.textDevice
+        raise Exception("unknown device")
+
     async def getSettings(self) -> list[Setting]:
         model = self.storage.getItem("model") or "Default"
         deviceIds = self.storage.getItem("deviceIds") or '["0"]'

plugins/onnx/src/ort/face_recognition.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import platform
+import sys
+import threading
+
+import numpy as np
+import onnxruntime
+from PIL import Image
+
+from predict.face_recognize import FaceRecognizeDetection
+
+
+class ONNXFaceRecognition(FaceRecognizeDetection):
+    def __init__(self, plugin, nativeId: str | None = None):
+        self.plugin = plugin
+
+        super().__init__(nativeId=nativeId)
+
+    def downloadModel(self, model: str):
+        onnxmodel = "best" if "scrypted" in model else model
+        model_version = "v1"
+        onnxfile = self.downloadFile(
+            f"https://raw.githubusercontent.com/koush/onnx-models/main/{model}/{onnxmodel}.onnx",
+            f"{model_version}/{model}/{onnxmodel}.onnx",
+        )
+        print(onnxfile)
+
+        compiled_models_array = []
+        compiled_models = {}
+        deviceIds = self.plugin.deviceIds
+
+        for deviceId in deviceIds:
+            sess_options = onnxruntime.SessionOptions()
+
+            providers: list[str] = []
+            if sys.platform == "darwin":
+                providers.append("CoreMLExecutionProvider")
+
+            if "linux" in sys.platform and platform.machine() == "x86_64":
+                deviceId = int(deviceId)
+                providers.append(("CUDAExecutionProvider", {"device_id": deviceId}))
+
+            providers.append("CPUExecutionProvider")
+
+            compiled_model = onnxruntime.InferenceSession(
+                onnxfile, sess_options=sess_options, providers=providers
+            )
+            compiled_models_array.append(compiled_model)
+
+            input = compiled_model.get_inputs()[0]
+            input_name = input.name
+
+        def executor_initializer():
+            thread_name = threading.current_thread().name
+            interpreter = compiled_models_array.pop()
+            compiled_models[thread_name] = interpreter
+            print("Runtime initialized on thread {}".format(thread_name))
+
+        executor = concurrent.futures.ThreadPoolExecutor(
+            initializer=executor_initializer,
+            max_workers=len(compiled_models_array),
+            thread_name_prefix="face",
+        )
+
+        prepareExecutor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=len(compiled_models_array),
+            thread_name_prefix="face-prepare",
+        )
+
+        return compiled_models, input_name, prepareExecutor, executor
+
+    async def predictDetectModel(self, input: Image.Image):
+        compiled_models, input_name, prepareExecutor, executor = self.detectModel
+
+        def prepare():
+            im = np.array(input)
+            im = np.expand_dims(input, axis=0)
+            im = im.transpose((0, 3, 1, 2))  # BHWC to BCHW, (n, 3, h, w)
+            im = im.astype(np.float32) / 255.0
+            im = np.ascontiguousarray(im)  # contiguous
+            return im
+
+        def predict(input_tensor):
+            compiled_model = compiled_models[threading.current_thread().name]
+            output_tensors = compiled_model.run(None, {input_name: input_tensor})
+            return output_tensors
+
+        input_tensor = await asyncio.get_event_loop().run_in_executor(
+            prepareExecutor, lambda: prepare()
+        )
+        objs = await asyncio.get_event_loop().run_in_executor(
+            executor, lambda: predict(input_tensor)
+        )
+
+        return objs[0][0]
+
+    async def predictFaceModel(self, input: np.ndarray):
+        compiled_models, input_name, prepareExecutor, executor = self.faceModel
+
+        def predict():
+            compiled_model = compiled_models[threading.current_thread().name]
+            output_tensors = compiled_model.run(None, {input_name: input})
+            return output_tensors
+
+        objs = await asyncio.get_event_loop().run_in_executor(
+            executor, lambda: predict()
+        )
+
+        return objs[0]