|
沙发
楼主 |
发表于 2021-3-3 23:41:34
|
只看该作者
补充代码如下:
- import numpy as np
- import cv2
- from rknn.api import RKNN
- # 设置模型存放路径
- ONNX_MODEL = 'classify.onnx'
- RKNN_MODEL = 'classify.rknn'
- # 显示TOP3结果
- def show_outputs(outputs):
- output = outputs[0][0]
- output_sorted = sorted(output, reverse=True)
- top3_str = 'mobilenetv2\n-----TOP 3-----\n'
- for i in range(3):
- value = output_sorted[i]
- index = np.where(output == value)
- for j in range(len(index)):
- if (i + j) >= 3:
- break
- if value > 0:
- topi = '{}: {}\n'.format(index[j], value)
- else:
- topi = '-1: 0.0\n'
- top3_str += topi
- print(top3_str)
- if __name__ == '__main__':
- #############################################################################
- #
- # 第一部分:通过onnx模型转换、量化为RKNN模型
- #
- #############################################################################
- # 1. 创建RKNN对象
- rknn = RKNN()
- # 2. 模型前处理配置
- print('--> config model')
- rknn.config(channel_mean_value='116.28 116.28 116.28 57.12', reorder_channel='0 1 2')
- rknn.config()
- print('done')
- # 3. 加载onnx模型文件
- print('--> Loading model')
- ret = rknn.load_onnx(model=ONNX_MODEL)
- if ret != 0:
- print('Load mobilenetv2 failed!')
- exit(ret)
- print('done')
- # 4. 创建模型文件(量化)
- print('--> Building model')
- ret = rknn.build(do_quantization=False,dataset='./eval_dir_180x320.txt')
- if ret != 0:
- print('Build mobilenetv2 failed!')
- exit(ret)
- print('done')
- # 5. 提取RKNN模型文件
- print('--> Export RKNN model')
- ret = rknn.export_rknn(RKNN_MODEL)
- if ret != 0:
- print('Export mobilenetv2.rknn failed!')
- exit(ret)
- print('done')
- #############################################################################
- #
- #第二部分:基于RKNN模型进行模型的前向推理
- #
- #############################################################################
-
- # 1. 读取图像文件&BGR2RGB
- img = cv2.imread('./1/5_361.jpg')
- img = cv2.resize(img,(320,180))
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- #img = cv2.normalize(img,img,alpha=0.017507,beta=2.035714)
- # 2. 初始化RKNN运行环境
- print('--> Init runtime environment')
- ret = rknn.init_runtime()
- if ret != 0:
- print('Init runtime environment failed')
- exit(ret)
- print('done')
- # 3. 前向推理
- print('--> Running model')
- outputs = rknn.inference(inputs=[img])
- x = outputs[0]
- output = np.exp(x)/np.sum(np.exp(x))
- outputs = [output]
- print(outputs)
- #show_outputs(outputs)
- print('done')
- rknn.release()
|
|