添加识别文字字典, 测试onnx模型中添加识别字典
This commit is contained in:
parent
6331efe449
commit
eacc373e10
@ -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()
|
||||
58
009PaddleOCR_ONNX转rknn.py
Normal file
58
009PaddleOCR_ONNX转rknn.py
Normal 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
7
011加载识别字典.py
Normal 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']))
|
||||
Loading…
Reference in New Issue
Block a user