1import numpy as np 2import cv2 3from rknn.api import RKNN 4 5 6def show_outputs(outputs): 7 np.save('./caffe_mobilenet_v2_0.npy', outputs[0]) 8 output = outputs[0].reshape(-1) 9 output_sorted = sorted(output, reverse=True) 10 top5_str = 'mobilenet_v2\n-----TOP 5-----\n' 11 for i in range(5): 12 value = output_sorted[i] 13 index = np.where(output == value) 14 for j in range(len(index)): 15 if (i + j) >= 5: 16 break 17 if value > 0: 18 topi = '{}: {}\n'.format(index[j], value) 19 else: 20 topi = '-1: 0.0\n' 21 top5_str += topi 22 print(top5_str) 23 24 25if __name__ == '__main__': 26 27 # Create RKNN object 28 rknn = RKNN(verbose=True) 29 30 # Pre-process config 31 print('--> Config model') 32 rknn.config(mean_values=[103.94, 116.78, 123.68], std_values=[58.82, 58.82, 58.82], quant_img_RGB2BGR=True) 33 print('done') 34 35 # Load model 36 print('--> Loading model') 37 ret = rknn.load_caffe(model='./mobilenet_v2.prototxt', 38 blobs='./mobilenet_v2.caffemodel') 39 if ret != 0: 40 print('Load model failed!') 41 exit(ret) 42 print('done') 43 44 # Build model 45 print('--> Building model') 46 ret = rknn.build(do_quantization=True, dataset='./dataset.txt') 47 if ret != 0: 48 print('Build model failed!') 49 exit(ret) 50 print('done') 51 52 # Export rknn model 53 print('--> Export rknn model') 54 ret = rknn.export_rknn('./mobilenet_v2.rknn') 55 if ret != 0: 56 print('Export rknn model failed!') 57 exit(ret) 58 print('done') 59 60 # Set inputs 61 img = cv2.imread('./dog_224x224.jpg') 62 63 # Init runtime environment 64 print('--> Init runtime environment') 65 ret = rknn.init_runtime() 66 if ret != 0: 67 print('Init runtime environment failed!') 68 exit(ret) 69 print('done') 70 71 # Inference 72 print('--> Running model') 73 outputs = rknn.inference(inputs=[img]) 74 show_outputs(outputs) 75 print('done') 76 77 rknn.release() 78