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

changes some cases

Browse Source
This commit is contained in:
yoni13 2025-01-17 19:25:01 +08:00
parent 9882b83cd4
commit f32d991131
4 changed files with 13 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
machine-learning/app/conftest.py
View File

@ -138,7 +138,7 @@ def ann_session() -> Iterator[mock.Mock]:
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
def rknn_session() -> Iterator[mock.Mock]: def rknn_session() -> Iterator[mock.Mock]:
with mock.patch("app.sessions.rknn.rknnPoolExecutor") as mocked: with mock.patch("app.sessions.rknn.RknnPoolExecutor") as mocked:
yield mocked yield mocked

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

@ -8,7 +8,7 @@ import onnxruntime as ort
from numpy.typing import NDArray from numpy.typing import NDArray
from app.schemas import SessionNode from app.schemas import SessionNode
from rknn.rknnpool import rknnPoolExecutor, soc_name from rknn.rknnpool import RknnPoolExecutor, soc_name
from ..config import log, settings from ..config import log, settings
@ -32,7 +32,7 @@ class RknnSession:
self.tpe = settings.rknn_facial_detection_threads self.tpe = settings.rknn_facial_detection_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.")
self.rknnpool = rknnPoolExecutor(rknnModel=self.model_path.as_posix(), TPEs=self.tpe, func=runInfrence) self.rknnpool = RknnPoolExecutor(rknnModel=self.model_path.as_posix(), TPEs=self.tpe, func=runInfrence)
log.info(f"Loaded RKNN model from {self.model_path} with {self.tpe} threads.") log.info(f"Loaded RKNN model from {self.model_path} with {self.tpe} threads.")
def __del__(self) -> None: def __del__(self) -> None:

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

@ -59,24 +59,24 @@ def initRKNN(rknnModel="./rknnModel/yolov5s.rknn", id=0):
return rknn_lite return rknn_lite
def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", TPEs=1): def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", tpes=1):
rknn_list = [] rknn_list = []
for i in range(TPEs): for i in range(tpes):
rknn_list.append(initRKNN(rknnModel, i % 3)) rknn_list.append(initRKNN(rknnModel, i % 3))
return rknn_list return rknn_list
class rknnPoolExecutor: class RknnPoolExecutor:
def __init__(self, rknnModel: str, TPEs: int, func): def __init__(self, rknnModel: str, tpes: int, func):
self.TPEs = TPEs self.tpes = tpes
self.queue = Queue() self.queue = Queue()
self.rknnPool = initRKNNs(rknnModel, TPEs) self.rknn_pool = initRKNNs(rknnModel, tpes)
self.pool = ThreadPoolExecutor(max_workers=TPEs) self.pool = ThreadPoolExecutor(max_workers=tpes)
self.func = func self.func = func
self.num = 0 self.num = 0
def put(self, frame) -> None: def put(self, frame) -> None:
self.queue.put(self.pool.submit(self.func, self.rknnPool[self.num % self.TPEs], frame)) self.queue.put(self.pool.submit(self.func, self.rknn_pool[self.num % self.tpes], frame))
self.num += 1 self.num += 1
def get(self) -> list[list[NDArray[np.float32]], bool]: def get(self) -> list[list[NDArray[np.float32]], bool]:
@ -87,5 +87,5 @@ class rknnPoolExecutor:
def release(self) -> None: def release(self) -> None:
self.pool.shutdown() self.pool.shutdown()
for rknn_lite in self.rknnPool: for rknn_lite in self.rknn_pool:
rknn_lite.release() rknn_lite.release()