Create ocr_text.py

ocr_text.py

@@ -0,0 +1,40 @@
+import os
+
+from surya.input.langs import replace_lang_with_code, get_unique_langs
+from surya.input.load import load_from_folder, load_from_file, load_lang_file
+from surya.model.detection.model import load_model as load_detection_model, load_processor as load_detection_processor
+from surya.model.recognition.model import load_model as load_recognition_model
+from surya.model.recognition.processor import load_processor as load_recognition_processor
+from surya.model.recognition.tokenizer import _tokenize
+from surya.ocr import run_ocr
+from surya.postprocessing.text import draw_text_on_image
+
+
+def main(input_path, max_pages=None, start_page=0, langs=None, lang_file=None):
+
+    assert langs or lang_file, "Must provide either langs or lang_file"
+
+    if os.path.isdir(input_path):
+        images, names = load_from_folder(input_path, max_pages, start_page)
+    else:
+        images, names = load_from_file(input_path, max_pages, start_page)
+
+
+    langs = langs.split(",")
+    replace_lang_with_code(langs)
+    image_langs = [langs] * len(images)
+
+    det_processor = load_detection_processor()
+    det_model = load_detection_model()
+
+    _, lang_tokens = _tokenize("", get_unique_langs(image_langs))
+    rec_model = load_recognition_model(langs=lang_tokens) # Prune model moe layer to only include languages we need
+    rec_processor = load_recognition_processor()
+
+    predictions_by_image = run_ocr(images, image_langs, det_model, det_processor, rec_model, rec_processor)
+
+    for idx, (name, image, pred, langs) in enumerate(zip(names, images, predictions_by_image, image_langs)):
+        bboxes = [l.bbox for l in pred.text_lines]
+        pred_text = [l.text for l in pred.text_lines]
+        page_image = draw_text_on_image(bboxes, pred_text, image.size, langs, has_math="_math" in langs)
+        return page_image