1import numpy as np 2import cv2 3from rknn.api import RKNN 4 5 6def show_outputs(outputs): 7 output = outputs[0][0] 8 output_sorted = sorted(output, reverse=True) 9 top5_str = 'mobilenet_v1\n-----TOP 5-----\n' 10 for i in range(5): 11 value = output_sorted[i] 12 index = np.where(output == value) 13 for j in range(len(index)): 14 if (i + j) >= 5: 15 break 16 if value > 0: 17 topi = '{}: {}\n'.format(index[j], value) 18 else: 19 topi = '-1: 0.0\n' 20 top5_str += topi 21 print(top5_str) 22 23 24if __name__ == '__main__': 25 26 # Create RKNN object 27 rknn = RKNN(verbose=True) 28 29 # Pre-process config 30 print('--> Config model') 31 rknn.config(mean_values=[128, 128, 128], std_values=[128, 128, 128], 32 quantized_method='layer', quantized_algorithm='mmse') 33 print('done') 34 35 # Load model 36 print('--> Loading model') 37 ret = rknn.load_tensorflow(tf_pb='mobilenet_v1.pb', 38 inputs=['input'], 39 input_size_list=[[1, 224, 224, 3]], 40 outputs=['MobilenetV1/Logits/SpatialSqueeze']) 41 if ret != 0: 42 print('Load model failed!') 43 exit(ret) 44 print('done') 45 46 # Build model 47 print('--> Building model') 48 ret = rknn.build(do_quantization=True, dataset='./dataset.txt') 49 if ret != 0: 50 print('Build model failed!') 51 exit(ret) 52 print('done') 53 54 # Accuracy analysis 55 print('--> Accuracy analysis') 56 ret = rknn.accuracy_analysis(inputs=['dog_224x224.jpg'], output_dir=None) 57 if ret != 0: 58 print('Accuracy analysis failed!') 59 exit(ret) 60 print('done') 61 62 f = open('./snapshot/error_analysis.txt') 63 lines = f.readlines() 64 cos = lines[-1].split()[2] 65 if float(cos) >= 0.963: 66 print('cos = {}, mmse work!'.format(cos)) 67 else: 68 print('cos = {} < 0.963, mmse abnormal!'.format(cos)) 69 f.close() 70 71 # Set inputs 72 img = cv2.imread('./dog_224x224.jpg') 73 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 74 img = np.expand_dims(img, 0) 75 76 # Init runtime environment 77 print('--> Init runtime environment') 78 ret = rknn.init_runtime() 79 if ret != 0: 80 print('Init runtime environment failed!') 81 exit(ret) 82 print('done') 83 84 # Inference 85 print('--> Running model') 86 outputs = rknn.inference(inputs=[img]) 87 np.save('./functions_mmse_0.npy', outputs[0]) 88 show_outputs(outputs) 89 print('done') 90 91 rknn.release() 92