fix typo and add a propper var name

This commit is contained in:
yoni13 2025-01-18 10:46:57 +08:00
parent 3634ae1f5b
commit f5de3de163
2 changed files with 6 additions and 6 deletions

View File

@ -13,7 +13,7 @@ from rknn.rknnpool import RknnPoolExecutor, soc_name
from ..config import log, settings from ..config import log, settings
def runInfrence(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]: def runInference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]:
outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=input, data_format="nchw") outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=input, data_format="nchw")
return outputs return outputs
@ -27,7 +27,7 @@ class RknnSession:
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.")
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=runInference)
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:

View File

@ -31,7 +31,7 @@ except (FileNotFoundError, ImportError):
soc_name = None soc_name = None
def initRKNN(rknnModel="./rknnModel/yolov5s.rknn", id=0): def initRKNN(rknnModel, id) -> RKNNLite:
if not is_available: if not is_available:
raise RuntimeError("rknn is not available!") raise RuntimeError("rknn is not available!")
rknn_lite = RKNNLite() rknn_lite = RKNNLite()
@ -59,7 +59,7 @@ 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, tpes) -> list[RKNNLite]:
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))
@ -75,8 +75,8 @@ class RknnPoolExecutor:
self.func = func self.func = func
self.num = 0 self.num = 0
def put(self, frame) -> None: def put(self, inputs) -> None:
self.queue.put(self.pool.submit(self.func, self.rknn_pool[self.num % self.tpes], frame)) self.queue.put(self.pool.submit(self.func, self.rknn_pool[self.num % self.tpes], inputs))
self.num += 1 self.num += 1
def get(self) -> list[list[NDArray[np.float32]], bool]: def get(self) -> list[list[NDArray[np.float32]], bool]: