YogeLiu/TexTeller
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

bugfix inference.py

Browse Source
This commit is contained in:
TonyLee1256
2024-05-07 13:19:43 +08:00
committed by GitHub
parent 7c5a547b1f
commit 8c85575260
1 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/inference.py
View File

@@ -1,4 +1,5 @@
import os import os
import sys
import argparse import argparse
import cv2 as cv import cv2 as cv
@@ -44,8 +45,17 @@ if __name__ == '__main__':
help='use mix mode' help='use mix mode'
) )
parser.add_argument(
'-lang',
type=str,
default='None'
)
args = parser.parse_args() args = parser.parse_args()
if args.mix and args.lang == "None":
print("When -mix is set, -lang must be set (support: ['zh', 'en'])")
sys.exit(-1)
# You can use your own checkpoint and tokenizer path. # You can use your own checkpoint and tokenizer path.
print('Loading model and tokenizer...') print('Loading model and tokenizer...')
latex_rec_model = TexTeller.from_pretrained() latex_rec_model = TexTeller.from_pretrained()
@@ -61,12 +71,12 @@ if __name__ == '__main__':
print(res) print(res)
else: else:
infer_config = PredictConfig("./models/det_model/model/infer_cfg.yml") infer_config = PredictConfig("./models/det_model/model/infer_cfg.yml")
latex_det_model = InferenceSession("./models/det_model/model/rtdetr_r50vd_6x_coco_IBEM_cnTextBook.onnx") latex_det_model = InferenceSession("./models/det_model/model/rtdetr_r50vd_6x_coco.onnx")
det_processor, det_model = segformer.load_processor(), segformer.load_model() det_processor, det_model = segformer.load_processor(), segformer.load_model()
rec_model, rec_processor = load_model(), load_processor() rec_model, rec_processor = load_model(), load_processor()
lang_ocr_models = [det_model, det_processor, rec_model, rec_processor] lang_ocr_models = [det_model, det_processor, rec_model, rec_processor]
latex_rec_models = [latex_rec_model, tokenizer] latex_rec_models = [latex_rec_model, tokenizer]
res = mix_inference(img_path, infer_config, latex_det_model, lang_ocr_models, latex_rec_models, args.inference_mode, args.num_beam) res = mix_inference(img_path, args.lang , infer_config, latex_det_model, lang_ocr_models, latex_rec_models, args.inference_mode, args.num_beam)
print(res) print(res)