131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
import cv2
|
|
import numpy as np
|
|
import platform
|
|
from synset_label import labels
|
|
from rknnlite.api import RKNNLite
|
|
|
|
# decice tree for RK356x/RK3588
|
|
DEVICE_COMPATIBLE_NODE = '/proc/device-tree/compatible'
|
|
|
|
def get_host():
|
|
# get platform and device type
|
|
system = platform.system()
|
|
machine = platform.machine()
|
|
os_machine = system + '-' + machine
|
|
if os_machine == 'Linux-aarch64':
|
|
try:
|
|
with open(DEVICE_COMPATIBLE_NODE) as f:
|
|
device_compatible_str = f.read()
|
|
if 'rk3588' in device_compatible_str:
|
|
host = 'RK3588'
|
|
elif 'rk3562' in device_compatible_str:
|
|
host = 'RK3562'
|
|
else:
|
|
host = 'RK3566_RK3568'
|
|
except IOError:
|
|
print('Read device node {} failed.'.format(DEVICE_COMPATIBLE_NODE))
|
|
exit(-1)
|
|
else:
|
|
host = os_machine
|
|
return host
|
|
|
|
INPUT_SIZE = 224
|
|
|
|
RK3566_RK3568_RKNN_MODEL = 'mobilenet_v2_for_rk3566_rk3568.rknn'
|
|
RK3588_RKNN_MODEL = 'mobilenet_v2_for_rk3588.rknn'
|
|
RK3562_RKNN_MODEL = 'mobilenet_v2_for_rk3562.rknn'
|
|
|
|
|
|
def show_top5(result):
|
|
output = result[0].reshape(-1)
|
|
# Get the indices of the top 5 largest values
|
|
output_sorted_indices = np.argsort(output)[::-1][:5]
|
|
top5_str = '-----TOP 5-----\n'
|
|
for i, index in enumerate(output_sorted_indices):
|
|
value = output[index]
|
|
if value > 0:
|
|
topi = '[{:>3d}] score:{:.6f} class:"{}"\n'.format(
|
|
index, value, labels[index])
|
|
else:
|
|
topi = '-1: 0.0\n'
|
|
top5_str += topi
|
|
print(top5_str)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# Get device information
|
|
host_name = get_host()
|
|
if host_name == 'RK3566_RK3568':
|
|
rknn_model = RK3566_RK3568_RKNN_MODEL
|
|
elif host_name == 'RK3562':
|
|
rknn_model = RK3562_RKNN_MODEL
|
|
elif host_name == 'RK3588':
|
|
rknn_model = RK3588_RKNN_MODEL
|
|
else:
|
|
print("This demo cannot run on the current platform: {}".format(host_name))
|
|
exit(-1)
|
|
|
|
dynamic_input = [
|
|
[[1, 3, 192, 192]],
|
|
[[1, 3, 256, 256]],
|
|
[[1, 3, 160, 160]],
|
|
[[1, 3, 224, 224]]
|
|
]
|
|
|
|
rknn_lite = RKNNLite()
|
|
|
|
# Load RKNN model
|
|
print('--> Load RKNN model')
|
|
ret = rknn_lite.load_rknn(rknn_model)
|
|
if ret != 0:
|
|
print('Load RKNN model failed')
|
|
exit(ret)
|
|
print('done')
|
|
|
|
img = cv2.imread('./dog_224x224.jpg')
|
|
|
|
# Init runtime environment
|
|
print('--> Init runtime environment')
|
|
# run on RK356x/RK3588 with Debian OS, do not need specify target.
|
|
if host_name == 'RK3588':
|
|
# For RK3588, specify which NPU core the model runs on through the core_mask parameter.
|
|
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
|
else:
|
|
ret = rknn_lite.init_runtime()
|
|
if ret != 0:
|
|
print('Init runtime environment failed')
|
|
exit(ret)
|
|
print('done')
|
|
|
|
# Inference
|
|
print('--> Running model')
|
|
print('model: mobilenet_v2\n')
|
|
print('input shape: 1,3,224,224')
|
|
real_img = cv2.resize(img, (224, 224))
|
|
real_img = np.expand_dims(real_img, 0)
|
|
real_img = np.transpose(real_img, (0, 3, 1, 2))
|
|
outputs = rknn_lite.inference(inputs=[real_img], data_format=['nchw'])
|
|
# Show the classification results
|
|
show_top5(outputs)
|
|
|
|
print('input shape: 1,3,160,160')
|
|
real_img = cv2.resize(img, (160, 160))
|
|
real_img = np.expand_dims(real_img, 0)
|
|
real_img = np.transpose(real_img, (0, 3, 1, 2))
|
|
outputs = rknn_lite.inference(inputs=[real_img], data_format=['nchw'])
|
|
# Show the classification results
|
|
show_top5(outputs)
|
|
|
|
print('input shape: 1,3,256,256')
|
|
real_img = cv2.resize(img, (256, 256))
|
|
real_img = np.expand_dims(real_img, 0)
|
|
real_img = np.transpose(real_img, (0, 3, 1, 2))
|
|
outputs = rknn_lite.inference(inputs=[real_img], data_format=['nchw'])
|
|
# Show the classification results
|
|
show_top5(outputs)
|
|
|
|
print('done')
|
|
|
|
rknn_lite.release()
|