mirror of
https://github.com/immich-app/immich.git
synced 2025-07-09 03:04:16 -04:00
ruff format
This commit is contained in:
parent
7f2af6f819
commit
d5e453a773
@ -12,24 +12,21 @@ from app.schemas import SessionNode
|
||||
|
||||
from ..config import log
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class RknnSession:
|
||||
def __init__(self, model_path: Path | str):
|
||||
|
||||
self.model_path = Path(model_path)
|
||||
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.")
|
||||
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)
|
||||
|
||||
self.ort_session = ort.InferenceSession(
|
||||
self.ort_model_path,
|
||||
@ -39,7 +36,6 @@ class RknnSession:
|
||||
|
||||
del self.ort_session
|
||||
|
||||
|
||||
def __del__(self):
|
||||
self.rknnpool.release()
|
||||
|
||||
@ -55,9 +51,7 @@ class RknnSession:
|
||||
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]
|
||||
self.rknnpool.put(input_data)
|
||||
outputs = self.rknnpool.get()[0]
|
||||
return outputs
|
||||
|
||||
|
@ -8,7 +8,8 @@ supported_socs = ["rk3566","rk3568","rk3588"]
|
||||
|
||||
try:
|
||||
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.
|
||||
# and this way of checking can't check if the rknpu driver is running or not.
|
||||
device_compatible_str = f.read()
|
||||
@ -43,7 +44,7 @@ def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", TPEs=1):
|
||||
return rknn_list
|
||||
|
||||
|
||||
class rknnPoolExecutor():
|
||||
class rknnPoolExecutor:
|
||||
def __init__(self, rknnModel, TPEs, func):
|
||||
self.TPEs = TPEs
|
||||
self.queue = Queue()
|
||||
@ -53,8 +54,7 @@ class rknnPoolExecutor():
|
||||
self.num = 0
|
||||
|
||||
def put(self, frame):
|
||||
self.queue.put(self.pool.submit(
|
||||
self.func, self.rknnPool[self.num % self.TPEs], frame))
|
||||
self.queue.put(self.pool.submit(self.func, self.rknnPool[self.num % self.TPEs], frame))
|
||||
self.num += 1
|
||||
|
||||
def get(self):
|
||||
|
Loading…
x
Reference in New Issue
Block a user