@@ -108,7 +108,7 @@ result.vis("vis_result.jpg")
108108
109109![ alt text] ( ../images/vis_result.jpg )
110110
111- #### 4. 模型精度测试
111+ ### 4. 模型精度测试
112112
113113该部分主要使用[ TextRecMetric] ( https://github.com/SWHL/TextRecMetric ) 和测试集[ text_rec_test_dataset] ( https://huggingface.co/datasets/SWHL/text_rec_test_dataset ) 来评测。
114114
@@ -122,72 +122,74 @@ result.vis("vis_result.jpg")
122122{'ExactMatch': 0.8097, 'CharMatch': 0.9444, 'avg_elapse': 0.0818 }
123123```
124124
125- ### 4 . 集成到rapidocr中
125+ ### 5 . 集成到rapidocr中
126126
127127该部分主要包括将字典文件写入到ONNX模型中、托管模型到魔搭、更改rapidocr中模型配置文件、编写对应单元测试等。
128128
129129#### 字典文件写入ONNX模型
130130
131- ``` python linenums="1"
132- from pathlib import Path
133- from typing import List, Union
134-
135- import onnx
136- import onnxruntime as ort
137- from onnx import ModelProto
138-
139-
140- def read_txt (txt_path : Union[Path, str ]) -> List[str ]:
141- with open (txt_path, " r" , encoding = " utf-8" ) as f:
142- data = [v.rstrip(" \n " ) for v in f]
143- return data
144-
145-
146- class ONNXMetaOp :
147- @ classmethod
148- def add_meta (
149- cls ,
150- model_path : Union[str , Path],
151- key : str ,
152- value : List[str ],
153- delimiter : str = " \n " ,
154- ) -> ModelProto:
155- model = onnx.load_model(model_path)
156- meta = model.metadata_props.add()
157- meta.key = key
158- meta.value = delimiter.join(value)
159- return model
160-
161- @ classmethod
162- def get_meta (
163- cls , model_path : Union[str , Path], key : str , split_sym : str = " \n "
164- ) -> List[str ]:
165- sess = ort.InferenceSession(model_path)
166- meta_map = sess.get_modelmeta().custom_metadata_map
167- key_content = meta_map.get(key)
168- key_list = key_content.split(split_sym)
169- return key_list
170-
171- @ classmethod
172- def del_meta (cls , model_path : Union[str , Path]) -> ModelProto:
173- model = onnx.load_model(model_path)
174- del model.metadata_props[:]
175- return model
176-
177- @ classmethod
178- def save_model (cls , save_path : Union[str , Path], model : ModelProto):
179- onnx.save_model(model, save_path)
180-
181-
182- dicts = read_txt(
183- " /Users/joshuawang/projects/_self/tmp/RapidOCR/paddle/PP-OCRv4/rec/ch_PP-OCRv4_rec_server_doc/ppocrv4_doc_dict.txt"
184- )
185- model_path = " models/PP-OCRv4_server_rec_doc.onnx"
186- model = ONNXMetaOp.add_meta(model_path, key = " character" , value = dicts)
187-
188- new_model_path = " models/PP-OCRv4_server_rec_doc_with_dict.onnx"
189- ONNXMetaOp.save_model(new_model_path, model)
190-
191- t = ONNXMetaOp.get_meta(new_model_path, key = " character" )
192- print (t)
193- ```
131+ !!! info "代码"
132+
133+ ```python linenums="1"
134+ from pathlib import Path
135+ from typing import List, Union
136+
137+ import onnx
138+ import onnxruntime as ort
139+ from onnx import ModelProto
140+
141+
142+ def read_txt(txt_path: Union[Path, str]) -> List[str]:
143+ with open(txt_path, "r", encoding="utf-8") as f:
144+ data = [v.rstrip("\n") for v in f]
145+ return data
146+
147+
148+ class ONNXMetaOp:
149+ @classmethod
150+ def add_meta(
151+ cls,
152+ model_path: Union[str, Path],
153+ key: str,
154+ value: List[str],
155+ delimiter: str = "\n",
156+ ) -> ModelProto:
157+ model = onnx.load_model(model_path)
158+ meta = model.metadata_props.add()
159+ meta.key = key
160+ meta.value = delimiter.join(value)
161+ return model
162+
163+ @classmethod
164+ def get_meta(
165+ cls, model_path: Union[str, Path], key: str, split_sym: str = "\n"
166+ ) -> List[str]:
167+ sess = ort.InferenceSession(model_path)
168+ meta_map = sess.get_modelmeta().custom_metadata_map
169+ key_content = meta_map.get(key)
170+ key_list = key_content.split(split_sym)
171+ return key_list
172+
173+ @classmethod
174+ def del_meta(cls, model_path: Union[str, Path]) -> ModelProto:
175+ model = onnx.load_model(model_path)
176+ del model.metadata_props[:]
177+ return model
178+
179+ @classmethod
180+ def save_model(cls, save_path: Union[str, Path], model: ModelProto):
181+ onnx.save_model(model, save_path)
182+
183+
184+ dicts = read_txt(
185+ "models/ppocrv4_doc_dict.txt"
186+ )
187+ model_path = "models/PP-OCRv4_server_rec_doc.onnx"
188+ model = ONNXMetaOp.add_meta(model_path, key="character", value=dicts)
189+
190+ new_model_path = "models/PP-OCRv4_server_rec_doc_with_dict.onnx"
191+ ONNXMetaOp.save_model(new_model_path, model)
192+
193+ t = ONNXMetaOp.get_meta(new_model_path, key="character")
194+ print(t)
195+ ```
0 commit comments