图像识别模型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 argparse

class onnx2rknn(object):
    def __init__(self,opt):
        self.opt = opt 
        self.input_w,self.input_h = opt.input_size 
        self.QUANTIZE_ON = True
        self.get_class_dict()
        
    def transformer(self):
        # Create RKNN object
        rknn = RKNN(verbose=True)

        # pre-process config
        print('--> Config model')
        rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]],target_platform='rk3588')
        print('done')

        # Load ONNX model
        print('--> Loading model')
        ret = rknn.load_onnx(model=self.opt.onnx_model)
        if ret != 0:
            print('Load model failed!')
            exit(ret)
        print('done')

        # Build model
        print('--> Building model')
        ret = rknn.build(do_quantization=False)
        if ret != 0:
            print('Build model failed!')
            exit(ret)
        print('done')

        # Export RKNN model
        print('--> 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()
        return


def 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的页面下载 

2.rk3588s可用的rknn模型(yolov5 640*640)

yolov5 rknn模型，和onnx2rknn python 源码