添加识别文字字典, 测试onnx模型中添加识别字典

008测试PadleCOR_ONNX固定输入形状模型全流程.py

@@ -1,4 +1,5 @@
 import cv2
+import yaml
 import numpy as np
 import onnxruntime as ort
 from PIL import Image, ImageDraw, ImageFont
@@ -22,14 +23,26 @@ class PaddleOCRONNX:
         self.rec_input_name = self.rec_session.get_inputs()[0].name
         
         # 字符集（根据您的模型调整）
-        self.character = ['blank', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', 
+        # self.character = ['blank', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', 
-                         ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', 
+        #                  ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', 
-                         '9', ':', ';', '<', '=', '>', '?', '@', 'A', 'B', 'C', 'D', 'E', 
+        #                  '9', ':', ';', '<', '=', '>', '?', '@', 'A', 'B', 'C', 'D', 'E', 
-                         'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 
+        #                  'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 
-                         'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', '\\', ']', '^', '_', 
+        #                  'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', '\\', ']', '^', '_', 
-                         '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 
+        #                  '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 
-                         'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 
+        #                  'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 
-                         'z', '{', '|', '}', '~'] + [chr(i) for i in range(19968, 40870)]  # 中文字符
+        #                  'z', '{', '|', '}', '~'] + [chr(i) for i in range(19968, 40870)]  # 中文字符
+        self.character = self.get_dict()
+        
+        if self.character is None:
+            raise ValueError('请检查字典文件是否存在！')
+    
+    def get_dict(self, dict_path='./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 resize_norm_img_det(self, img, input_shape=(640, 640)):
         """
@@ -345,7 +358,9 @@ class PaddleOCRONNX:
 # 使用示例
 def main():
     # 初始化OCR
-    ocr = PaddleOCRONNX('/home/admin-root/haotian/康达瑞贝斯机器狗/det_shape.onnx', '/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape.onnx')
+    # ocr = PaddleOCRONNX('/home/admin-root/haotian/康达瑞贝斯机器狗/det_shape.onnx', '/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape.onnx')
+    
+    ocr = PaddleOCRONNX('/home/admin-root/haotian/康达瑞贝斯机器狗/det_shape_20250814.onnx', '/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape_20250815.onnx')
     
     # 执行OCR
     image_path = '/home/admin-root/haotian/康达瑞贝斯机器狗/data_image/001读表图片/3aee64cc1f90d93a5a45979f7b17cb4b_frame_001460.jpg'
@@ -377,7 +392,7 @@ def visualize_results(image_path, results):
                    (box[0][0], box[0][1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
     
-    cv2.imwrite('result_shape.jpg', image)
+    cv2.imwrite('result_shape_20250815.jpg', image)
 
 if __name__ == "__main__":
     main()

009PaddleOCR_ONNX转rknn.py

@@ -0,0 +1,58 @@
+from rknn.api import RKNN
+import cv2
+import numpy as np
+
+
+# 初始化 RKNN
+rknn = RKNN()
+
+# 配置参数（关键！）
+rknn.config(
+    target_platform="rk3588",          # 根据实际芯片型号修改
+    mean_values=[[123.675, 116.28, 103.53]],  # 默认输入是255, 使用这些值归一化到0-1
+    std_values=[[58.395, 57.12, 57.375]],
+    quant_img_RGB2BGR=True,
+    optimization_level=3,          # 最高优化级别
+)
+
+# 加载 ONNX, 只支持固定输入的onnx模型
+ret = rknn.load_onnx(model="/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape.onnx")
+assert ret == 0, "加载 ONNX 失败！"
+
+# 转换模型
+ret = rknn.build(
+    do_quantization=False,          # 启用量化
+    # dataset="dataset.txt",         # 校准数据路径
+)
+assert ret == 0, "转换 RKNN 失败！"
+
+# 导出 RKNN
+ret = rknn.export_rknn("/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape.rknn")
+assert ret == 0, "导出 RKNN 失败！"
+
+
+# Set inputs
+img = cv2.imread('/home/admin-root/haotian/rk3588/pytorch模型转rknn/images/bus.jpg')
+img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+img.resize((3, 640, 640))
+img = np.expand_dims(img, 0)
+
+# Init runtime environment
+print('--> Init runtime environment')
+ret = rknn.init_runtime()
+if ret != 0:
+    print('Init runtime environment failed!')
+    exit(ret)
+print('done')
+
+# Inference
+print('--> Running model')
+outputs = rknn.inference(inputs=[img], data_format=['nchw'])
+np.save('./tflite_mobilenet_v1_0.npy', outputs[0])
+print(len(outputs))
+print('done')
+
+rknn.release()
+
+
+

011加载识别字典.py

@@ -0,0 +1,7 @@
+import yaml 
+
+
+with open('dict.yaml', 'r', encoding='utf-8') as f:
+    dict_rec = yaml.safe_load(f)
+    
+print(len(dict_rec['character_dict']))