Cutlery/immich
1
0
Fork 1
mirror of https://github.com/immich-app/immich.git synced 2025-07-09 03:04:16 -04:00
Code Issues Packages Projects Releases Wiki Activity

update mapping

Browse Source
This commit is contained in:
yoni13 2025-01-19 13:15:38 +00:00
parent ac4ce3ea9c
commit 20ba9f97e9
1 changed files with 18 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
machine-learning/app/sessions/rknn.py
View File

@ -18,8 +18,7 @@ def runInference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArr
return outputs return outputs
input_output_mapping = { input_output_mapping: dict[str, dict[str, Any]] = {
"buffalo_l": {
"detection": { "detection": {
"input": {"norm_tensor:0": (1, 3, 640, 640)}, "input": {"norm_tensor:0": (1, 3, 640, 640)},
"output": { "output": {
@ -36,13 +35,12 @@ input_output_mapping = {
}, },
"recognition": {"input": {"norm_tensor:0": (1, 3, 112, 112)}, "output": {"norm_tensor:1": (1, 512)}}, "recognition": {"input": {"norm_tensor:0": (1, 3, 112, 112)}, "output": {"norm_tensor:1": (1, 512)}},
} }
}
class RknnSession: class RknnSession:
def __init__(self, model_path: Path | str): def __init__(self, model_path: Path | str):
self.model_path = Path(str(model_path).replace("model", soc_name)) self.model_path = Path(str(model_path).replace("model", soc_name))
self.model_type = "detection" if "detection" in self.model_path.as_posix() else "recognition"
self.tpe = settings.rknn_threads self.tpe = settings.rknn_threads
log.info(f"Loading RKNN model from {self.model_path} with {self.tpe} threads.") log.info(f"Loading RKNN model from {self.model_path} with {self.tpe} threads.")
@ -53,24 +51,10 @@ class RknnSession:
self.rknnpool.release() self.rknnpool.release()
def get_inputs(self) -> list[SessionNode]: def get_inputs(self) -> list[SessionNode]:
for model_name in input_output_mapping: return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["input"].items()]
if model_name in self.model_path.as_posix():
model_type = "detection" if "detection" in self.model_path.as_posix() else "recognition"
return [
RknnNode(name=k, shape=v)
for k, v in input_output_mapping[model_name][model_type]["input"].items()
]
raise ValueError(f"Model {self.model_path} not found in input_output_mapping.")
def get_outputs(self) -> list[SessionNode]: def get_outputs(self) -> list[SessionNode]:
for model_name in input_output_mapping: return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["output"].items()]
if model_name in self.model_path.as_posix():
model_type = "detection" if "detection" in self.model_path.as_posix() else "recognition"
return [
RknnNode(name=k, shape=v)
for k, v in input_output_mapping[model_name][model_type]["output"].items()
]
raise ValueError(f"Model {self.model_path} not found in input_output_mapping.")
def run( def run(
self, self,