Cutlery/immich
1
0
Fork 1
mirror of https://github.com/immich-app/immich.git synced 2025-06-23 15:30:51 -04:00
Code Issues 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
machine-learning/app/sessions/rknn.py
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

10
machine-learning/rknn/rknnpool.py
View File

@ -4,6 +4,8 @@
import os
from concurrent.futures import ThreadPoolExecutor
from queue import Queue
import numpy as np
from numpy.typing import NDArray
supported_socs = ["rk3562", "rk3566", "rk3568", "rk3576", "rk3588"]
coremask_supported_socs = ["rk3576","rk3588"]
@ -63,7 +65,7 @@ def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", TPEs=1):
class rknnPoolExecutor:
def __init__(self, rknnModel, TPEs, func):
def __init__(self, rknnModel: str, TPEs: int, func):
self.TPEs = TPEs
self.queue = Queue()
self.rknnPool = initRKNNs(rknnModel, TPEs)
@ -71,17 +73,17 @@ class rknnPoolExecutor:
self.func = func
self.num = 0
def put(self, frame):
def put(self, frame) -> None:
self.queue.put(self.pool.submit(self.func, self.rknnPool[self.num % self.TPEs], frame))
self.num += 1
def get(self):
def get(self) -> list[list[NDArray[np.float32]], bool]:
if self.queue.empty():
return None, False
fut = self.queue.get()
return fut.result(), True
def release(self):
def release(self) -> None:
self.pool.shutdown()
for rknn_lite in self.rknnPool:
rknn_lite.release()