onnx篇---onnx转换成rknn

多张图预测

# -- coding: utf-8 --
import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from matplotlib import pyplot as plt
from rknn.api import RKNN
 

def show_outputs(outputs):
    # output = outputs[0][0]
    output = outputs[0].reshape(-1)
    print(type(output))
    print(output.shape)
    output_sorted = sorted(output, reverse=True)
    top5_str = 'yolov4\n-----TOP 5-----\n'
    for i in range(5):
        value = output_sorted[i]
        index = np.where(output == value)
        for j in range(len(index)):
            if (i + j) >= 5:
                break
            if value > 0:
                topi = '{}: {}\n'.format(index[j], value)
            else:
                topi = '-1: 0.0\n'
            top5_str += topi
    print(top5_str)


if __name__ == '__main__':
    ONNX_MODEL = './weights/Yolo_bin.onnx'
    RKNN_MODEL_PATH = './weights/Yolo_bin_quantization_mul.rknn'
    im_file = './data/0617_Bin_026.jpg'

    # Create RKNN object
    rknn = RKNN(verbose=True, verbose_file='./inference.log')

    NEED_BUILD_MODEL = True

    if NEED_BUILD_MODEL:

        start_time = time.time()
        # pre-process config
        print('--> Config model')
        rknn.config(batch_size=8, reorder_channel='2 1 0', mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]])
        print('done')

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

        # Build model
        print('--> Building model')
        ret = rknn.build(do_quantization=True, dataset='./data/dataset1.txt')
        if ret != 0:
            print('Build yolov4 failed!')
            exit(ret)
        print('done')

        # Export RKNN model
        print('--> Export RKNN model')
        ret = rknn.export_rknn(RKNN_MODEL_PATH)
        if ret != 0:
            print('Export yolov4.rknn failed!')
            exit(ret)
        print('done')

    else:
        # Direct load rknn model
        print('Loading RKNN model')
        ret = rknn.load_rknn(RKNN_MODEL_PATH)
        if ret != 0:
            print('Load RKNN model failed.')
            exit(ret)
        print('done')

    # init runtime environment
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')

    # Set inputs
    img = cv2.imread(im_file)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # Inference
    print('--> Running model')
    outputs = rknn.inference(inputs=[img], data_format="nchw")
    print("output", outputs)
    print('done')

    x = outputs[0]
    output = np.exp(x) / np.sum(np.exp(x))
    outputs = [output]
    show_outputs(outputs)
    print('done')

    # perf
    print('--> Begin evaluate model performance')
    perf_results = rknn.eval_perf(inputs=[img])
    print(perf_results)
    print('done')

    rknn.release()

rknn.build中的dataset.txt设置

单张图的预测

只需要改dataset.txt