File size: 4,258 Bytes
477da44 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import cv2
import numpy as np
import platform
from synset_label import labels
from rknnlite.api import RKNNLite
# decice tree for RK356x/RK3576/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'
elif 'rk3576' in device_compatible_str:
host = 'RK3576'
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'
RK3576_RKNN_MODEL = 'mobilenet_v2_for_rk3576.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 == 'RK3576':
rknn_model = RK3576_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 / RK3576 / RK3588 with Debian OS, do not need specify target.
if host_name in ['RK3576', 'RK3588']:
# For RK3576 / 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()
|