[reading order] de-skew rotated layout geoms

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

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 @@
                 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,13 +293,79 @@
     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",
     labels: Sequence[str | None] | None = None,
     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 @@
             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 @@
     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 @@
             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 @@
     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 @@
         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 @@
         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 @@
                 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 @@
             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,
         )

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]