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.
86 lines
2.3 KiB
86 lines
2.3 KiB
import numpy as np |
|
import cv2 |
|
from rknn.api import RKNN |
|
|
|
|
|
def show_outputs(outputs): |
|
output_ = outputs[0].reshape((-1, 1000)) |
|
for output in output_: |
|
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) |
|
|
|
|
|
def show_perfs(perfs): |
|
perfs = 'perfs: {}\n'.format(outputs) |
|
print(perfs) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
# Create RKNN object |
|
rknn = RKNN(verbose=True) |
|
|
|
# Pre-process config |
|
print('--> Config model') |
|
rknn.config(mean_values=[103.94, 116.78, 123.68], std_values=[58.82, 58.82, 58.82], quant_img_RGB2BGR=True) |
|
print('done') |
|
|
|
# Load model |
|
print('--> Loading model') |
|
ret = rknn.load_caffe(model='../../caffe/mobilenet_v2/mobilenet_v2.prototxt', |
|
blobs='../../caffe/mobilenet_v2/mobilenet_v2.caffemodel') |
|
if ret != 0: |
|
print('Load model failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Build model |
|
print('--> Building model') |
|
ret = rknn.build(do_quantization=True, dataset='./dataset.txt', rknn_batch_size=4) |
|
if ret != 0: |
|
print('Build model failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Export rknn model |
|
print('--> Export rknn model') |
|
ret = rknn.export_rknn('./mobilenet_v2.rknn') |
|
if ret != 0: |
|
print('Export rknn model failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Set inputs |
|
img = cv2.imread('./dog_224x224.jpg') |
|
|
|
img = np.expand_dims(img, 0) |
|
img = np.concatenate((img, img, img, img), axis=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]) |
|
np.save('./functions_batch_size_0.npy', outputs[0]) |
|
show_outputs(outputs) |
|
print('done') |
|
|
|
rknn.release()
|
|
|