YOLOV8-DET转ONNX和RKNN

news2024/11/24 0:05:40

1. 前言

2.环境配置

(1) RK3588开发板Python环境

(2) PC转onnx和rknn的环境

3.PT模型转onnx

4. ONNX模型转RKNN

6.测试结果

1. 前言

yolov8就不介绍了，详细的请见YOLOV8详细对比，本文章注重实际的使用，从拿到yolov8的pt检测模型，怎样转onnx和rknn的使用，配置好相应的环境，就可以开心的做yolov8的检测啦！！！

2.环境配置

(1) RK3588开发板Python环境

        numpy ==1.26.0
        onnxruntime == 1.16.0
        opencv-python== 4.8.0.76
        rknn-toolkit-lite2 ==1.5.2

(2) PC转onnx和rknn的环境

Python == 3.10

ultralytics ==8.0.147

rknn_toolkit2 == 1.5.2

torch ==2.0.1
torchvision==0.15.2

相关的库缺少什么就安装，很简单的

3.PT模型转onnx

根据自己模型路径设置model_path，输入的shape也自己设置，我这里设置的是640*640，这样就得到了onnx模型，按道理是可以做推理了

from ultralytics import YOLO

# pip install ultralytics

input_width = 640
input_height = 640
model_path = "./models/yolov8x.pt"


model = YOLO(model_path)
model.export(format="onnx", imgsz=[input_height, input_width], opset=12)

4. ONNX模型转RKNN

rknn用FP16，用INT8设置了，但是量化效果精度有误差，用混合量化或者增加量化数据集，都可以试试(我偷懒就不尝试了)

from rknn.api import RKNN


ONNX_MODEL = 'yolov8n-sim.onnx'
RKNN_MODEL = 'yolov8n-sim-fp16.rknn'
IMG_PATH = './bus.jpg'
DATASET = './dataset.txt'


if __name__ == '__main__':
    # Create RKNN object
    rknn = RKNN()
    
    if not os.path.exists(ONNX_MODEL):
        print('model not exist')
        exit(-1)
    
    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[1, 1, 1]], target_platform='rk3588')
    print('done')
    
    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load yolov8 failed!')
        exit(ret)
    print('done')
    
    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=False, dataset=DATASET)
    if ret != 0:
        print('Build yolov8 failed!')
        exit(ret)
    print('done')
    
    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export yolov8rknn failed!')
        exit(ret)
    print('done')
    
    # Ret = rknn.accuracy_analysis(inputs=['./bus.jpg'])
    # if ret != 0:
    #     print('Accuracy analysis failed!')
    # exit(ret)
    # print('done')

5. RK3588推理

代码使用的是https://github.com/ibaiGorordo/ONNX-YOLOv8-Object-DetectionibaiGorord博主的基础代码，在YOLOv8中集成了rknn模型的推理

import cv2
from yolov8 import YOLOv8


# Initialize yolov8 object detector
# model_path = "./models/yolov8n-sim.onnx"
model_path = "./models/yolov8n-sim-fp16.rknn"
yolov8_detector = YOLOv8(model_path, conf_thres=0.3, iou_thres=0.5)
# Read image
img = cv2.imread('./bus.jpg')
# Detect Objects
boxes, scores, class_ids = yolov8_detector(img)
# Draw detections
combined_img = yolov8_detector.draw_detections(img)
cv2.imwrite("./detected_objects.png", combined_img)