export cli

machine-learning/app/models/base.py

@@ -33,7 +33,6 @@ class InferenceModel(ABC):
         self.model_name = clean_name(model_name)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else self._cache_dir_default
         self.model_format = model_format if model_format is not None else self._model_format_default
-        self.model_path_prefix = rknn.model_prefix if self.model_format == ModelFormat.RKNN else None
         if session is not None:
             self.session = session
 
@@ -121,15 +120,19 @@ class InferenceModel(ABC):
                 raise ValueError(f"Unsupported model file type: {model_path.suffix}")
         return session
 
+    def model_path_for_format(self, model_format: ModelFormat) -> Path:
+        model_path_prefix = rknn.model_prefix if model_format == ModelFormat.RKNN else None
+        if model_path_prefix:
+            return self.model_dir / model_path_prefix / f"model.{model_format}"
+        return self.model_dir / f"model.{model_format}"
+
     @property
     def model_dir(self) -> Path:
         return self.cache_dir / self.model_type.value
 
     @property
     def model_path(self) -> Path:
-        if self.model_path_prefix:
-            return self.model_dir / self.model_path_prefix / f"model.{self.model_format}"
-        return self.model_dir / f"model.{self.model_format}"
+        return self.model_path_for_format(self.model_format)
 
     @property
     def model_task(self) -> ModelTask:

machine-learning/app/models/constants.py

@@ -65,6 +65,9 @@ _INSIGHTFACE_MODELS = {
 
 SUPPORTED_PROVIDERS = ["CUDAExecutionProvider", "OpenVINOExecutionProvider", "CPUExecutionProvider"]
 
+RKNN_SUPPORTED_SOCS = ["rk3566", "rk3588"]
+RKNN_COREMASK_SUPPORTED_SOCS = ["rk3576", "rk3588"]
+
 
 def get_model_source(model_name: str) -> ModelSource | None:
     cleaned_name = clean_name(model_name)

machine-learning/app/models/facial_recognition/recognition.py

@@ -31,7 +31,7 @@ class FaceRecognizer(InferenceModel):
             self._add_batch_axis(self.model_path)
             session = self._make_session(self.model_path)
         self.model = ArcFaceONNX(
-            self.model_path.with_suffix(".onnx").as_posix(),
+            self.model_path_for_format(ModelFormat.ONNX).as_posix(),
             session=session,
         )
         return session

machine-learning/app/sessions/rknn/rknnpool.py

@@ -9,16 +9,14 @@ import numpy as np
 from numpy.typing import NDArray
 
 from app.config import log
-
-supported_socs = ["rk3566", "rk3588"]
-coremask_supported_socs = ["rk3576", "rk3588"]
+from app.models.constants import RKNN_COREMASK_SUPPORTED_SOCS, RKNN_SUPPORTED_SOCS
 
 
 def get_soc(device_tree_path: Path | str) -> str | None:
     try:
         with Path(device_tree_path).open() as f:
             device_compatible_str = f.read()
-            for soc in supported_socs:
+            for soc in RKNN_SUPPORTED_SOCS:
                 if soc in device_compatible_str:
                     return soc
             log.warning("Device is not supported for RKNN")
@@ -46,7 +44,7 @@ def init_rknn(model_path: str) -> "RKNNLite":
     if ret != 0:
         raise RuntimeError("Load RKNN rknnModel failed")
 
-    if soc_name in coremask_supported_socs:
+    if soc_name in RKNN_COREMASK_SUPPORTED_SOCS:
         ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
     else:
         ret = rknn_lite.init_runtime()  # Please do not set this parameter on other platforms.