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

Author: VikParuchuri

surya/ocr.py

@@ -60,10 +60,10 @@ def run_recognition(images: List[Image.Image], langs: List[List[str] | None], re
     return predictions_by_image
 
 
-def run_ocr(images: List[Image.Image], langs: List[List[str] | None], det_model, det_processor, rec_model, rec_processor, batch_size=None, highres_images: List[Image.Image] | None = None) -> List[OCRResult]:
+def run_ocr(images: List[Image.Image], langs: List[List[str] | None], det_model, det_processor, rec_model, rec_processor, detection_batch_size=None, recognition_batch_size=None, highres_images: List[Image.Image] | None = None) -> List[OCRResult]:
     images = convert_if_not_rgb(images)
     highres_images = convert_if_not_rgb(highres_images) if highres_images is not None else [None] * len(images)
-    det_predictions = batch_text_detection(images, det_model, det_processor)
+    det_predictions = batch_text_detection(images, det_model, det_processor, batch_size=detection_batch_size)
 
     all_slices = []
     slice_map = []
@@ -82,7 +82,7 @@ def run_ocr(images: List[Image.Image], langs: List[List[str] | None], det_model,
         all_langs.extend([lang] * len(slices))
         all_slices.extend(slices)
 
-    rec_predictions, confidence_scores = batch_recognition(all_slices, all_langs, rec_model, rec_processor, batch_size=batch_size)
+    rec_predictions, confidence_scores = batch_recognition(all_slices, all_langs, rec_model, rec_processor, batch_size=recognition_batch_size)
 
     predictions_by_image = []
     slice_start = 0