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

ruff format

Browse Source
This commit is contained in:
yoni13 2025-01-11 16:26:17 +08:00
parent 7f2af6f819
commit d5e453a773
3 changed files with 13 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
machine-learning/app/models/base.py
View File

@ -162,7 +162,7 @@ class InferenceModel(ABC):
def _model_format_default(self) -> ModelFormat: def _model_format_default(self) -> ModelFormat:
if rknn.rknnpool.is_available and settings.rknn: if rknn.rknnpool.is_available and settings.rknn:
return ModelFormat.RKNN return ModelFormat.RKNN
elif ann.ann.is_available and settings.ann : elif ann.ann.is_available and settings.ann:
return ModelFormat.ARMNN return ModelFormat.ARMNN
else: else:
return ModelFormat.ONNX return ModelFormat.ONNX

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

@ -12,24 +12,21 @@ from app.schemas import SessionNode
from ..config import log from ..config import log
def runInfrence(rknn_lite, input): def runInfrence(rknn_lite, input):
outputs = rknn_lite.inference(inputs=[input], data_format='nchw') outputs = rknn_lite.inference(inputs=[input], data_format="nchw")
return outputs return outputs
class RknnSession: class RknnSession:
def __init__(self, model_path: Path | str): def __init__(self, model_path: Path | str):
self.model_path = Path(model_path) self.model_path = Path(model_path)
self.ort_model_path = str(self.model_path).replace(".rknn", ".onnx") self.ort_model_path = str(self.model_path).replace(".rknn", ".onnx")
self.tpe = 1 if 'textual' in str(self.model_path) else 2 self.tpe = 1 if "textual" in str(self.model_path) else 2
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( self.rknnpool = rknnPoolExecutor(rknnModel=self.model_path.as_posix(), TPEs=self.tpe, func=runInfrence)
rknnModel=self.model_path.as_posix(),
TPEs= self.tpe,
func=runInfrence)
self.ort_session = ort.InferenceSession( self.ort_session = ort.InferenceSession(
self.ort_model_path, self.ort_model_path,
@ -39,7 +36,6 @@ class RknnSession:
del self.ort_session del self.ort_session
def __del__(self): def __del__(self):
self.rknnpool.release() self.rknnpool.release()
@ -53,11 +49,9 @@ class RknnSession:
self, self,
output_names: list[str] | None, output_names: list[str] | None,
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]], input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
run_options: Any = None, run_options: Any = None,
): ):
input_data = [np.ascontiguousarray(v) for v in input_feed.values()][0] input_data = [np.ascontiguousarray(v) for v in input_feed.values()][0]
self.rknnpool.put(input_data) self.rknnpool.put(input_data)
outputs = self.rknnpool.get()[0] outputs = self.rknnpool.get()[0]
return outputs return outputs

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

@ -4,11 +4,12 @@
from queue import Queue from queue import Queue
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
supported_socs = ["rk3566","rk3568","rk3588"] supported_socs = ["rk3566", "rk3568", "rk3588"]
try: try:
from rknnlite.api import RKNNLite from rknnlite.api import RKNNLite
with open('/proc/device-tree/compatible') as f:
with open("/proc/device-tree/compatible") as f:
# Keep in mind that this is not in container by default. # Keep in mind that this is not in container by default.
# and this way of checking can't check if the rknpu driver is running or not. # and this way of checking can't check if the rknpu driver is running or not.
device_compatible_str = f.read() device_compatible_str = f.read()
@ -43,7 +44,7 @@ def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", TPEs=1):
return rknn_list return rknn_list
class rknnPoolExecutor(): class rknnPoolExecutor:
def __init__(self, rknnModel, TPEs, func): def __init__(self, rknnModel, TPEs, func):
self.TPEs = TPEs self.TPEs = TPEs
self.queue = Queue() self.queue = Queue()
@ -53,8 +54,7 @@ class rknnPoolExecutor():
self.num = 0 self.num = 0
def put(self, frame): def put(self, frame):
self.queue.put(self.pool.submit( self.queue.put(self.pool.submit(self.func, self.rknnPool[self.num % self.TPEs], frame))
self.func, self.rknnPool[self.num % self.TPEs], frame))
self.num += 1 self.num += 1
def get(self): def get(self):
@ -66,4 +66,4 @@ class rknnPoolExecutor():
def release(self): def release(self):
self.pool.shutdown() self.pool.shutdown()
for rknn_lite in self.rknnPool: for rknn_lite in self.rknnPool:
rknn_lite.release() rknn_lite.release()