Adjust rknn to use RKNNRunner

frigate/detectors/base_runner.py

@@ -6,8 +6,7 @@ from typing import Any
 import onnxruntime as ort
 
 from frigate.detectors.plugins.openvino import OpenVINOModelRunner
-from frigate.detectors.plugins.onnx import CudaGraphRunner
+from frigate.detectors.plugins.rknn import RKNNModelRunner
-from frigate.embeddings.onnx.runner import RKNNModelRunner
 from frigate.util.model import get_ort_providers
 from frigate.util.rknn_converter import auto_convert_model, is_rknn_compatible
 
@@ -61,7 +60,7 @@ def get_optimized_runner(model_path: str, device: str, **kwargs) -> BaseModelRun
         rknn_path = auto_convert_model(model_path)
 
         if rknn_path:
-            return RKNNModelRunner(rknn_path, device)
+            return RKNNModelRunner(rknn_path)
 
     providers, options = get_ort_providers(device == "CPU", device, **kwargs)
 
@@ -74,9 +73,4 @@ def get_optimized_runner(model_path: str, device: str, **kwargs) -> BaseModelRun
         provider_options=options,
     )
 
-    cuda_idx = providers.index("CUDAExecutionProvider")
-
-    if cuda_idx == 0:
-        return CudaGraphRunner(ort, options[cuda_idx].get("device_id", 0))
-
     return ONNXModelRunner(model_path, device, **kwargs)

frigate/detectors/plugins/rknn.py

@@ -9,6 +9,7 @@ import numpy as np
 from pydantic import Field
 
 from frigate.const import MODEL_CACHE_DIR
+from frigate.detectors.base_runner import BaseModelRunner
 from frigate.detectors.detection_api import DetectionApi
 from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
 from frigate.util.model import post_process_yolo
@@ -34,12 +35,11 @@ class RknnDetectorConfig(BaseDetectorConfig):
     num_cores: int = Field(default=0, ge=0, le=3, title="Number of NPU cores to use.")
 
 
-class RKNNModelRunner:
+class RKNNModelRunner(BaseModelRunner):
     """Run RKNN models for embeddings."""
 
-    def __init__(self, model_path: str, device: str = "AUTO", model_type: str = None):
+    def __init__(self, model_path: str, model_type: str = None):
         self.model_path = model_path
-        self.device = device
         self.model_type = model_type
         self.rknn = None
         self._load_model()
@@ -70,6 +70,7 @@ class RKNNModelRunner:
 
     def get_input_names(self) -> list[str]:
         """Get input names for the model."""
+        # For detection models, we typically use "input" as the default input name
         # For CLIP models, we need to determine the model type from the path
         model_name = os.path.basename(self.model_path).lower()
 
@@ -94,6 +95,8 @@ class RKNNModelRunner:
             return 224  # CLIP V1 uses 224x224
         elif "arcface" in model_name:
             return 112
+        # For detection models, we can't easily determine this from the RKNN model
+        # The calling code should provide this information
         return -1
 
     def run(self, inputs: dict[str, Any]) -> Any:
@@ -161,18 +164,18 @@ class Rknn(DetectionApi):
                     "For more information, see: https://docs.deci.ai/super-gradients/latest/LICENSE.YOLONAS.html"
                 )
 
-        from rknnlite.api import RKNNLite
+        # Initialize the RKNN model runner
-
+        self.runner = RKNNModelRunner(
-        self.rknn = RKNNLite(verbose=False)
+            model_path=model_props["path"],
-        if self.rknn.load_rknn(model_props["path"]) != 0:
+            model_type=config.model.model_type.value
-            logger.error("Error initializing rknn model.")
+            if config.model.model_type
-        if self.rknn.init_runtime(core_mask=core_mask) != 0:
+            else None,
-            logger.error(
+        )
-                "Error initializing rknn runtime. Do you run docker in privileged mode?"
-            )
 
     def __del__(self):
-        self.rknn.release()
+        if hasattr(self, "runner") and self.runner:
+            # The runner's __del__ method will handle cleanup
+            pass
 
     def get_soc(self):
         try:
@@ -405,9 +408,7 @@ class Rknn(DetectionApi):
             )
 
     def detect_raw(self, tensor_input):
-        output = self.rknn.inference(
+        # Prepare input for the runner
-            [
+        inputs = {"input": tensor_input}
-                tensor_input,
+        output = self.runner.run(inputs)
-            ]
-        )
         return self.post_process(output)