1import os 2import urllib 3import traceback 4import time 5import sys 6import numpy as np 7import cv2 8from rknn.api import RKNN 9 10ONNX_MODEL = 'resnet50v2.onnx' 11RKNN_MODEL = 'resnet50v2.rknn' 12 13 14def show_outputs(outputs): 15 output = outputs[0][0] 16 output_sorted = sorted(output, reverse=True) 17 top5_str = 'resnet50v2\n-----TOP 5-----\n' 18 for i in range(5): 19 value = output_sorted[i] 20 index = np.where(output == value) 21 for j in range(len(index)): 22 if (i + j) >= 5: 23 break 24 if value > 0: 25 topi = '{}: {}\n'.format(index[j], value) 26 else: 27 topi = '-1: 0.0\n' 28 top5_str += topi 29 print(top5_str) 30 31 32def readable_speed(speed): 33 speed_bytes = float(speed) 34 speed_kbytes = speed_bytes / 1024 35 if speed_kbytes > 1024: 36 speed_mbytes = speed_kbytes / 1024 37 if speed_mbytes > 1024: 38 speed_gbytes = speed_mbytes / 1024 39 return "{:.2f} GB/s".format(speed_gbytes) 40 else: 41 return "{:.2f} MB/s".format(speed_mbytes) 42 else: 43 return "{:.2f} KB/s".format(speed_kbytes) 44 45 46def show_progress(blocknum, blocksize, totalsize): 47 speed = (blocknum * blocksize) / (time.time() - start_time) 48 speed_str = " Speed: {}".format(readable_speed(speed)) 49 recv_size = blocknum * blocksize 50 51 f = sys.stdout 52 progress = (recv_size / totalsize) 53 progress_str = "{:.2f}%".format(progress * 100) 54 n = round(progress * 50) 55 s = ('#' * n).ljust(50, '-') 56 f.write(progress_str.ljust(8, ' ') + '[' + s + ']' + speed_str) 57 f.flush() 58 f.write('\r\n') 59 60 61if __name__ == '__main__': 62 63 # Create RKNN object 64 rknn = RKNN(verbose=True) 65 66 # If resnet50v2 does not exist, download it. 67 # Download address: 68 # https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet50v2/resnet50v2.onnx 69 if not os.path.exists(ONNX_MODEL): 70 print('--> Download {}'.format(ONNX_MODEL)) 71 url = 'https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet50v2/resnet50v2.onnx' 72 download_file = ONNX_MODEL 73 try: 74 start_time = time.time() 75 urllib.request.urlretrieve(url, download_file, show_progress) 76 except: 77 print('Download {} failed.'.format(download_file)) 78 print(traceback.format_exc()) 79 exit(-1) 80 print('done') 81 82 # pre-process config 83 print('--> config model') 84 rknn.config(mean_values=[123.675, 116.28, 103.53], std_values=[58.82, 58.82, 58.82]) 85 print('done') 86 87 # Load model 88 print('--> Loading model') 89 ret = rknn.load_onnx(model=ONNX_MODEL) 90 if ret != 0: 91 print('Load model failed!') 92 exit(ret) 93 print('done') 94 95 # Build model 96 print('--> Building model') 97 ret = rknn.build(do_quantization=True, dataset='./dataset.txt') 98 if ret != 0: 99 print('Build model failed!') 100 exit(ret) 101 print('done') 102 103 # Export rknn model 104 print('--> Export rknn model') 105 ret = rknn.export_rknn(RKNN_MODEL) 106 if ret != 0: 107 print('Export rknn model failed!') 108 exit(ret) 109 print('done') 110 111 # Set inputs 112 img = cv2.imread('./dog_224x224.jpg') 113 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 114 115 # Init runtime environment 116 print('--> Init runtime environment') 117 ret = rknn.init_runtime() 118 if ret != 0: 119 print('Init runtime environment failed!') 120 exit(ret) 121 print('done') 122 123 # Inference 124 print('--> Running model') 125 outputs = rknn.inference(inputs=[img]) 126 np.save('./onnx_resnet50v2_0.npy', outputs[0]) 127 x = outputs[0] 128 output = np.exp(x)/np.sum(np.exp(x)) 129 outputs = [output] 130 show_outputs(outputs) 131 print('done') 132 133 rknn.release() 134