Merge remote-tracking branch 'origin/pre_release' into pre_release

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
 

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张图片上训练得到
 
 ### 公式检测
 

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

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',