From ce201cf39a6fec2f02d657d3cec8af81583cb433 Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 21 Jul 2026 14:36:01 +0200 Subject: [PATCH] [reading order] de-skew rotated layout geoms --- docs/source/modules/models.rst | 24 ++++++ doctr/models/reading_order/base.py | 99 +++++++++++++++++++++++ tests/common/test_models_reading_order.py | 92 ++++++++++++++++++++- 3 files changed, 214 insertions(+), 1 deletion(-) diff --git a/docs/source/modules/models.rst b/docs/source/modules/models.rst index 661cd0b14a..8083f9922c 100644 --- a/docs/source/modules/models.rst +++ b/docs/source/modules/models.rst @@ -146,3 +146,27 @@ doctr.models.utils .. autofunction:: export_model_to_onnx .. autofunction:: add_whitelist + + +doctr.models.reading_order +-------------------------- + +Language-independent estimation of the reading order of the elements of a page, combining a column-following +topological sort of their spatial relations with optional layout semantics (headers, footers, captions). The +reading-order-aware export of a :class:`~doctr.io.elements.Page` / :class:`~doctr.io.elements.Document` to +Markdown or AsciiDoc is exposed through the ``export_as_markdown`` / ``export_as_asciidoc`` methods of +:mod:`doctr.io.elements`. + +.. currentmodule:: doctr.models.reading_order + +.. autofunction:: detect_text_direction + +.. autofunction:: sort_reading_order + +.. autofunction:: resolve_reading_segments + +.. autofunction:: assign_layout_labels + +.. autofunction:: deskew_reading_geometries + +.. autoclass:: ReadingOrderPredictor diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index 7acabcd84a..8ffa55e437 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -9,11 +9,13 @@ import numpy as np +from doctr.utils.geometry import estimate_page_angle, order_points from doctr.utils.repr import NestedObject __all__ = [ "ReadingOrderPredictor", "assign_layout_labels", + "deskew_reading_geometries", "detect_text_direction", "layout_label_role", "normalize_layout_label", @@ -238,6 +240,8 @@ def _find(node: int) -> int: if last >= 0 else np.empty(0, dtype=int) ) + if candidates.size == 0 and last >= 0: + candidates = ready[y_overlap[last, ready] > y_overlap_threshold] if candidates.size == 0 and last >= 0 and not spanning[last]: # Continuation broken (gap, fragment): stay in the same column before switching to another same_column = ready[component[ready] == component[last]] @@ -289,6 +293,70 @@ def _attach_captions( return order +def deskew_reading_geometries( + geoms: Sequence[Any] | np.ndarray, + region_geoms: Sequence[Any] | np.ndarray | None = None, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, + min_angle: float = 1.0, +) -> tuple[list[Any], list[Any]]: + """De-skew rotated geometries so the reading order can be computed in an upright frame. + + Args: + geoms: geometries of the elements to order, in any docTR format (cf. `sort_reading_order`) + region_geoms: optional geometries of the layout regions, de-skewed together with the elements + page_shape: the page dimensions (height, width). Required for an exact angle on non-square pages with + relative coordinates, since relative coordinates distort angles by the aspect ratio. + angle_geoms: optional reading-oriented 4-point polygons used to estimate the page angle + min_angle: minimum estimated angle (in degrees) to trigger the de-skew. Ordering is affected as soon + as the drift along a line approaches the line height (about 1 degree for a page-wide line), and a + small rigid rotation cannot change the order of an upright page, hence the low default. Beyond 45 + degrees the corner identification is ambiguous (an upstream orientation correction is needed) and + nothing is done. + + Returns: + the (possibly de-skewed) element and region geometries + """ + region_geoms = list(region_geoms) if region_geoms is not None else [] + pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in geoms] + if len(pts) == 0 or any(p.shape[0] != 4 for p in pts): + return list(geoms), region_geoms # straight geometries: nothing to de-skew + height, width = page_shape if page_shape is not None else (1, 1) + scale = np.array([width, height], dtype=np.float64) + angle_source = angle_geoms if angle_geoms is not None else [] + angle_pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in angle_source] + if len(angle_pts) > 0 and all(p.shape[0] == 4 for p in angle_pts): + # Detection polygons are already reading-oriented (cf. `estimate_page_angle`): keep their vertex order + angle = estimate_page_angle(np.stack(angle_pts) * scale) + else: + # Normalize the vertex order (TL, TR, BR, BL) so the estimation does not depend on the vertex + angle = estimate_page_angle(np.stack([order_points(p * scale) for p in pts])) + if not np.isfinite(angle) or abs(angle) < min_angle or abs(angle) >= 45: + return list(geoms), region_geoms + # Rigid rotation zeroing the estimated angle; the center is irrelevant for ordering purposes + theta = np.deg2rad(angle) + rot = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) + all_pts = np.concatenate(pts, axis=0) * scale + center = all_pts.mean(axis=0) + + def _rotate(points: np.ndarray) -> np.ndarray: + return ((points * scale - center) @ rot.T + center) / scale + + deskewed = [_rotate(p) for p in pts] + + def _corners(points: np.ndarray) -> np.ndarray: + # Straight ((xmin, ymin), (xmax, ymax)) regions must be expanded to their 4 corners before rotating, + # otherwise only the diagonal would be rotated and the resulting extent would be underestimated + if points.shape[0] == 4: + return points + (x0, y0), (x1, y1) = points.min(axis=0), points.max(axis=0) + return np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) + + region_pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in region_geoms] + regions_out = [_rotate(_corners(p)) for p in region_pts] + return deskewed, regions_out + + def sort_reading_order( geoms: Sequence[Any] | np.ndarray, direction: str = "ltr", @@ -296,6 +364,8 @@ def sort_reading_order( x_overlap_threshold: float = 0.2, y_overlap_threshold: float = 0.5, caption_max_distance: float = 0.1, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, ) -> list[int]: """Compute the reading order of document elements from their geometries (and optionally, layout labels). @@ -324,12 +394,17 @@ def sort_reading_order( on the same visual row caption_max_distance: maximum relative distance between a caption and a float (table or figure) for the caption to be attached to it + page_shape: the page dimensions (height, width), used to de-skew rotated pages exactly (cf. + `deskew_reading_geometries`) + angle_geoms: optional reading-oriented 4-point polygons (typically the page's word polygons) used to + estimate the page angle on rotated pages (cf. `deskew_reading_geometries`) Returns: the permutation of the input indices which sorts the elements in reading order """ if direction not in SUPPORTED_DIRECTIONS[1:]: raise ValueError(f"invalid reading direction '{direction}', should be one of {SUPPORTED_DIRECTIONS[1:]}") + geoms, _ = deskew_reading_geometries(geoms, page_shape=page_shape, angle_geoms=angle_geoms) boxes = _to_boxes(geoms) num_boxes = boxes.shape[0] if labels is not None and len(labels) != num_boxes: @@ -368,6 +443,8 @@ def resolve_reading_segments( y_overlap_threshold: float = 0.5, caption_max_distance: float = 0.1, paragraph_gap: float = 0.8, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, ) -> list[list[int]]: """Order elements in reading order and group consecutive ones into segments (paragraphs or regions). @@ -388,11 +465,17 @@ def resolve_reading_segments( the caption to be attached to it paragraph_gap: maximum vertical gap between two consecutive elements to belong to the same segment, as a multiple of the median element height + page_shape: the page dimensions (height, width), used to de-skew rotated pages exactly (cf. + `deskew_reading_geometries`) + angle_geoms: optional reading-oriented 4-point polygons (typically the page's word polygons) used to + estimate the page angle on rotated pages (cf. `deskew_reading_geometries`) Returns: a partition of the input indices into reading-ordered segments (each segment being itself in reading order) """ + # De-skew once so the ordering and the vertical-gap measurement share the same upright frame + geoms, _ = deskew_reading_geometries(geoms, page_shape=page_shape, angle_geoms=angle_geoms) order = sort_reading_order( geoms, direction=direction, @@ -431,6 +514,8 @@ def assign_layout_labels( layout_geoms: Sequence[Any] | np.ndarray, layout_labels: Sequence[str], min_coverage: float = 0.5, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, ) -> list[str | None]: """Assign a layout label to each element based on its overlap with the detected layout regions. @@ -442,10 +527,16 @@ def assign_layout_labels( layout_geoms: geometries of the layout regions, in any docTR format layout_labels: labels of the layout regions (e.g. `[region.type for region in page.layout]`) min_coverage: minimum share of an element's area a region must cover to assign its label + page_shape: the page dimensions (height, width), used to de-skew rotated pages exactly (cf. + `deskew_reading_geometries`) + angle_geoms: optional reading-oriented 4-point polygons (typically the page's word polygons) used to + estimate the page angle on rotated pages (cf. `deskew_reading_geometries`) Returns: the label of each element (None when no region covers it enough) """ + # De-skew elements and regions together so the coverage is measured on tight boxes in the same frame + geoms, layout_geoms = deskew_reading_geometries(geoms, layout_geoms, page_shape=page_shape, angle_geoms=angle_geoms) boxes, regions = _to_boxes(geoms), _to_boxes(layout_geoms) if len(layout_labels) != regions.shape[0]: raise ValueError(f"Incompatible number of labels ({len(layout_labels)}) and regions ({regions.shape[0]})") @@ -523,6 +614,8 @@ def __call__( texts: Sequence[str] | None = None, labels: Sequence[str | None] | None = None, language: str | None = None, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, ) -> list[int]: """Compute the reading order of document elements. @@ -532,6 +625,10 @@ def __call__( direction detection labels: optional layout labels (one per geometry), used to handle page furniture & captions language: optional ISO 639 language code used as a fallback hint for the direction detection + page_shape: the page dimensions (height, width), used to de-skew rotated pages exactly on + non-square pages + angle_geoms: optional reading-oriented 4-point polygons (typically the page's word polygons) used + to estimate the page angle on rotated pages Returns: the permutation of the input indices which sorts the elements in reading order @@ -543,4 +640,6 @@ def __call__( x_overlap_threshold=self.x_overlap_threshold, y_overlap_threshold=self.y_overlap_threshold, caption_max_distance=self.caption_max_distance, + page_shape=page_shape, + angle_geoms=angle_geoms, ) diff --git a/tests/common/test_models_reading_order.py b/tests/common/test_models_reading_order.py index 7490b7f550..d4c6ffc2dd 100644 --- a/tests/common/test_models_reading_order.py +++ b/tests/common/test_models_reading_order.py @@ -4,6 +4,7 @@ from doctr.models.reading_order import ( ReadingOrderPredictor, assign_layout_labels, + deskew_reading_geometries, detect_text_direction, layout_label_role, normalize_layout_label, @@ -196,7 +197,82 @@ def test_reading_order_predictor(): ReadingOrderPredictor(direction="bottom-up") -# regression test for a bug where a stray fragment of a split line was read after the next line in the same column +def _rotated_box(box, deg, width=800, height=1000): + angle = np.deg2rad(deg) + rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) + center = np.array([width / 2, height / 2]) + (x0, y0), (x1, y1) = box + pts = np.array([ + [x0 * width, y0 * height], + [x1 * width, y0 * height], + [x1 * width, y1 * height], + [x0 * width, y1 * height], + ]) + return ((pts - center) @ rot.T + center) / [width, height] + + +def test_sort_reading_order_rotated_pages(): + title = [((0.1, 0.06), (0.9, 0.09))] + left = [((0.1, 0.12 + 0.05 * idx), (0.47, 0.15 + 0.05 * idx)) for idx in range(5)] + right = [((0.53, 0.12 + 0.05 * idx), (0.9, 0.15 + 0.05 * idx)) for idx in range(5)] + geoms = title + left + right + expected = list(range(11)) + for deg in (-35, -15, 15, 35): + rotated = [_rotated_box(box, deg) for box in geoms] + assert sort_reading_order(rotated) == expected + assert sort_reading_order(rotated, page_shape=(1000, 800)) == expected + # straight polygons are untouched (angle below the threshold, behavior identical to 2-point boxes) + straight = np.asarray([[(x0, y0), (x1, y0), (x1, y1), (x0, y1)] for ((x0, y0), (x1, y1)) in geoms]) + assert sort_reading_order(straight) == expected + + +def test_deskew_reading_geometries(): + geoms = [((0.1, 0.12), (0.47, 0.15)), ((0.53, 0.12), (0.9, 0.15))] + rotated = [_rotated_box(box, 25) for box in geoms] + # straight 2-point boxes are returned unchanged + out, regions = deskew_reading_geometries(geoms, [((0.0, 0.0), (1.0, 0.5))]) + assert out == list(geoms) and len(regions) == 1 + # rotated polygons are de-skewed: the two boxes end up on the same visual row + out, _ = deskew_reading_geometries(rotated, page_shape=(1000, 800)) + y_centers = [np.asarray(poly)[:, 1].mean() for poly in out] + assert abs(y_centers[0] - y_centers[1]) < 0.005 + # a straight region is expanded to its corners and rotated with the elements + out, regions = deskew_reading_geometries(rotated, [((0.0, 0.1), (1.0, 0.2))], page_shape=(1000, 800)) + assert np.asarray(regions[0]).shape == (4, 2) + # the operation is idempotent + again, _ = deskew_reading_geometries(out, page_shape=(1000, 800)) + assert all(np.allclose(a, b) for a, b in zip(out, again)) + # angle_geoms as the estimation source + out, _ = deskew_reading_geometries(rotated, page_shape=(1000, 800), angle_geoms=np.stack(rotated)) + y_centers = [np.asarray(poly)[:, 1].mean() for poly in out] + assert abs(y_centers[0] - y_centers[1]) < 0.005 + + +def test_reading_order_predictor_rotated(): + left = [_rotated_box(((0.1, 0.1 + 0.2 * idx), (0.3, 0.2 + 0.2 * idx)), 25) for idx in range(3)] + right = [_rotated_box(((0.6, 0.1 + 0.2 * idx), (0.8, 0.2 + 0.2 * idx)), 25) for idx in range(3)] + order = ReadingOrderPredictor()(left + right, page_shape=(1000, 800)) + assert order == [0, 1, 2, 3, 4, 5] + + +def test_deskew_strong_rotation_non_square_page(): + layout = [(80, 40, 670, 110), (80, 150, 360, 900), (390, 150, 670, 900)] # title + 2 columns + for height, width in [(1000, 750), (700, 2000)]: + sx, sy = width / 750, height / 1000 + for angle in (-44, 30, 44): + theta = np.deg2rad(angle) + rot = np.array([[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]]) + center = np.array([width / 2, height / 2]) + polys = [] + for x0, y0, x1, y1 in layout: + pts = np.array([[x0 * sx, y0 * sy], [x1 * sx, y0 * sy], [x1 * sx, y1 * sy], [x0 * sx, y1 * sy]]) + polys.append(((pts - center) @ rot.T + center) / np.array([width, height])) + assert sort_reading_order(polys, page_shape=(height, width)) == [0, 1, 2], (height, width, angle) + + +# Auto generated regression tests for known failures of the reading order algorithm + + def test_sort_reading_order_fragmented_columns(): left = [ ((0.10, 0.10), (0.45, 0.13)), # 0 wide @@ -210,3 +286,17 @@ def test_sort_reading_order_fragmented_columns(): order = sort_reading_order(left + right) # every left element (0..5) is read before every right element (6..11) assert max(order.index(i) for i in range(6)) < min(order.index(i) for i in range(6, 12)) + + +def test_fragmented_row_with_merged_column_components(): + geoms = [ + ((0.35, 0.05), (0.65, 0.10)), # 0 gutter-straddling element (bridges both columns) + ((0.10, 0.15), (0.45, 0.20)), # 1 left col, row 1 + ((0.10, 0.22), (0.16, 0.27)), # 2 left col, row 2, fragment A + ((0.17, 0.22), (0.24, 0.27)), # 3 left col, row 2, fragment B + ((0.25, 0.22), (0.45, 0.27)), # 4 left col, row 2, fragment C + ((0.10, 0.29), (0.45, 0.34)), # 5 left col, row 3 + ((0.55, 0.15), (0.90, 0.20)), # 6 right col, row 1 + ((0.55, 0.22), (0.90, 0.27)), # 7 right col, row 2 + ] + assert sort_reading_order(geoms) == [0, 1, 2, 3, 4, 5, 6, 7]