263 lines
8.1 KiB
Python
263 lines
8.1 KiB
Python
#!/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@<rk3588_ip>:/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())
|