#!/usr/bin/env python3 """ 将 YOLOv8 ONNX 模型转换为 RKNN 格式 适用于 RK3588 / RK3568 / RK3576 等平台 环境要求: - Ubuntu x86_64 / Docker - Python 3.8 / 3.9 / 3.10 / 3.11 - rknn-toolkit2 (pip install rknn-toolkit2==2.2.0) 使用方法: # FP16 模式（推荐，速度快精度高） python 04_convert_rknn.py best.onnx -o shoe_detector.rknn -t rk3588 # INT8 量化（模型更小，需要校准数据集） python 04_convert_rknn.py best.onnx -o shoe_detector.rknn -t rk3588 -q -d dataset.txt 支持的 target_platform: - rk3588 / rk3588s - rk3568 / rk3566 - rk3576 - rv1106 / rv1103 / rv1103b - rv1126 """ import argparse import os import sys from pathlib import Path def check_environment(): """检查运行环境""" try: from rknn.api import RKNN print("✓ RKNN Toolkit2 已安装") return True except ImportError: print("✗ 错误: 未安装 RKNN Toolkit2") print("\n请安装:") print(" pip install rknn-toolkit2==2.2.0") print("\n或从源码安装:") print(" https://github.com/airockchip/rknn-toolkit2") return False def create_sample_dataset(onnx_path: str, output_path: str = "dataset.txt", num_samples: int = 20): """ 创建示例量化校准数据集 用于 INT8 量化时提供校准图片路径 """ print(f"\n创建示例校准数据集: {output_path}") print("注意: 请用实际图片替换这些示例路径") sample_content = f"""# RKNN INT8 量化校准数据集 # 每行一个图片路径，建议使用 20-100 张典型场景图片 # 图片格式: JPG, PNG, BMP 等 # 示例路径（请替换为实际路径）: # /path/to/train/images/img001.jpg # /path/to/train/images/img002.jpg # /path/to/valid/images/img001.jpg # 提示: # 1. 图片应与实际部署场景相似 # 2. 包含各种光照、角度、背景的样本 # 3. 建议 20-100 张，越多越慢但可能更准 """ with open(output_path, 'w') as f: f.write(sample_content) print(f"✓ 示例数据集已创建: {output_path}") print(" 请编辑此文件，添加实际的图片路径") return output_path def convert_onnx_to_rknn( onnx_path: str, output_path: str = None, target_platform: str = "rk3588", do_quantization: bool = False, dataset_path: str = None, verbose: bool = True ): """ 转换 ONNX 模型到 RKNN Args: onnx_path: ONNX 模型文件路径 output_path: 输出 RKNN 文件路径，默认与 ONNX 同名 target_platform: 目标平台，默认 rk3588 do_quantization: 是否启用 INT8 量化 dataset_path: 量化校准数据集路径（txt 文件，每行一张图片路径） verbose: 是否打印详细信息 """ if output_path is None: output_path = onnx_path.replace(".onnx", ".rknn") # 确保输出目录存在 output_dir = os.path.dirname(output_path) if output_dir and not os.path.exists(output_dir): os.makedirs(output_dir) print("="*70) print(f"ONNX 转 RKNN") print("="*70) print(f"输入: {onnx_path}") print(f"输出: {output_path}") print(f"目标: {target_platform}") print(f"量化: {'INT8' if do_quantization else 'FP16 (无量化)'}") if do_quantization: print(f"校准: {dataset_path}") print("="*70) # 检查输入文件 if not os.path.exists(onnx_path): print(f"\n✗ 错误: 找不到 ONNX 文件: {onnx_path}") return False # 检查数据集（如果需要量化） if do_quantization: if dataset_path is None: print("\n✗ 错误: INT8 量化需要提供校准数据集") print(" 使用 --dataset 指定数据集文件路径") print(" 或运行 --create-dataset 创建示例") return False if not os.path.exists(dataset_path): print(f"\n✗ 错误: 找不到数据集文件: {dataset_path}") return False from rknn.api import RKNN # 创建 RKNN 对象 rknn = RKNN(verbose=verbose) try: # 配置模型 print("\n[1/4] 配置模型...") rknn.config( mean_values=[[0, 0, 0]], # YOLOv8 使用 0-255 输入 std_values=[[255, 255, 255]], # 归一化到 0-1 target_platform=target_platform ) print(" ✓ 完成") # 加载 ONNX print("\n[2/4] 加载 ONNX 模型...") ret = rknn.load_onnx(model=onnx_path) if ret != 0: print(" ✗ 加载失败!") return False print(" ✓ 完成") # 构建模型 print("\n[3/4] 构建 RKNN 模型...") if do_quantization: print(f" 使用 INT8 量化，校准数据集: {dataset_path}") ret = rknn.build(do_quantization=True, dataset=dataset_path) else: print(" 使用 FP16 模式（无量化）") ret = rknn.build(do_quantization=False) if ret != 0: print(" ✗ 构建失败!") return False print(" ✓ 完成") # 导出 RKNN print("\n[4/4] 导出 RKNN 模型...") ret = rknn.export_rknn(output_path) if ret != 0: print(" ✗ 导出失败!") return False print(" ✓ 完成") finally: rknn.release() # 验证输出 if os.path.exists(output_path): size_mb = os.path.getsize(output_path) / (1024 * 1024) print("\n" + "="*70) print(f"✓ 转换成功!") print(f" 输出文件: {output_path}") print(f" 文件大小: {size_mb:.2f} MB") print("="*70) print("\n下一步:") print(f" 1. 复制到 RK3588:") print(f" scp {output_path} orangepi@:/home/orangepi/apps/OrangePi3588Media/models/") print(f" 2. 更新配置文件中的模型路径") return True else: print("\n✗ 错误: 输出文件未生成") return False def main(): parser = argparse.ArgumentParser( description="将 YOLOv8 ONNX 模型转换为 RKNN", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: # FP16 模式（推荐） python 04_convert_rknn.py best.onnx -o shoe_detector.rknn # 指定目标平台 python 04_convert_rknn.py best.onnx -t rk3568 # INT8 量化 python 04_convert_rknn.py best.onnx -q -d dataset.txt # 创建示例校准数据集 python 04_convert_rknn.py --create-dataset """ ) parser.add_argument("onnx", nargs="?", help="ONNX 模型文件路径") parser.add_argument("-o", "--output", help="输出 RKNN 文件路径") parser.add_argument("-t", "--target", default="rk3588", help="目标平台 (默认: rk3588)") parser.add_argument("-q", "--quantize", action="store_true", help="启用 INT8 量化") parser.add_argument("-d", "--dataset", help="量化校准数据集路径 (txt 文件)") parser.add_argument("--create-dataset", action="store_true", help="创建示例校准数据集并退出") parser.add_argument("-v", "--verbose", action="store_true", default=True, help="显示详细信息") args = parser.parse_args() # 创建示例数据集 if args.create_dataset: create_sample_dataset("best.onnx", "dataset.txt") return 0 # 检查参数 if args.onnx is None: parser.print_help() print("\n错误: 请提供 ONNX 文件路径") return 1 # 检查环境 if not check_environment(): return 1 # 执行转换 success = convert_onnx_to_rknn( onnx_path=args.onnx, output_path=args.output, target_platform=args.target, do_quantization=args.quantize, dataset_path=args.dataset, verbose=args.verbose ) return 0 if success else 1 if __name__ == "__main__": sys.exit(main())