export cli

This commit is contained in:
mertalev
2025-03-12 22:50:29 -04:00
parent ce2a41826e
commit c57c562166
30 changed files with 1735 additions and 4781 deletions
+7 -4
View File
@@ -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:
+3
View File
@@ -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)
@@ -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