yolov8检测模型pt转rknn

前言：

步骤是pt转onnx再转rknn。以下有些要点，请看代码备注。

一、pt2onnx.py

import argparse
import cv2
import numpy as np
import onnxruntime as ort
import torch
from ultralytics import YOLO

model = YOLO("runs/detect/train3/weights/best.pt")
# 导出为 ONNX 格式，并指定保存路径
output_path = "runs/detect/train3/weights"  # 自定义路径
 # 指定保存目录  特别注意点：opset = 12
model.export(format="onnx", save_dir="custom_directory",opset = 12) 

二、onnx转rknn

onnx2rknn.py

注意点

1. 准备校准图片
2. 注意写好对应的相机型号
3. 此转换过程是在pc端进行的

from rknn.api import RKNN
import os


# ===================== 配置（v2.3.2 严格兼容版） =====================
ONNX_MODEL = "./runs/detect/train3/weights/best.onnx"    # 新导出的标准ONNX
RKNN_MODEL = "./runs/detect/train3/weights/obstacle_det.rknn"
TARGET_PLATFORM = "rk3588"   # 根据你使用的相机型号进行替换
CALIB_DIR = "./calib_images/"  # INT8校准图目录（10-50张同分布图片）
CALIB_LIST = "./calib_list.txt"  # 🔧 必须是.txt文件路径，不能是列表

# ===================== 自动生成校准数据集（关键修复） =====================
def get_calib_dataset(calib_dir, CALIB_LIST):
    # 自动遍历文件夹，获取所有图片路径
    img_ext = [".jpg", ".jpeg", ".png", ".bmp"]
    calib_data = []
    for root, _, files in os.walk(calib_dir):
        for file in files:
            if any(file.lower().endswith(ext) for ext in img_ext):
                calib_data.append(os.path.join(root, file))
    if len(calib_dataset) == 0:
        print("❌ calib_images目录下无图片，请放入10-50张校准图")
        exit(1)
    # 写入列表文件
    with open(OUTPUT_LIST, "w") as f:
        for path in calib_dataset:
            f.write(path + "\n")
    return CALIB_LIST
    # return calib_data



# ===================== 初始化RKNN =====================
rknn = RKNN(verbose=True)

# mean_values长度匹配3通道，删除input_size_list等手动指定参数
print("=== 配置RKNN模型（RK3588 v2.3.2兼容） ===")
rknn.config(
    mean_values=[[0, 0, 0]],
    std_values=[[255, 255, 255]],
    target_platform=TARGET_PLATFORM,
    quantized_dtype="w8a8",  # INT8量化
    optimization_level=3
)

# 删除inputs/outputs/input_size_list，让RKNN自动读取ONNX维度
print("=== 加载ONNX模型 ===")
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
    print("❌ 加载ONNX失败")
    exit(ret)

# 自动生成校准数据集，传入build函数
print("=== 加载校准图 ===")
CALIB_LIST = get_calib_dataset(CALIB_DIR, CALIB_LIST)

# ===================== 构建INT8 RKNN模型 =====================
print("=== 构建INT8 RKNN模型 ===")
ret = rknn.build(
    do_quantization=True,
    dataset=CALIB_LIST  # 校准图目录
)
if ret != 0:
    print("❌ 构建RKNN失败")
    exit(ret)

# ===================== 导出RKNN模型 =====================
print("=== 导出RKNN模型 ===")
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    print("❌ 导出RKNN失败")
    exit(ret)

print(f"\n🎉 RK3588专用INT8 RKNN模型生成完成：{RKNN_MODEL}")
rknn.release()

三、最终验证步骤

1. 放图：把 10~50 张图放进 calib_images/
2. 生成列表：运行上面的生成脚本
3. 转换：运行 python onnx2rknn.py
4. 结果：生成 yolov8s_rk3588_int8.rknn