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

This commit is contained in:
haotian 2025-08-15 16:44:27 +08:00
parent 6331efe449
commit eacc373e10
4 changed files with 18475 additions and 10 deletions

View File

@ -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', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+',
',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8',
'9', ':', ';', '<', '=', '>', '?', '@', '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', 'Z', '[', '\\', ']', '^', '_',
'`', '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',
'z', '{', '|', '}', '~'] + [chr(i) for i in range(19968, 40870)] # 中文字符
# self.character = ['blank', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+',
# ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8',
# '9', ':', ';', '<', '=', '>', '?', '@', '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', 'Z', '[', '\\', ']', '^', '_',
# '`', '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',
# '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()

View File

@ -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()

7
011加载识别字典.py Normal file
View File

@ -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']))

18385
dict.yaml Normal file

File diff suppressed because it is too large Load Diff