From 1ad4413b9bc76bbd001780267d1b6cf3b62c8a4b Mon Sep 17 00:00:00 2001 From: haotian <2421912570@qq.com> Date: Tue, 26 Aug 2025 15:45:14 +0800 Subject: [PATCH] =?UTF-8?q?1.=E5=AE=9E=E7=8E=B0016=E6=96=B0=E7=9A=84?= =?UTF-8?q?=E5=AF=BC=E5=87=BArknn=E6=A8=A1=E5=9E=8B=E8=84=9A=E6=9C=AC\n2.0?= =?UTF-8?q?17=E5=AE=9E=E7=8E=B0=E4=BD=BF=E7=94=A8rknn=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E7=9A=84=E8=84=9A=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- 016新的导出rknn模型脚本.py | 304 ++++++++++++++++++ 017能用的PaddleOCR_rknn脚本.py | 559 +++++++++++++++++++++++++++++++++ 2 files changed, 863 insertions(+) create mode 100644 016新的导出rknn模型脚本.py create mode 100644 017能用的PaddleOCR_rknn脚本.py diff --git a/016新的导出rknn模型脚本.py b/016新的导出rknn模型脚本.py new file mode 100644 index 0000000..1838c50 --- /dev/null +++ b/016新的导出rknn模型脚本.py @@ -0,0 +1,304 @@ +#!/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() \ No newline at end of file diff --git a/017能用的PaddleOCR_rknn脚本.py b/017能用的PaddleOCR_rknn脚本.py new file mode 100644 index 0000000..95b0625 --- /dev/null +++ b/017能用的PaddleOCR_rknn脚本.py @@ -0,0 +1,559 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +RK3588 PaddleOCR RKNN推理程序 +使用转换后的RKNN模型在RK3588上进行OCR文本检测和识别 +""" + +import cv2 +import yaml +import numpy as np +import math +from rknn.api import RKNN +import argparse +import os +import time + +class RK3588OCR: + def __init__(self, det_model_path, rec_model_path): + """ + 初始化OCR推理器 + Args: + det_model_path: 文本检测RKNN模型路径 + rec_model_path: 文本识别RKNN模型路径 + """ + self.det_model_path = det_model_path + self.rec_model_path = rec_model_path + + # 初始化RKNN实例 + self.det_rknn = RKNN(verbose=False) + self.rec_rknn = RKNN(verbose=False) + + # 模型输入尺寸 + self.det_input_size = (640, 640) + self.rec_input_size = (320, 48) + + # 文本检测相关参数 + self.det_threshold = 0.3 + self.det_box_threshold = 0.6 + self.det_unclip_ratio = 1.5 + + # 加载模型 + self._load_models() + + self.character = self.get_dict() + + def get_dict(self, dict_path='/home/orangepi/Desktop/kangda_robotic_dog/机器狗后台服务/dict.yaml'): + """ + 加载字典 + """ + with open(dict_path, 'r', encoding='utf-8') as f: + dict_rec = yaml.safe_load(f) + return dict_rec.get('character_dict', []) + + def _load_models(self): + """加载RKNN模型""" + print("加载文本检测模型...") + ret = self.det_rknn.load_rknn(self.det_model_path) + if ret != 0: + raise Exception(f"加载检测模型失败: {ret}") + + ret = self.det_rknn.init_runtime(target='rk3588') + if ret != 0: + raise Exception(f"初始化检测模型运行环境失败: {ret}") + + print("加载文本识别模型...") + ret = self.rec_rknn.load_rknn(self.rec_model_path) + if ret != 0: + raise Exception(f"加载识别模型失败: {ret}") + + ret = self.rec_rknn.init_runtime(target='rk3588') + if ret != 0: + raise Exception(f"初始化识别模型运行环境失败: {ret}") + + print("模型加载完成!") + + def resize_norm_img_det(self, img, input_shape=(640, 640)): + """ + 检测模型的图像预处理 - 固定输入形状 [1, 3, 640, 640] + """ + h, w, _ = img.shape + target_h, target_w = input_shape + + # 计算缩放比例 - 保持宽高比 + ratio_h = target_h / h + ratio_w = target_w / w + ratio = min(ratio_h, ratio_w) + + # 计算缩放后的尺寸 + new_h = int(h * ratio) + new_w = int(w * ratio) + + # 调整图像大小 + resized_img = cv2.resize(img, (new_w, new_h)) + + # 创建目标尺寸的图像，用灰色填充 + padded_img = np.ones((target_h, target_w, 3), dtype=np.float32) * 114.0 # 直接用float32 + + # 计算居中位置 + top = (target_h - new_h) // 2 + left = (target_w - new_w) // 2 + + # 将缩放后的图像放到居中位置 + padded_img[top:top+new_h, left:left+new_w] = resized_img.astype(np.float32) + + # 归一化 + img = (padded_img / 255.0 - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array([0.229, 0.224, 0.225], dtype=np.float32) + img = img.transpose(2, 0, 1).astype(np.float32) + img = np.expand_dims(img, axis=0).astype(np.float32) + + return img, ratio, (top, left) + + def post_process_det(self, dt_boxes, ratio, padding_info, ori_shape): + """ + 检测结果后处理 - 适配固定输入形状 + """ + if dt_boxes is None: + return None + + ori_h, ori_w = ori_shape + top, left = padding_info + + # 将坐标从模型输出空间转换回原图空间 + dt_boxes[:, :, 0] = (dt_boxes[:, :, 0] - left) / ratio + dt_boxes[:, :, 1] = (dt_boxes[:, :, 1] - top) / ratio + + # 裁剪到原图范围内 + dt_boxes[:, :, 0] = np.clip(dt_boxes[:, :, 0], 0, ori_w) + dt_boxes[:, :, 1] = np.clip(dt_boxes[:, :, 1], 0, ori_h) + + return dt_boxes + + def boxes_from_bitmap(self, pred, bitmap, dest_width, dest_height, max_candidates=1000, box_thresh=0.6): + """ + 从位图中提取文本框 + """ + bitmap = bitmap.astype(np.uint8) + height, width = bitmap.shape + + # 查找轮廓 + contours, _ = cv2.findContours(bitmap, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + + num_contours = min(len(contours), max_candidates) + boxes = [] + scores = [] + + for i in range(num_contours): + contour = contours[i] + points, sside = self.get_mini_boxes(contour) + if sside < 5: + continue + + points = np.array(points) + score = self.box_score_fast(pred, points.reshape(-1, 2)) + if box_thresh > score: + continue + + # 扩展box + box = self.unclip(points, 1.5).reshape(-1, 1, 2) + box, sside = self.get_mini_boxes(box) + if sside < 5 + 2: + continue + + box = np.array(box) + box[:, 0] = np.clip(box[:, 0] / width * dest_width, 0, dest_width) + box[:, 1] = np.clip(box[:, 1] / height * dest_height, 0, dest_height) + + boxes.append(box.astype(np.int16)) + scores.append(score) + + return np.array(boxes), scores + + def get_mini_boxes(self, contour): + """获取最小外接矩形""" + bounding_box = cv2.minAreaRect(contour) + points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0]) + + index_1, index_2, index_3, index_4 = 0, 1, 2, 3 + if points[1][1] > points[0][1]: + index_1 = 0 + index_4 = 1 + else: + index_1 = 1 + index_4 = 0 + + if points[3][1] > points[2][1]: + index_2 = 2 + index_3 = 3 + else: + index_2 = 3 + index_3 = 2 + + box = [points[index_1], points[index_2], points[index_3], points[index_4]] + return box, min(bounding_box[1]) + + def box_score_fast(self, bitmap, _box): + """快速计算box得分""" + h, w = bitmap.shape[:2] + box = _box.copy() + xmin = np.clip(np.floor(box[:, 0].min()).astype(int), 0, w - 1) + xmax = np.clip(np.ceil(box[:, 0].max()).astype(int), 0, w - 1) + ymin = np.clip(np.floor(box[:, 1].min()).astype(int), 0, h - 1) + ymax = np.clip(np.ceil(box[:, 1].max()).astype(int), 0, h - 1) + + mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8) + box[:, 0] = box[:, 0] - xmin + box[:, 1] = box[:, 1] - ymin + cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1) + + return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] + + def unclip(self, box, unclip_ratio): + """扩展文本框""" + from shapely.geometry import Polygon + import pyclipper + + poly = Polygon(box) + distance = poly.area * unclip_ratio / poly.length + + offset = pyclipper.PyclipperOffset() + offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON) + expanded = offset.Execute(distance) + + if len(expanded) == 0: + return box + else: + return np.array(expanded[0]) + + def resize_norm_img_rec(self, img, input_shape=(320, 48)): + """ + 识别模型的图像预处理 - 固定输入形状 [1, 3, 48, 320] + """ + target_w, target_h = input_shape # 注意：宽度在前 + + h, w = img.shape[:2] + + # 计算缩放比例，保持宽高比 + ratio_h = target_h / h + ratio_w = target_w / w + ratio = min(ratio_h, ratio_w) + + # 计算缩放后的尺寸 + new_h = int(h * ratio) + new_w = int(w * ratio) + + # 调整图像大小 + resized_image = cv2.resize(img, (new_w, new_h)) + + # 创建目标尺寸的图像，用黑色填充 + padded_image = np.zeros((target_h, target_w, 3), dtype=np.float32) # 直接用float32 + + # 将缩放后的图像放到左上角（识别模型通常左对齐） + padded_image[:new_h, :new_w] = resized_image.astype(np.float32) + + # 归一化 + # padded_image = (padded_image / 255.0 - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array([0.229, 0.224, 0.225], dtype=np.float32) + + # 不缩放反而会将识别结果再移后一个?? + padded_image = (padded_image / 255.0).astype(np.float32) + padded_image = padded_image.transpose((2, 0, 1)).astype(np.float32) + + return np.expand_dims(padded_image, axis=0).astype(np.float32) + + def decode_rec_result(self, preds_prob): + """ + 解码识别结果 + """ + + # preds_idx = preds_idx[0] + preds_prob = preds_prob[0] + + preds_idx = np.argmax(preds_prob, axis=1) + preds_prob = np.max(preds_prob, axis=1) + + # CTC解码 + last_idx = 0 + preds_text = [] + preds_conf = [] + + + + # print("preds_id", len(preds_idx[0])) + + for i, idx in enumerate(preds_idx): + if idx != last_idx and idx != 0: # 0是blank + if idx < len(self.character): + # print("self.character[idx]", self.character[idx]) + # print("preds_prob[i]", preds_prob[i]) + preds_text.append(self.character[idx]) + preds_conf.append(preds_prob[i]) + last_idx = idx + + text = ''.join(preds_text) + conf = np.mean(preds_conf) if preds_conf else 0.0 + + return text, conf + + def detect_text(self, image): + """ + 文本检测 - 适配固定输入形状 [1, 3, 640, 640] + """ + ori_h, ori_w = image.shape[:2] + + # 预处理 + det_img, ratio, padding_info = self.resize_norm_img_det(image) + + # 推理 + det_output = self.det_rknn.inference(inputs=[det_img], data_format="nchw")[0] + + # 后处理 + mask = det_output[0, 0, :, :] + threshold = 0.3 + bitmap = (mask > threshold).astype(np.uint8) * 255 + + # 从位图中提取文本框（坐标是在640x640空间中的） + boxes, scores = self.boxes_from_bitmap(mask, bitmap, 640, 640) + + # 将坐标转换回原图空间 + if len(boxes) > 0: + boxes = self.post_process_det(boxes, ratio, padding_info, (ori_h, ori_w)) + + + + print("*"*100, len(boxes)) + + return boxes, scores + + def visualize_det_results(self, image_path, boxes): + image = cv2.imread(image_path) + + for box in boxes: + box = np.array(box, dtype=np.int32) + cv2.polylines(image, [box], True, (0, 255, 0), 2) + cv2.imwrite('./visual_det.jpg', image) + + def recognize_text(self, image): + """ + 文本识别 + """ + # 预处理 + rec_img = self.resize_norm_img_rec(image) + + # 推理 + rec_output = self.rec_rknn.inference(inputs=[rec_img], data_format="nchw") + + # 解码 + text, conf = self.decode_rec_result(rec_output[0]) + + # print("") + + return text, conf + + def get_rotate_crop_image(self, img, points): + """ + 根据四个点坐标裁剪并矫正图像 + """ + img_crop_width = int( + max( + np.linalg.norm(points[0] - points[1]), + np.linalg.norm(points[2] - points[3]))) + img_crop_height = int( + max( + np.linalg.norm(points[0] - points[3]), + np.linalg.norm(points[1] - points[2]))) + pts_std = np.float32([[0, 0], [img_crop_width, 0], + [img_crop_width, img_crop_height], + [0, img_crop_height]]) + M = cv2.getPerspectiveTransform(points, pts_std) + dst_img = cv2.warpPerspective( + img, + M, (img_crop_width, img_crop_height), + borderMode=cv2.BORDER_REPLICATE, + flags=cv2.INTER_CUBIC) + dst_img_height, dst_img_width = dst_img.shape[0:2] + if dst_img_height * 1.0 / dst_img_width >= 1.5: + dst_img = np.rot90(dst_img) + return dst_img + + def ocr(self, image_path): + """ + 完整的OCR流程 + """ + # 读取图像 + image = cv2.imread(image_path) + if image is None: + return [] + + # 1. 文本检测 + dt_boxes, scores = self.detect_text(image) + + # 可视化检测框 + self.visualize_det_results(image_path, dt_boxes) + + if dt_boxes is None or len(dt_boxes) == 0: + return [] + + # 2. 文本识别 + ocr_results = [] + + text_list = [] + confidence_list = [] + for i, box in enumerate(dt_boxes): + # 裁剪文本区域 + box_points = box.astype(np.float32) + crop_img = self.get_rotate_crop_image(image, box_points) + + # 识别文本 + text, conf = self.recognize_text(crop_img) + + if conf > 0.4: # 置信度过滤 + ocr_results.append({ + 'text': text, + 'confidence': conf, + 'box': box.tolist(), + 'score': scores[i] if i < len(scores) else 0.0 + }) + + text_list.append(text) + confidence_list.append(round(conf.item(), 2)) + + # return ocr_results + return [text_list, confidence_list] + + def release(self): + """释放资源""" + self.det_rknn.release() + self.rec_rknn.release() + +def main(): + parser = argparse.ArgumentParser(description='RK3588 PaddleOCR RKNN推理') + parser.add_argument('--det_model', type=str, required=True, help='文本检测RKNN模型路径') + parser.add_argument('--rec_model', type=str, required=True, help='文本识别RKNN模型路径') + parser.add_argument('--image', type=str, help='输入图像路径') + parser.add_argument('--video', type=str, help='输入视频路径') + parser.add_argument('--camera', type=int, help='摄像头设备ID') + parser.add_argument('--output', type=str, help='输出路径') + parser.add_argument('--show', action='store_true', help='显示结果') + + args = parser.parse_args() + + # 检查模型文件 + if not os.path.exists(args.det_model): + print(f"检测模型文件不存在: {args.det_model}") + return + + if not os.path.exists(args.rec_model): + print(f"识别模型文件不存在: {args.rec_model}") + return + + # 初始化OCR + print("初始化RK3588 OCR...") + ocr = RK3588OCR(args.det_model, args.rec_model) + + try: + if args.image: + # 图像模式 + print(f"处理图像: {args.image}") + + + # 进行OCR + start_time = time.time() + text, confidence = ocr.ocr(args.image) + + print("text", text) + print("confidence", confidence) + + total_time = time.time() - start_time + + # 打印结果 + print(f"\n总耗时: {total_time:.3f}s") + print(f"识别结果:") + for i in range(len(text)): + print(f"{i+1}. 文本: '{text[i]}', 置信度: {confidence[i]:.3f}") + + # 绘制结果 + # annotated_image = ocr.draw_results(image, results) + + # 保存或显示结果 + # if args.output: + # cv2.imwrite(args.output, annotated_image) + # print(f"结果已保存到: {args.output}") + + # if args.show: + # cv2.imshow('OCR结果', annotated_image) + # cv2.waitKey(0) + # cv2.destroyAllWindows() + + elif args.video or args.camera is not None: + # 视频或摄像头模式 + if args.video: + cap = cv2.VideoCapture(args.video) + print(f"处理视频: {args.video}") + else: + cap = cv2.VideoCapture(args.camera) + print(f"使用摄像头: {args.camera}") + + if not cap.isOpened(): + print("无法打开视频源") + return + + while True: + ret, frame = cap.read() + if not ret: + break + + # 进行OCR + results = ocr.ocr(frame) + + # 绘制结果 + annotated_frame = ocr.draw_results(frame, results) + + # 显示结果 + cv2.imshow('Real-time OCR', annotated_frame) + + # 按'q'退出 + if cv2.waitKey(1) & 0xFF == ord('q'): + break + + cap.release() + cv2.destroyAllWindows() + + else: + print("请指定输入源: --image, --video 或 --camera") + + finally: + ocr.release() + +if __name__ == "__main__": + # 如果直接运行脚本，提供示例用法 + ''' + 启动命令示例 + python 010使用PaddleOCR_rknn.py --det_model ./text_detection.rknn --rec_model ./text_recognition.rknn --image ./image_test/632e474452d560edd7004f745319ff00_frame_000730.jpg --output ./result.jpg + 注: + 导出的额rknn模型没有进行归一化, 归一化参数mean=0,std=1 + ''' + if len(os.sys.argv) == 1: + print("RK3588 PaddleOCR RKNN推理程序") + print("\n使用示例:") + print("# 处理单张图像") + print("python rk3588_ocr.py \\") + print(" --det_model ./rknn_models/text_detection.rknn \\") + print(" --rec_model ./rknn_models/text_recognition.rknn \\") + print(" --image ./test.jpg \\") + print(" --output ./result.jpg \\") + print(" --show") + print() + print("# 实时摄像头OCR") + print("python rk3588_ocr.py \\") + print(" --det_model ./rknn_models/text_detection.rknn \\") + print(" --rec_model ./rknn_models/text_recognition.rknn \\") + print(" --camera 0") + print() + print("# 处理视频文件") + print("python rk3588_ocr.py \\") + print(" --det_model ./rknn_models/text_detection.rknn \\") + print(" --rec_model ./rknn_models/text_recognition.rknn \\") + print(" --video ./input_video.mp4") + else: + main() \ No newline at end of file