图像识别模型onnx2rknn转换程序（可用于鲁班猫4

图像识别模型onnx2rknn转换程序（可用于鲁班猫4 - rk3588s）

本文主要是介绍图像识别模型onnx2rknn转换程序（可用于鲁班猫4 - rk3588s），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

鲁班猫4中不含rk3588 yolov5的rknn模型库。在rknn-toolkit2中包含的那个.rknn是针对rk3568的。这里给出了将一个onnx模型转换为rknn模型的代码，并附上了可用于rk3588s的yolov5.rknn模型。

1.模型转换代码

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
import argparseclass onnx2rknn(object):def __init__(self,opt):self.opt = opt self.input_w,self.input_h = opt.input_size self.QUANTIZE_ON = Trueself.get_class_dict()def transformer(self):# Create RKNN objectrknn = RKNN(verbose=True)# pre-process configprint('--> Config model')rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]],target_platform='rk3588')print('done')# Load ONNX modelprint('--> Loading model')ret = rknn.load_onnx(model=self.opt.onnx_model)if ret != 0:print('Load model failed!')exit(ret)print('done')# Build modelprint('--> Building model')ret = rknn.build(do_quantization=False)if ret != 0:print('Build model failed!')exit(ret)print('done')# Export RKNN modelprint('--> Export rknn model')ret = rknn.export_rknn(self.opt.rnkk_model)if ret != 0:print('Export rknn model failed!')exit(ret)print('done')rknn.release()returndef set_option():parser = argparse.ArgumentParser()parser.add_argument('--input_size',type=list,default=[640,640]) #[288,32]parser.add_argument('--onnx_model',type=str,default='./yolov5s.onnx')parser.add_argument('--rnkk_model',type=str,default='./yolov5_rk3588.rknn')parser.add_argument('--class_file',type=str,default='./data/class.txt')parser.add_argument('--test_img',type=str,default='./test_imgs/test.jpg')opt = parser.parse_args()return opt if __name__ == '__main__':opt = set_option()onnx2rknn = onnx2rknn(opt)onnx2rknn.transformer()

注意yolo5s.onnx在鲁班猫4的toolkits里面。class.txt可以自行创建，也可以到yolov5的页面下载