Merge pull request #311 from VikParuchuri/hotfix/height

Fix height issue and add test

Author: VikParuchuri

pyproject.toml

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

surya/detection/util.py

@@ -27,10 +27,10 @@ def split_image(img, height):
             if bottom > img_height:
                 bottom = img_height
             cropped = img.crop((0, top, img.size[0], bottom))
-            height = bottom - top
-            if height < height:
+            chunk_height = bottom - top
+            if chunk_height < height:
                 cropped = ImageOps.pad(cropped, (img.size[0], height), color=255, centering=(0, 0))
             splits.append(cropped)
-            split_heights.append(height)
+            split_heights.append(chunk_height)
         return splits, split_heights
     return [img.copy()], [img_height]

tests/conftest.py

@@ -58,3 +58,11 @@ def test_image():
               font_size=24)
     return image
 
+@pytest.fixture()
+def test_image_tall():
+    image = Image.new("RGB", (4096, 4096), "white")
+    draw = ImageDraw.Draw(image)
+    draw.text((10, 10), "Hello World", fill="black", font_size=72)
+    draw.text((4000, 4000), "This is a sentence of text.\n\nNow it is a paragraph.\n\nA three-line one.", fill="black",  font_size=24)
+    return image
+

tests/test_detection.py

@@ -5,4 +5,15 @@ def test_detection(detection_predictor, test_image):
     assert detection_results[0].image_bbox == [0, 0, 1024, 1024]
 
     bboxes = detection_results[0].bboxes
-    assert len(bboxes) == 4
+    assert len(bboxes) == 4
+
+
+def test_detection_chunking(detection_predictor, test_image_tall):
+    detection_results = detection_predictor([test_image_tall])
+
+    assert len(detection_results) == 1
+    assert detection_results[0].image_bbox == [0, 0, 4096, 4096]
+
+    bboxes = detection_results[0].bboxes
+    assert len(bboxes) >= 3 # Sometimes merges into 3
+    assert abs(4000 - bboxes[1].polygon[0][0]) < 50