新增公式检测模块

src/rec_infer_from_crop_imgs.py

@@ -0,0 +1,59 @@
+import os
+import argparse
+import cv2 as cv
+from pathlib import Path
+from utils import to_katex
+from models.ocr_model.utils.inference import inference as latex_inference
+from models.ocr_model.model.TexTeller import TexTeller
+
+
+if __name__ == '__main__':
+    os.chdir(Path(__file__).resolve().parent)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '-img_dir',
+        type=str,
+        default="./subimages",
+        help='path to the directory containing input images'
+    )
+    parser.add_argument(
+        '-output_dir',
+        type=str,
+        default="./results",
+        help='path to the output directory for storing recognition results'
+    )
+    parser.add_argument(
+        '-cuda',
+        default=False,
+        action='store_true',
+        help='use cuda or not'
+    )
+
+    args = parser.parse_args()
+
+    print('Loading model and tokenizer...')
+    latex_rec_model = TexTeller.from_pretrained()
+    tokenizer = TexTeller.get_tokenizer()
+    print('Model and tokenizer loaded.')
+
+    # Create the output directory if it doesn't exist
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Loop through all images in the input directory
+    for filename in os.listdir(args.img_dir):
+        img_path = os.path.join(args.img_dir, filename)
+        img = cv.imread(img_path)
+
+        if img is not None:
+            print(f'Inference for {filename}...')
+            res = latex_inference(latex_rec_model, tokenizer, [img], args.cuda)
+            res = to_katex(res[0])
+
+            # Save the recognition result to a text file
+            output_file = os.path.join(args.output_dir, os.path.splitext(filename)[0] + '.txt')
+            with open(output_file, 'w') as f:
+                f.write(res)
+
+            print(f'Result saved to {output_file}')
+        else:
+            print(f"Warning: Could not read image {img_path}. Skipping...")