27 lines
674 B
Python
27 lines
674 B
Python
from rknnlite.api import RKNNLite as RKNN
|
|
|
|
class RKNN_model_container():
|
|
def __init__(self, model_path, target=None, device_id=None) -> None:
|
|
rknn = RKNN()
|
|
rknn.load_rknn(model_path)
|
|
ret = rknn.init_runtime()
|
|
self.rknn = rknn
|
|
|
|
def run(self, inputs):
|
|
if self.rknn is None:
|
|
print("ERROR: rknn has been released")
|
|
return []
|
|
|
|
if isinstance(inputs, list) or isinstance(inputs, tuple):
|
|
pass
|
|
else:
|
|
inputs = [inputs]
|
|
|
|
result = self.rknn.inference(inputs=inputs)
|
|
|
|
return result
|
|
|
|
def release(self):
|
|
self.rknn.release()
|
|
self.rknn = None
|