fix(ocr): rotate image to the natural orientation before layout prediction

Signed-off-by: Clément Doumouro <[email protected]>

Author: ClemDoum

docling/models/layout_model.py

@@ -19,6 +19,7 @@
 from docling.models.utils.hf_model_download import download_hf_model
 from docling.utils.accelerator_utils import decide_device
 from docling.utils.layout_postprocessor import LayoutPostprocessor
+from docling.utils.orientation import detect_orientation
 from docling.utils.profiling import TimeRecorder
 from docling.utils.visualization import draw_clusters
 
@@ -164,6 +165,10 @@ def __call__(
             page_image = page.get_image(scale=1.0)
             assert page_image is not None
 
+            page_orientation = detect_orientation(page.cells)
+            if page_orientation:
+                page_image = page_image.rotate(-page_orientation, expand=True)
+
             valid_pages.append(page)
             valid_page_images.append(page_image)
 

docling/utils/orientation.py

@@ -1,3 +1,20 @@
+from collections import Counter
+from operator import itemgetter
+
+from docling_core.types.doc.page import TextCell
+
+_ORIENTATIONS = [0, 90, 180, 270]
+
+
+def _clipped_orientation(angle: float) -> int:
+    return min((abs(angle - o) % 360, o) for o in _ORIENTATIONS)[1]
+
+
+def detect_orientation(cells: list[TextCell]) -> int:
+    if not cells:
+        return 0
+    orientation_counter = Counter(_clipped_orientation(c.rect.angle_360) for c in cells)
+    return max(orientation_counter.items(), key=itemgetter(1))[0]
 from typing import Tuple
 
 from docling_core.types.doc import BoundingBox, CoordOrigin

tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_180.doctags.txt

@@ -1,4 +1,4 @@
 <document>
-<paragraph><location><page_1><loc_74><loc_16><loc_88><loc_18></location>package</paragraph>
+<paragraph><location><page_1><loc_75><loc_16><loc_88><loc_18></location>package</paragraph>
 <paragraph><location><page_1><loc_15><loc_9><loc_88><loc_15></location>Docling bundles PDF document conversion to JSON and Markdown in an easy self contained</paragraph>
 </document>

tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_270.doctags.txt

@@ -1,3 +1,3 @@
 <document>
-<paragraph><location><page_1><loc_82><loc_74><loc_84><loc_88></location>package</paragraph>
+<paragraph><location><page_1><loc_82><loc_75><loc_84><loc_88></location>package</paragraph>
 </document>

tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.doctags.txt

@@ -1,3 +1,4 @@
 <document>
-<paragraph><location><page_1><loc_16><loc_12><loc_18><loc_26></location>package</paragraph>
+<paragraph><location><page_1><loc_9><loc_12><loc_11><loc_85></location>Docling bundles PDF document conversion to</paragraph>
+<paragraph><location><page_1><loc_12><loc_12><loc_15><loc_85></location><location><page_1><loc_12><loc_12><loc_15><loc_85></location>JSON and Markdown in an easy self contained package</paragraph>
 </document>

tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.md

@@ -1 +1,3 @@
-package
+Docling bundles PDF document conversion to
+
+JSON and Markdown in an easy self contained package

tests/data_scanned/groundtruth/docling_v2/ocr_test_rotated_90.md

@@ -1 +1,3 @@
-package
+Docling bundles PDF document conversion to
+
+JSON and Markdown in an easy self contained package