304 lines
11 KiB
Python
304 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
PaddleOCR ONNX to RKNN Model Converter
|
|
将PaddleOCR的文本检测和文本识别ONNX模型转换为RKNN模型
|
|
"""
|
|
|
|
import os
|
|
import numpy as np
|
|
from rknn.api import RKNN
|
|
import argparse
|
|
|
|
class PaddleOCRConverter:
|
|
def __init__(self):
|
|
self.rknn = RKNN(verbose=True)
|
|
|
|
def get_quantized_dtype(self, platform, quantize):
|
|
"""
|
|
根据平台和是否量化选择合适的quantized_dtype
|
|
"""
|
|
if platform.lower() in ['rk3588']:
|
|
# RK3588支持的量化类型
|
|
return 'w8a8' if quantize else None
|
|
else:
|
|
# 其他平台 (rk3566, rk3568等)
|
|
return 'w8a8' if quantize else 'w8a16'
|
|
|
|
def convert_detection_model(self, onnx_path, rknn_path, platform='rk3588', quantize=False, dataset_path=None):
|
|
"""
|
|
转换文本检测模型
|
|
Args:
|
|
onnx_path: 输入的ONNX模型路径
|
|
rknn_path: 输出的RKNN模型路径
|
|
platform: 目标平台 (rk3588, rk3566, rk3568等)
|
|
quantize: 是否进行量化
|
|
dataset_path: 量化数据集路径
|
|
"""
|
|
print(f"开始转换文本检测模型: {onnx_path}")
|
|
|
|
# 获取适合的quantized_dtype
|
|
quantized_dtype = self.get_quantized_dtype(platform, quantize)
|
|
|
|
# 配置RKNN
|
|
config_params = {
|
|
# 'mean_values': [[123.675, 116.28, 103.53]], # PaddleOCR标准化参数
|
|
# 'std_values': [[58.395, 57.12, 57.375]],
|
|
'target_platform': platform
|
|
}
|
|
|
|
# 只有在需要时才添加quantized_dtype参数
|
|
if quantized_dtype is not None:
|
|
config_params['quantized_dtype'] = quantized_dtype
|
|
|
|
self.rknn.config(**config_params)
|
|
|
|
# 加载ONNX模型
|
|
print("加载ONNX模型...")
|
|
ret = self.rknn.load_onnx(model=onnx_path)
|
|
if ret != 0:
|
|
print(f"加载ONNX模型失败: {ret}")
|
|
return False
|
|
|
|
# 构建模型
|
|
print("构建RKNN模型...")
|
|
ret = self.rknn.build(do_quantization=quantize, dataset=dataset_path)
|
|
if ret != 0:
|
|
print(f"构建RKNN模型失败: {ret}")
|
|
return False
|
|
|
|
# 导出模型
|
|
print(f"导出RKNN模型到: {rknn_path}")
|
|
ret = self.rknn.export_rknn(rknn_path)
|
|
if ret != 0:
|
|
print(f"导出RKNN模型失败: {ret}")
|
|
return False
|
|
|
|
print("文本检测模型转换完成!")
|
|
return True
|
|
|
|
def convert_recognition_model(self, onnx_path, rknn_path, platform='rk3588', quantize=False, dataset_path=None):
|
|
"""
|
|
转换文本识别模型
|
|
Args:
|
|
onnx_path: 输入的ONNX模型路径
|
|
rknn_path: 输出的RKNN模型路径
|
|
platform: 目标平台
|
|
quantize: 是否进行量化
|
|
dataset_path: 量化数据集路径
|
|
"""
|
|
print(f"开始转换文本识别模型: {onnx_path}")
|
|
|
|
# 重新初始化RKNN实例
|
|
self.rknn.release()
|
|
self.rknn = RKNN(verbose=True)
|
|
|
|
# 获取适合的quantized_dtype
|
|
quantized_dtype = self.get_quantized_dtype(platform, quantize)
|
|
|
|
# 配置RKNN (文本识别模型的预处理参数)
|
|
config_params = {
|
|
# 'mean_values': [[127.5, 127.5, 127.5]], # 文本识别模型标准化参数
|
|
# 'std_values': [[127.5, 127.5, 127.5]],
|
|
'target_platform': platform
|
|
}
|
|
|
|
# 只有在需要时才添加quantized_dtype参数
|
|
if quantized_dtype is not None:
|
|
config_params['quantized_dtype'] = quantized_dtype
|
|
|
|
self.rknn.config(**config_params)
|
|
|
|
# 加载ONNX模型
|
|
print("加载ONNX模型...")
|
|
ret = self.rknn.load_onnx(model=onnx_path)
|
|
if ret != 0:
|
|
print(f"加载ONNX模型失败: {ret}")
|
|
return False
|
|
|
|
# 构建模型
|
|
print("构建RKNN模型...")
|
|
ret = self.rknn.build(do_quantization=quantize, dataset=dataset_path)
|
|
if ret != 0:
|
|
print(f"构建RKNN模型失败: {ret}")
|
|
return False
|
|
|
|
# 导出模型
|
|
print(f"导出RKNN模型到: {rknn_path}")
|
|
ret = self.rknn.export_rknn(rknn_path)
|
|
if ret != 0:
|
|
print(f"导出RKNN模型失败: {ret}")
|
|
return False
|
|
|
|
print("文本识别模型转换完成!")
|
|
return True
|
|
|
|
def generate_quantization_dataset(self, input_shape, num_samples=100, output_path="./dataset.txt"):
|
|
"""
|
|
生成量化数据集
|
|
Args:
|
|
input_shape: 输入形状 (N, C, H, W)
|
|
num_samples: 样本数量
|
|
output_path: 输出路径
|
|
"""
|
|
print(f"生成量化数据集: {output_path}")
|
|
|
|
# 创建数据集目录
|
|
dataset_dir = os.path.dirname(output_path)
|
|
if dataset_dir and not os.path.exists(dataset_dir):
|
|
os.makedirs(dataset_dir)
|
|
|
|
# 生成随机数据并保存
|
|
dataset_files = []
|
|
for i in range(num_samples):
|
|
# 生成随机数据 (0-255)
|
|
data = np.random.randint(0, 256, input_shape, dtype=np.uint8)
|
|
|
|
# 保存为npy文件
|
|
npy_path = f"{dataset_dir}/sample_{i}.npy"
|
|
np.save(npy_path, data)
|
|
dataset_files.append(npy_path)
|
|
|
|
# 创建数据集列表文件
|
|
with open(output_path, 'w') as f:
|
|
for file_path in dataset_files:
|
|
f.write(f"{file_path}\n")
|
|
|
|
print(f"生成了 {num_samples} 个样本的量化数据集")
|
|
return output_path
|
|
|
|
def test_model(self, rknn_path, input_shape):
|
|
"""
|
|
测试转换后的RKNN模型
|
|
Args:
|
|
rknn_path: RKNN模型路径
|
|
input_shape: 输入形状
|
|
"""
|
|
print(f"测试RKNN模型: {rknn_path}")
|
|
|
|
# 重新初始化RKNN实例用于推理
|
|
test_rknn = RKNN(verbose=False)
|
|
|
|
# 加载RKNN模型
|
|
ret = test_rknn.load_rknn(rknn_path)
|
|
if ret != 0:
|
|
print(f"加载RKNN模型失败: {ret}")
|
|
return False
|
|
|
|
# 初始化运行环境
|
|
ret = test_rknn.init_runtime()
|
|
if ret != 0:
|
|
print(f"初始化运行环境失败: {ret}")
|
|
return False
|
|
|
|
# 生成测试数据
|
|
test_data = np.random.randint(0, 256, input_shape, dtype=np.uint8)
|
|
|
|
# 进行推理
|
|
outputs = test_rknn.inference(inputs=[test_data])
|
|
|
|
print(f"推理成功! 输出形状: {[output.shape for output in outputs]}")
|
|
|
|
# 释放资源
|
|
test_rknn.release()
|
|
return True
|
|
|
|
def release(self):
|
|
"""释放RKNN资源"""
|
|
self.rknn.release()
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='将PaddleOCR ONNX模型转换为RKNN模型')
|
|
parser.add_argument('--det_onnx', type=str, required=True, help='文本检测ONNX模型路径')
|
|
parser.add_argument('--rec_onnx', type=str, required=True, help='文本识别ONNX模型路径')
|
|
parser.add_argument('--output_dir', type=str, default='./rknn_models', help='RKNN模型输出目录')
|
|
parser.add_argument('--platform', type=str, default='rk3588', help='目标平台')
|
|
parser.add_argument('--quantize', action='store_true', help='是否进行量化')
|
|
parser.add_argument('--test', action='store_true', help='是否测试转换后的模型')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# 创建输出目录
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
# 初始化转换器
|
|
converter = PaddleOCRConverter()
|
|
|
|
try:
|
|
# 定义输入形状
|
|
det_input_shape = (1, 3, 640, 640) # 文本检测模型输入形状
|
|
rec_input_shape = (1, 3, 48, 320) # 文本识别模型输入形状
|
|
|
|
# 生成量化数据集（如果需要量化）
|
|
det_dataset_path = None
|
|
rec_dataset_path = None
|
|
|
|
if args.quantize:
|
|
print("生成量化数据集...")
|
|
det_dataset_path = converter.generate_quantization_dataset(
|
|
det_input_shape, num_samples=100,
|
|
output_path=os.path.join(args.output_dir, "det_dataset.txt")
|
|
)
|
|
rec_dataset_path = converter.generate_quantization_dataset(
|
|
rec_input_shape, num_samples=100,
|
|
output_path=os.path.join(args.output_dir, "rec_dataset.txt")
|
|
)
|
|
|
|
# 转换文本检测模型
|
|
det_rknn_path = os.path.join(args.output_dir, "text_detection.rknn")
|
|
success = converter.convert_detection_model(
|
|
args.det_onnx, det_rknn_path, args.platform, args.quantize, det_dataset_path
|
|
)
|
|
|
|
if not success:
|
|
print("文本检测模型转换失败!")
|
|
return
|
|
|
|
# 转换文本识别模型
|
|
rec_rknn_path = os.path.join(args.output_dir, "text_recognition.rknn")
|
|
success = converter.convert_recognition_model(
|
|
args.rec_onnx, rec_rknn_path, args.platform, args.quantize, rec_dataset_path
|
|
)
|
|
|
|
if not success:
|
|
print("文本识别模型转换失败!")
|
|
return
|
|
|
|
# 测试模型（如果指定）
|
|
if args.test:
|
|
print("\n开始测试转换后的模型...")
|
|
converter.test_model(det_rknn_path, det_input_shape)
|
|
converter.test_model(rec_rknn_path, rec_input_shape)
|
|
|
|
print(f"\n所有模型转换完成!")
|
|
print(f"文本检测模型: {det_rknn_path}")
|
|
print(f"文本识别模型: {rec_rknn_path}")
|
|
|
|
finally:
|
|
converter.release()
|
|
|
|
if __name__ == "__main__":
|
|
# 如果直接运行脚本，提供示例用法
|
|
'''
|
|
python 016新的导出rknn模型脚本.py --det_onnx /home/admin-root/haotian/康达瑞贝斯机器狗/det_mobile_14_shape.onnx --rec_onnx /home/admin-root/haotian/康达瑞贝斯机 器狗/rec_mobile_14_shape.onnx --platform rk3588 --output_dir ./models
|
|
'''
|
|
if len(os.sys.argv) == 1:
|
|
print("PaddleOCR ONNX to RKNN Model Converter")
|
|
print("\n使用示例:")
|
|
print("python paddleocr_converter.py \\")
|
|
print(" --det_onnx ./text_detection.onnx \\")
|
|
print(" --rec_onnx ./text_recognition.onnx \\")
|
|
print(" --output_dir ./rknn_models \\")
|
|
print(" --platform rk3588 \\")
|
|
print(" --quantize \\")
|
|
print(" --test")
|
|
print("\n参数说明:")
|
|
print(" --det_onnx: 文本检测ONNX模型路径")
|
|
print(" --rec_onnx: 文本识别ONNX模型路径")
|
|
print(" --output_dir: RKNN模型输出目录")
|
|
print(" --platform: 目标平台 (rk3588, rk3566, rk3568等)")
|
|
print(" --quantize: 是否进行量化")
|
|
print(" --test: 是否测试转换后的模型")
|
|
else:
|
|
main() |