diff --git a/README.md b/README.md index 72f00b5..3b9209e 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,7 @@ TexTeller also supports **formula detection and recognition** on full images, al ### Download Weights -Chinese-English documentation formula detection [[link](https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco.onnx?download=true)]: Trained on 3415 Chinese textbook images (130+ layouts). +Chinese and English document formula detection [[link](https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco.onnx?download=true)]: Trained on a total of 11,867 images, consisting of 3,415 images from Chinese textbooks (130+ layouts) and 8,272 images from the [IBEM dataset](https://zenodo.org/records/4757865). ### Formula Detection diff --git a/assets/README_zh.md b/assets/README_zh.md index c0243fd..443b9fe 100644 --- a/assets/README_zh.md +++ b/assets/README_zh.md @@ -115,7 +115,7 @@ TexTeller还支持对整张图片进行**公式检测+公式识别**,从而对 ### 下载权重 -中文英文文档公式检测 [[link](https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco.onnx?download=true)]:在3415张中文教材数据(130+版式)上训练得到 +中文英文文档公式检测 [[link](https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco.onnx?download=true)]:在3415张中文教材数据(130+版式)和8272张[IBEM数据集](https://zenodo.org/records/4757865)上,共11867张图片上训练得到 ### 公式检测 diff --git a/src/infer_det.py b/src/infer_det.py index a5047fc..b90a1c8 100644 --- a/src/infer_det.py +++ b/src/infer_det.py @@ -1,9 +1,11 @@ import os import argparse import glob +import subprocess from onnxruntime import InferenceSession from pathlib import Path + from models.det_model.inference import PredictConfig, predict_image @@ -12,8 +14,8 @@ parser.add_argument("--infer_cfg", type=str, help="infer_cfg.yml", default="./models/det_model/model/infer_cfg.yml") parser.add_argument('--onnx_file', type=str, help="onnx model file path", default="./models/det_model/model/rtdetr_r50vd_6x_coco.onnx") -parser.add_argument("--image_dir", type=str) -parser.add_argument("--image_file", type=str, required=True) +parser.add_argument("--image_dir", type=str, default='./testImgs') +parser.add_argument("--image_file", type=str) parser.add_argument("--imgsave_dir", type=str, default="./detect_results") @@ -47,6 +49,10 @@ def get_test_images(infer_dir, infer_img): return images +def download_file(url, filename): + print(f"Downloading {filename}...") + subprocess.run(["wget", "-q", "--show-progress", "-O", filename, url], check=True) + print("Download complete.") if __name__ == '__main__': cur_path = os.getcwd() @@ -54,6 +60,15 @@ if __name__ == '__main__': os.chdir(script_dirpath) FLAGS = parser.parse_args() + + if not os.path.exists(FLAGS.infer_cfg): + infer_cfg_url = "https://huggingface.co/TonyLee1256/texteller_det/resolve/main/infer_cfg.yml?download=true" + download_file(infer_cfg_url, FLAGS.infer_cfg) + + if not os.path.exists(FLAGS.onnx_file): + onnx_file_url = "https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco.onnx?download=true" + download_file(onnx_file_url, FLAGS.onnx_file) + # load image list img_list = get_test_images(FLAGS.image_dir, FLAGS.image_file) # load predictor diff --git a/src/rec_infer_from_crop_imgs.py b/src/rec_infer_from_crop_imgs.py index 73bfa73..89bef18 100644 --- a/src/rec_infer_from_crop_imgs.py +++ b/src/rec_infer_from_crop_imgs.py @@ -11,10 +11,16 @@ if __name__ == '__main__': os.chdir(Path(__file__).resolve().parent) parser = argparse.ArgumentParser() parser.add_argument( - '-img', + '-img_dir', type=str, - required=True, - help='path to the input image' + help='path to the input image', + default='./detect_results/subimages' + ) + parser.add_argument( + '-output_dir', + type=str, + help='path to the output dir', + default='./rec_results' ) parser.add_argument( '--inference-mode',