1import numpy as np 2import cv2 3from rknn.api import RKNN 4 5 6def show_outputs(outputs): 7 output_ = outputs[0].reshape((-1, 1000)) 8 for output in output_: 9 output_sorted = sorted(output, reverse=True) 10 top5_str = 'mobilenet_v1\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 25def show_perfs(perfs): 26 perfs = 'perfs: {}\n'.format(outputs) 27 print(perfs) 28 29 30if __name__ == '__main__': 31 32 # Create RKNN object 33 rknn = RKNN(verbose=True) 34 35 # Pre-process config 36 print('--> Config model') 37 rknn.config(mean_values=[103.94, 116.78, 123.68], std_values=[58.82, 58.82, 58.82], quant_img_RGB2BGR=True) 38 print('done') 39 40 # Load model 41 print('--> Loading model') 42 ret = rknn.load_caffe(model='../../caffe/mobilenet_v2/mobilenet_v2.prototxt', 43 blobs='../../caffe/mobilenet_v2/mobilenet_v2.caffemodel') 44 if ret != 0: 45 print('Load model failed!') 46 exit(ret) 47 print('done') 48 49 # Build model 50 print('--> Building model') 51 ret = rknn.build(do_quantization=True, dataset='./dataset.txt', rknn_batch_size=4) 52 if ret != 0: 53 print('Build model failed!') 54 exit(ret) 55 print('done') 56 57 # Export rknn model 58 print('--> Export rknn model') 59 ret = rknn.export_rknn('./mobilenet_v2.rknn') 60 if ret != 0: 61 print('Export rknn model failed!') 62 exit(ret) 63 print('done') 64 65 # Set inputs 66 img = cv2.imread('./dog_224x224.jpg') 67 68 img = np.expand_dims(img, 0) 69 img = np.concatenate((img, img, img, img), axis=0) 70 71 # Init runtime environment 72 print('--> Init runtime environment') 73 ret = rknn.init_runtime() 74 if ret != 0: 75 print('Init runtime environment failed!') 76 exit(ret) 77 print('done') 78 79 # Inference 80 print('--> Running model') 81 outputs = rknn.inference(inputs=[img]) 82 np.save('./functions_batch_size_0.npy', outputs[0]) 83 show_outputs(outputs) 84 print('done') 85 86 rknn.release() 87