Merge pull request #341 from VikParuchuri/dev

Author: tarun-menta

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surya-ocr"
-version = "0.13.0"
+version = "0.13.1"
 description = "OCR, layout, reading order, and table recognition in 90+ languages"
 authors = ["Vik Paruchuri <[email protected]>"]
 readme = "README.md"

surya/detection/__init__.py

@@ -135,26 +135,22 @@ def batch_detection(
 
 class InlineDetectionPredictor(DetectionPredictor):
     model_loader_cls = InlineDetectionModelLoader
-    
-    def batch_generator(self, iterable, batch_size=None):
-        if batch_size is None:
-            batch_size = self.get_batch_size()
-
-        for i in range(0, len(iterable), batch_size):
-            yield iterable[i:i+batch_size]
 
     def __call__(self, images, text_boxes: List[List[List[float]]], batch_size=None, include_maps=False) -> List[TextDetectionResult]:
         detection_generator = self.batch_detection(images, batch_size=batch_size, static_cache=settings.DETECTOR_STATIC_CACHE)
-        text_box_generator = self.batch_generator(text_boxes, batch_size=batch_size)
 
         postprocessing_futures = []
         max_workers = min(settings.DETECTOR_POSTPROCESSING_CPU_WORKERS, len(images))
         parallelize = not settings.IN_STREAMLIT and len(images) >= settings.DETECTOR_MIN_PARALLEL_THRESH
         executor = ThreadPoolExecutor if parallelize else FakeExecutor
+
+        image_idx = 0
         with executor(max_workers=max_workers) as e:
-            for (preds, orig_sizes), batch_text_boxes in zip(detection_generator, text_box_generator):
-                for pred, orig_size, image_text_boxes in zip(preds, orig_sizes, batch_text_boxes):
-                    postprocessing_futures.append(e.submit(parallel_get_inline_boxes, pred, orig_size, image_text_boxes, include_maps))
+            for (preds, orig_sizes) in detection_generator:
+                for pred, orig_size in zip(preds, orig_sizes):
+                    postprocessing_futures.append(e.submit(parallel_get_inline_boxes, pred, orig_size, text_boxes[image_idx], include_maps))
+                    image_idx += 1
 
         assert len(postprocessing_futures) == len(images) == len(text_boxes) # Ensure we have a 1:1 mapping
+
         return [future.result() for future in postprocessing_futures]

surya/detection/heatmap.py

@@ -156,6 +156,7 @@ def parallel_get_boxes(preds, orig_sizes, include_maps=False):
         # Skip for vertical boxes
         if box.height < 3 * box.width:
             box.expand(x_margin=0, y_margin=settings.DETECTOR_BOX_Y_EXPAND_MARGIN)
+            box.fit_to_bounds([0, 0, orig_sizes[0], orig_sizes[1]])     # Fix any bad expands
 
     result = TextDetectionResult(
         bboxes=bboxes,