You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
77 lines
1.9 KiB
77 lines
1.9 KiB
import numpy as np |
|
import cv2 |
|
from rknn.api import RKNN |
|
|
|
|
|
def show_outputs(outputs): |
|
output = outputs[0][0] |
|
output_sorted = sorted(output, reverse=True) |
|
top5_str = 'mobilenet_v1\n-----TOP 5-----\n' |
|
for i in range(5): |
|
value = output_sorted[i] |
|
index = np.where(output == value) |
|
for j in range(len(index)): |
|
if (i + j) >= 5: |
|
break |
|
if value > 0: |
|
topi = '{}: {}\n'.format(index[j], value) |
|
else: |
|
topi = '-1: 0.0\n' |
|
top5_str += topi |
|
print(top5_str) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
# Create RKNN object |
|
rknn = RKNN() |
|
|
|
# pre-process config |
|
print('--> config model') |
|
rknn.config(mean_values=[128, 128, 128], std_values=[128, 128, 128]) |
|
print('done') |
|
|
|
# Load tensorflow model |
|
print('--> Loading model') |
|
ret = rknn.load_tflite(model='mobilenet_v1_1.0_224.tflite') |
|
if ret != 0: |
|
print('Load mobilenet_v1 failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Build model |
|
print('--> Building model') |
|
ret = rknn.build(do_quantization=True, dataset='./dataset.txt') |
|
if ret != 0: |
|
print('Build mobilenet_v1 failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Export rknn model |
|
print('--> Export RKNN model') |
|
ret = rknn.export_rknn('./mobilenet_v1.rknn') |
|
if ret != 0: |
|
print('Export mobilenet_v1.rknn failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Set inputs |
|
img = cv2.imread('./dog_224x224.jpg') |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
img = np.expand_dims(img, 0) |
|
|
|
# init runtime environment |
|
print('--> Init runtime environment') |
|
ret = rknn.init_runtime() |
|
if ret != 0: |
|
print('Init runtime environment failed') |
|
exit(ret) |
|
print('done') |
|
|
|
# Inference |
|
print('--> Running model') |
|
outputs = rknn.inference(inputs=[img]) |
|
show_outputs(outputs) |
|
print('done') |
|
|
|
rknn.release()
|
|
|