From f32d991131bbe3773ade47b78bab7b8a447724ce Mon Sep 17 00:00:00 2001 From: yoni13 Date: Fri, 17 Jan 2025 19:25:01 +0800 Subject: [PATCH] changes some cases --- machine-learning/app/conftest.py | 2 +- machine-learning/app/sessions/rknn.py | 4 ++-- machine-learning/pyproject.toml | 2 +- machine-learning/rknn/rknnpool.py | 18 +++++++++--------- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/machine-learning/app/conftest.py b/machine-learning/app/conftest.py index 8db3ba2d72..50c084215a 100644 --- a/machine-learning/app/conftest.py +++ b/machine-learning/app/conftest.py @@ -138,7 +138,7 @@ def ann_session() -> Iterator[mock.Mock]: @pytest.fixture(scope="function") 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 diff --git a/machine-learning/app/sessions/rknn.py b/machine-learning/app/sessions/rknn.py index 90bea605ca..585129c0ea 100644 --- a/machine-learning/app/sessions/rknn.py +++ b/machine-learning/app/sessions/rknn.py @@ -8,7 +8,7 @@ import onnxruntime as ort from numpy.typing import NDArray from app.schemas import SessionNode -from rknn.rknnpool import rknnPoolExecutor, soc_name +from rknn.rknnpool import RknnPoolExecutor, soc_name from ..config import log, settings @@ -32,7 +32,7 @@ class RknnSession: self.tpe = settings.rknn_facial_detection_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.") def __del__(self) -> None: diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml index c50a4b7fef..74e0861477 100644 --- a/machine-learning/pyproject.toml +++ b/machine-learning/pyproject.toml @@ -104,4 +104,4 @@ line-length = 120 target-version = ['py311'] [tool.pytest.ini_options] -markers = ["providers", "ov_device_ids"] \ No newline at end of file +markers = ["providers", "ov_device_ids"] diff --git a/machine-learning/rknn/rknnpool.py b/machine-learning/rknn/rknnpool.py index fd2209a240..73b40093d7 100644 --- a/machine-learning/rknn/rknnpool.py +++ b/machine-learning/rknn/rknnpool.py @@ -59,24 +59,24 @@ def initRKNN(rknnModel="./rknnModel/yolov5s.rknn", id=0): return rknn_lite -def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", TPEs=1): +def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", tpes=1): rknn_list = [] - for i in range(TPEs): + for i in range(tpes): rknn_list.append(initRKNN(rknnModel, i % 3)) return rknn_list -class rknnPoolExecutor: - def __init__(self, rknnModel: str, TPEs: int, func): - self.TPEs = TPEs +class RknnPoolExecutor: + def __init__(self, rknnModel: str, tpes: int, func): + self.tpes = tpes self.queue = Queue() - self.rknnPool = initRKNNs(rknnModel, TPEs) - self.pool = ThreadPoolExecutor(max_workers=TPEs) + self.rknn_pool = initRKNNs(rknnModel, tpes) + self.pool = ThreadPoolExecutor(max_workers=tpes) self.func = func self.num = 0 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 def get(self) -> list[list[NDArray[np.float32]], bool]: @@ -87,5 +87,5 @@ class rknnPoolExecutor: def release(self) -> None: self.pool.shutdown() - for rknn_lite in self.rknnPool: + for rknn_lite in self.rknn_pool: rknn_lite.release()