diff --git a/docling/models/layout_model.py b/docling/models/layout_model.py index 44e7286d..f1fffb54 100644 --- a/docling/models/layout_model.py +++ b/docling/models/layout_model.py @@ -18,6 +18,7 @@ from docling.models.base_model import BasePageModel 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 @@ -155,7 +156,9 @@ class LayoutModel(BasePageModel): assert page.size is not None 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) clusters = [] for ix, pred_item in enumerate( self.layout_predictor.predict(page_image) diff --git a/docling/utils/orientation.py b/docling/utils/orientation.py index 29c02ff7..f9c30096 100644 --- a/docling/utils/orientation.py +++ b/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 diff --git a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_180.doctags.txt b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_180.doctags.txt index 029be08d..50f50834 100644 --- a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_180.doctags.txt +++ b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_180.doctags.txt @@ -1,4 +1,4 @@ -package +package Docling bundles PDF document conversion to JSON and Markdown in an easy self contained \ No newline at end of file diff --git a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_270.doctags.txt b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_270.doctags.txt index d5c2972a..8350737b 100644 --- a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_270.doctags.txt +++ b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_270.doctags.txt @@ -1,3 +1,3 @@ -package +package \ No newline at end of file diff --git a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.doctags.txt b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.doctags.txt index 0b7a3a14..c1068b56 100644 --- a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.doctags.txt +++ b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.doctags.txt @@ -1,3 +1,4 @@ -package +Docling bundles PDF document conversion to +JSON and Markdown in an easy self contained package \ No newline at end of file diff --git a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.md b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.md index 597acc76..8d77a437 100644 --- a/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.md +++ b/tests/data_scanned/groundtruth/docling_v1/ocr_test_rotated_90.md @@ -1 +1,3 @@ -package \ No newline at end of file +Docling bundles PDF document conversion to + +JSON and Markdown in an easy self contained package \ No newline at end of file diff --git a/tests/data_scanned/groundtruth/docling_v2/ocr_test_rotated_90.md b/tests/data_scanned/groundtruth/docling_v2/ocr_test_rotated_90.md index 597acc76..8d77a437 100644 --- a/tests/data_scanned/groundtruth/docling_v2/ocr_test_rotated_90.md +++ b/tests/data_scanned/groundtruth/docling_v2/ocr_test_rotated_90.md @@ -1 +1,3 @@ -package \ No newline at end of file +Docling bundles PDF document conversion to + +JSON and Markdown in an easy self contained package \ No newline at end of file