cjones/immich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

typing be happy.

Browse Source
This commit is contained in:
yoni13
2025-01-13 18:24:12 +08:00
parent ebdfe1b7b6
commit 2f7e44aa63
2 changed files with 30 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
machine-learning/app/sessions/rknn.py
+24 -24
View File
@@ -13,8 +13,8 @@ from rknn.rknnpool import rknnPoolExecutor, soc_name
from ..config import log, settings
def runInfrence(rknn_lite, input):
outputs = rknn_lite.inference(inputs=[input], data_format="nchw")
def runInfrence(rknn_lite: Any, input: NDArray[np.float32]) -> list[NDArray[np.float32]]:
outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=[input], data_format="nchw")
return outputs
@@ -23,8 +23,6 @@ class RknnSession:
def __init__(self, model_path: Path | str):
self.model_path = Path(str(model_path).replace("model", soc_name))
self.ort_model_path = Path(str(self.model_path).replace(f"{soc_name}.rknn", "model.onnx"))
self.inputs = None
self.outputs = None
if "textual" in str(self.model_path):
self.tpe = settings.rknn_textual_threads
@@ -36,36 +34,38 @@ class RknnSession:
log.info(f"Loading RKNN model from {self.model_path} with {self.tpe} threads.")
self.rknnpool = rknnPoolExecutor(rknnModel=self.model_path.as_posix(), TPEs=self.tpe, func=runInfrence)
def __del__(self):
def __del__(self) -> None:
self.rknnpool.release()
def _load_ort_session(self) -> None:
self.ort_session = ort.InferenceSession(
self.ort_model_path.as_posix(),
)
self.inputs: list[SessionNode] = self.ort_session.get_inputs()
self.outputs: list[SessionNode] = self.ort_session.get_outputs()
del self.ort_session
def get_inputs(self) -> list[SessionNode]:
if not self.inputs:
self.ort_session = ort.InferenceSession(
self.ort_model_path.as_posix(),
)
self.inputs = self.ort_session.get_inputs()
self.outputs = self.ort_session.get_outputs()
del self.ort_session
return self.inputs
try:
return self.inputs
except AttributeError:
self._load_ort_session()
return self.inputs
def get_outputs(self) -> list[SessionNode]:
if not self.outputs:
self.ort_session = ort.InferenceSession(
self.ort_model_path.as_posix(),
)
self.inputs = self.ort_session.get_inputs()
self.outputs = self.ort_session.get_outputs()
del self.ort_session
return self.outputs
try:
return self.outputs
except AttributeError:
self._load_ort_session()
return self.outputs
def run(
self,
output_names: list[str] | None,
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
run_options: Any = None,
):
input_data = [np.ascontiguousarray(v) for v in input_feed.values()][0]
) -> list[NDArray[np.float32]]:
input_data: NDArray[np.float32] = np.ascontiguousarray(list(input_feed.values())[0], dtype=np.float32)
self.rknnpool.put(input_data)
outputs = self.rknnpool.get()[0]
outputs: list[NDArray[np.float32]] = self.rknnpool.get()[0]
return outputs