Refactor rec predictor initialization in scripts + benchmarks

Author: tarun-menta

README.md

@@ -132,11 +132,13 @@ Setting the `RECOGNITION_BATCH_SIZE` env var properly will make a big difference
 
 ```python
 from PIL import Image
+from surya.foundation import FoudnationPredictor
 from surya.recognition import RecognitionPredictor
 from surya.detection import DetectionPredictor
 
 image = Image.open(IMAGE_PATH)
-recognition_predictor = RecognitionPredictor()
+foundation_predictor = FoundationPredictor()
+recognition_predictor = RecognitionPredictor(foundation_predictor)
 detection_predictor = DetectionPredictor()
 
 predictions = recognition_predictor([image], det_predictor=detection_predictor)

benchmark/recognition.py

@@ -7,6 +7,7 @@
 from benchmark.utils.scoring import overlap_score, overlap_score_exact
 from surya.input.processing import convert_if_not_rgb
 from surya.debug.text import draw_text_on_image
+from surya.foundation import FoundationPredictor
 from surya.recognition import RecognitionPredictor
 from surya.settings import settings
 from surya.recognition.languages import CODE_TO_LANGUAGE
@@ -112,7 +113,8 @@ def main(
     textract_cpus: int,
     languages: str | None,
 ):
-    rec_predictor = RecognitionPredictor()
+    foundation_predictor = FoundationPredictor()
+    rec_predictor = RecognitionPredictor(foundation_predictor)
 
     split = "train"
     dataset = datasets.load_dataset(

benchmark/texify.py

@@ -11,6 +11,7 @@
 
 from surya.common.surya.schema import TaskNames
 from surya.settings import settings
+from surya.foundation import FoundationPredictor
 from surya.recognition import RecognitionPredictor, OCRResult
 import json
 from rapidfuzz.distance import Levenshtein
@@ -77,7 +78,8 @@ def inference_texify(
     "--line_mode", is_flag=True, help="Use line mode for texify.", default=False
 )
 def main(ds_name: str, results_dir: str, max_rows: int, line_mode: bool):
-    predictor = RecognitionPredictor()
+    foundation_predictor = FoundationPredictor()
+    predictor = RecognitionPredictor(foundation_predictor)
     ds = datasets.load_dataset(ds_name, split="train")
 
     if max_rows:

surya/models.py

@@ -7,6 +7,7 @@
 from surya.layout import LayoutPredictor
 from surya.logging import configure_logging
 from surya.ocr_error import OCRErrorPredictor
+from surya.foundation import FoundationPredictor
 from surya.recognition import RecognitionPredictor
 from surya.table_rec import TableRecPredictor
 
@@ -16,10 +17,11 @@
 def load_predictors(
     device: str | torch.device | None = None, dtype: torch.dtype | str | None = None
 ) -> Dict[str, BasePredictor]:
+    foundation_predictor = FoundationPredictor(device=device, dtype=dtype)
     return {
         "layout": LayoutPredictor(device=device, dtype=dtype),
         "ocr_error": OCRErrorPredictor(device=device, dtype=dtype),
-        "recognition": RecognitionPredictor(device=device, dtype=dtype),
+        "recognition": RecognitionPredictor(foundation_predictor),
         "detection": DetectionPredictor(device=device, dtype=dtype),
         "table_rec": TableRecPredictor(device=device, dtype=dtype),
     }

surya/scripts/ocr_latex.py

@@ -7,6 +7,7 @@
 
 from surya.logging import configure_logging, get_logger
 from surya.scripts.config import CLILoader
+from surya.foundation import FoundationPredictor
 from surya.recognition import RecognitionPredictor
 from surya.common.surya.schema import TaskNames
 
@@ -19,7 +20,8 @@
 def ocr_latex_cli(input_path: str, **kwargs):
     loader = CLILoader(input_path, kwargs, highres=True)
 
-    texify_predictor = RecognitionPredictor()
+    foundation_predictor = FoundationPredictor()
+    texify_predictor = RecognitionPredictor(foundation_predictor)
     tasks = [TaskNames.block_without_boxes] * len(loader.images)
     bboxes = [[[0, 0, image.width, image.height]] for image in loader.images]
 

surya/scripts/ocr_text.py

@@ -8,6 +8,7 @@
 from surya.detection import DetectionPredictor
 from surya.debug.text import draw_text_on_image
 from surya.logging import configure_logging, get_logger
+from surya.foundation import FoundationPredictor
 from surya.recognition import RecognitionPredictor
 from surya.scripts.config import CLILoader
 
@@ -25,8 +26,9 @@ def ocr_text_cli(input_path: str, task_name: str, disable_math: bool, **kwargs):
     loader = CLILoader(input_path, kwargs, highres=True)
     task_names = [task_name] * len(loader.images)
 
+    foundation_predictor = FoundationPredictor()
     det_predictor = DetectionPredictor()
-    rec_predictor = RecognitionPredictor()
+    rec_predictor = RecognitionPredictor(foundation_predictor)
 
     start = time.time()
     predictions_by_image = rec_predictor(

surya/scripts/texify_app.py

@@ -3,6 +3,7 @@
 from typing import List
 
 from surya.recognition import RecognitionPredictor
+from surya.foundation import FoundationPredictor
 from surya.common.surya.schema import TaskNames
 
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = (
@@ -33,7 +34,8 @@ def replace_fences(text):
 
 @st.cache_resource()
 def load_predictor():
-    return RecognitionPredictor()
+    foundation_predictor = FoundationPredictor()
+    return RecognitionPredictor(foundation_predictor)
 
 
 @st.cache_data()

tests/conftest.py

@@ -9,6 +9,7 @@
 from surya.ocr_error import OCRErrorPredictor
 from surya.layout import LayoutPredictor
 from surya.recognition import RecognitionPredictor
+from surya.foundation import FoundationPredictor
 from surya.table_rec import TableRecPredictor
 
 
@@ -35,7 +36,8 @@ def detection_predictor() -> DetectionPredictor:
 
 @pytest.fixture(scope="session")
 def recognition_predictor() -> RecognitionPredictor:
-    recognition_predictor = RecognitionPredictor()
+    foundation_predictor = FoundationPredictor()
+    recognition_predictor = RecognitionPredictor(foundation_predictor)
     yield recognition_predictor
     del recognition_predictor