#!/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()