From f1608a71b0e4c6c3a0aaaa12ee4ff368f4dfb39c Mon Sep 17 00:00:00 2001 From: felix Date: Mon, 12 Oct 2026 13:31:38 +0200 Subject: [PATCH 1/4] test reading order updates --- doctr/io/exporters.py | 61 +++++- doctr/models/builder.py | 246 ++++++++++++++++++---- doctr/models/reading_order/base.py | 112 +++++++++- tests/common/test_io_exporters.py | 107 ++++++++++ tests/common/test_models_builder.py | 157 ++++++++++++++ tests/common/test_models_reading_order.py | 76 +++++++ 6 files changed, 702 insertions(+), 57 deletions(-) diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py index 40961892fb..9a90535d22 100644 --- a/doctr/io/exporters.py +++ b/doctr/io/exporters.py @@ -89,6 +89,25 @@ def _covering_region_indices( return _covering_regions(_to_boxes(geoms), _to_boxes(region_geoms), min_coverage, region_labels) +def _merge_overlapping_regions(elt_regions: list[int], region_geoms: list[Any], region_labels: list[str]) -> list[int]: + """Map the layout region of each element to a group shared by the overlapping regions of the same label + + Args: + elt_regions: the layout region index of each element (-1 for none) + region_geoms: the geometries of the layout regions + region_labels: the labels of the layout regions + + Returns: + the region group of each element (-1 for none) + """ + from doctr.models.reading_order.base import _region_groups, _to_boxes + + if len(region_geoms) < 2: + return list(elt_regions) + groups = _region_groups(_to_boxes(region_geoms), region_labels) + return [groups[reg] if reg >= 0 else reg for reg in elt_regions] + + def _xyxy(geometry: Any) -> tuple[float, float, float, float]: """Return the enclosing box (xmin, ymin, xmax, ymax) of a box or a polygon""" pts = np.asarray(geometry, dtype=np.float64).reshape(-1, 2) @@ -155,7 +174,7 @@ def page_reading_order( normalize_layout_label, resolve_reading_segments, ) - from doctr.models.reading_order.base import _FURNITURE_ROLES, _to_boxes + from doctr.models.reading_order.base import _FURNITURE_ROLES, _float_text_label, _to_boxes signature = _reading_order_signature(page, direction) cached = getattr(page, "_reading_order_cache", None) @@ -212,8 +231,21 @@ def page_reading_order( column_voters = [ not (isinstance(elt, LayoutElement) and box[2] - box[0] > 0.5 * span) for elt, box in zip(elements, boxes) ] + # A text line lying in a table or a figure region keeps the label of its region, but it is not a float itself: + # it is grouped into paragraphs like regular text (and still read right after its figure) + order_labels = [ + _float_text_label(str(label)) + if not isinstance(elt, (Table, LayoutElement)) and layout_label_role(label) == "float" + else label + for elt, label in zip(elements, elt_labels) + ] segments = resolve_reading_segments( - elt_geoms, direction=direction, labels=elt_labels, column_voters=column_voters, region_groups=elt_regions + elt_geoms, + direction=direction, + labels=order_labels, + page_shape=page.dimensions, + column_voters=column_voters, + region_groups=_merge_overlapping_regions(elt_regions, region_geoms, region_labels), ) def _is_float_item(idx: int) -> bool: @@ -467,14 +499,19 @@ def _plan_figures( if not figure_idcs or not (render or not include_furniture): return _FigurePlan({}, set(), set()) - # Same upright frame as the reading order + # Same upright frame as the reading order (the lines of the blocks too) + line_owners = [idx for idx, item in enumerate(items) if isinstance(item, Block) for _ in item.lines] upright, upright_regions = deskew_reading_geometries( - [item.geometry for item in items], + [item.geometry for item in items] + + [line.geometry for item in items if isinstance(item, Block) for line in item.lines], [region.geometry for region in page.layout], page_shape=page.dimensions, angle_geoms=[word.geometry for block in page.blocks for line in block.lines for word in line.words], ) - boxes = [_xyxy(geom) for geom in upright] + boxes = [_xyxy(geom) for geom in upright[: len(items)]] + line_boxes: dict[int, list[tuple[float, ...]]] = {} + for owner, geom in zip(line_owners, upright[len(items) :]): + line_boxes.setdefault(owner, []).append(_xyxy(geom)) def _coverage(box: tuple[float, ...], target: tuple[float, ...]) -> float: """Share of `box` covered by `target`""" @@ -483,6 +520,10 @@ def _coverage(box: tuple[float, ...], target: tuple[float, ...]) -> float: ) return inter / max((box[2] - box[0]) * (box[3] - box[1]), 1e-9) + def _inside(idx: int, fig_idx: int) -> bool: + """Whether each line of a block lies in a figure""" + return all(_coverage(box, boxes[fig_idx]) >= 0.5 for box in line_boxes.get(idx, [boxes[idx]])) + def _figure_text(fig_idx: int) -> set[int]: """Blocks covered by a figure and labeled as part of it (or as furniture, for a figure in furniture)""" in_furniture = layout_label_role(labels[fig_idx]) in _FURNITURE_ROLES @@ -494,7 +535,7 @@ def _figure_text(fig_idx: int) -> set[int]: is_picture_label(labels[idx]) or (in_furniture and layout_label_role(labels[idx]) in _FURNITURE_ROLES) ) - and _coverage(boxes[idx], boxes[fig_idx]) >= 0.5 + and _inside(idx, fig_idx) } hidden: set[int] = set() @@ -512,15 +553,15 @@ def _figure_text(fig_idx: int) -> set[int]: for idx in figure_idcs: if sources[idx] is not None: hidden |= _figure_text(idx) - region_of = _covering_region_indices(upright, upright_regions, region_labels=[r.type for r in page.layout]) + region_of = _covering_region_indices( + upright[: len(items)], upright_regions, region_labels=[r.type for r in page.layout] + ) def _is_caption(idx: int) -> bool: return isinstance(items[idx], Block) and normalize_layout_label(labels[idx]) == "caption" def _is_inner_text(idx: int, fig_idx: int) -> bool: - return ( - isinstance(items[idx], Block) and not _is_caption(idx) and _coverage(boxes[idx], boxes[fig_idx]) >= 0.5 - ) + return isinstance(items[idx], Block) and not _is_caption(idx) and _inside(idx, fig_idx) def _caption_run(start: int) -> tuple[int, ...]: # A caption region can span several blocks diff --git a/doctr/models/builder.py b/doctr/models/builder.py index ca265300ea..118b5ac199 100644 --- a/doctr/models/builder.py +++ b/doctr/models/builder.py @@ -4,6 +4,8 @@ # See LICENSE or go to for full license details. +import bisect +import statistics from collections.abc import Sequence from typing import Any @@ -34,6 +36,13 @@ __all__ = ["DocumentBuilder"] +# Median word height of body text on an A4 page, relative to the page width: `paragraph_break` is scaled up on pages +# whose type is larger than that (e.g. a crop of a page) +_BODY_WORD_HEIGHT = 0.02 +# Cost of the horizontal distance between a box and a row, relative to their vertical offset, when several rows are +# level with the box: 50 units of horizontal distance weigh as much as 1 unit of vertical offset +_ROW_DISTANCE_WEIGHT = 0.02 + class DocumentBuilder(NestedObject): """Implements a document builder @@ -41,7 +50,9 @@ class DocumentBuilder(NestedObject): Args: resolve_lines: whether words should be automatically grouped into lines resolve_blocks: whether lines should be automatically grouped into blocks - paragraph_break: relative length of the minimum space separating paragraphs + paragraph_break: relative length of the minimum space separating paragraphs. It is scaled up on pages with + large type (e.g. a crop of a page), and lines can be split at a smaller distance adapted to the word + spacing of the page. export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. keep_reading_order: if True, arrange the content of every page in reading order (cf. @@ -105,7 +116,11 @@ def _sort_boxes(boxes: np.ndarray, shape: tuple[int, int] | None = None) -> tupl return (boxes[:, 0] + 2 * boxes[:, 3] / med_height).argsort(), boxes def _resolve_sub_lines( - self, boxes: np.ndarray, word_idcs: list[int], break_dist: float | None = None + self, + boxes: np.ndarray, + word_idcs: list[int], + break_dist: float | None = None, + word_regions: tuple[np.ndarray, np.ndarray] | None = None, ) -> list[list[int]]: """Split a line in sub_lines @@ -113,6 +128,10 @@ def _resolve_sub_lines( boxes: bounding boxes of shape (N, 4) word_idcs: list of indexes for the words of the line break_dist: horizontal distance above which the line is split (defaults to `self.paragraph_break`) + word_regions: optional layout region (group) of each box (-1 for none) and coverage of each box by each + region group (cf. `_word_regions`): the line is also split between two boxes of distinct regions (e.g. + two columns separated by a gutter as narrow as a word gap), unless one also lies in the region of the + other (nested or overlapping regions) Returns: A list of (sub-)lines computed from the original line (words) @@ -127,31 +146,186 @@ def _resolve_sub_lines( lines.append(word_idcs) else: sub_line = [word_idcs[0]] + # Right edge of the current sub-line: a box can end before an earlier, wider one (e.g. a punctuation box + # nested in a word box), so the running max is tracked instead of the right edge of the last box + right = boxes[word_idcs[0], 2] + groups, coverage = word_regions if word_regions is not None else (None, None) + region = -1 if groups is None else int(groups[word_idcs[0]]) # region of the sub-line + prev = word_idcs[0] for i in word_idcs[1:]: horiz_break = True - prev_box = boxes[sub_line[-1]] - # Compute distance between boxes - dist = boxes[i, 0] - prev_box[2] + # Compute distance between the sub-line and the box + dist = boxes[i, 0] - right # If distance between boxes is lower than the break distance, same sub-line if dist < break_dist: horiz_break = False + # Unless the box lies in another layout region than the sub-line, and neither the box lies in the + # region of the sub-line nor the previous box in the region of the box + box_region = -1 if groups is None else int(groups[i]) + if ( + coverage is not None + and box_region >= 0 + and region >= 0 + and box_region != region + and coverage[i, region] < 0.5 + and coverage[prev, box_region] < 0.5 + ): + horiz_break = True if horiz_break: lines.append(sub_line) sub_line = [] + right = boxes[i, 2] + region = -1 sub_line.append(i) + right = max(right, boxes[i, 2]) + if region < 0: + region = box_region + prev = i lines.append(sub_line) return lines - def _resolve_lines(self, boxes: np.ndarray, shape: tuple[int, int] | None = None) -> list[list[int]]: + @staticmethod + def _word_regions( + boxes: np.ndarray, + region_geoms: list[Any], + region_labels: list[str], + shape: tuple[int, int] | None = None, + ) -> tuple[np.ndarray, np.ndarray] | None: + """Assign each box to the layout region covering it, overlapping regions of the same label being grouped + + A box covered by no region by half (e.g. a word sticking out of a region slightly too short) is assigned to + the smallest region holding it horizontally and covering it at least partly vertically. + + Args: + boxes: bounding boxes of shape (N, 4) or (N, 4, 2) + region_geoms: the geometries of the layout regions + region_labels: the labels of the layout regions + shape: the page dimensions (height, width), used to de-skew rotated pages exactly + + Returns: + the region group of each box (-1 for none) and the coverage of each box by each region group (indexed by + group), None without layout regions + """ + from doctr.models.reading_order.base import ( + _covering_regions, + _region_groups, + _to_boxes, + deskew_reading_geometries, + ) + + if len(region_geoms) == 0 or boxes.shape[0] == 0: + return None + word_geoms: Any = boxes + if boxes.ndim == 3: # de-skew the words and the regions together + word_geoms, region_geoms = deskew_reading_geometries( + list(boxes), region_geoms, page_shape=shape, angle_geoms=list(boxes) + ) + regions, words = _to_boxes(region_geoms), _to_boxes(word_geoms) + groups = np.asarray(_region_groups(regions, region_labels)) + covering = np.asarray(_covering_regions(words, regions, 0.5, region_labels)) + # Coverage of each box by each region, on both axes and in area + inter_w = np.minimum(words[:, None, 2], regions[None, :, 2]) - np.maximum( + words[:, None, 0], regions[None, :, 0] + ) + inter_h = np.minimum(words[:, None, 3], regions[None, :, 3]) - np.maximum( + words[:, None, 1], regions[None, :, 1] + ) + x_cov = np.clip(inter_w, 0, None) / np.clip(words[:, 2:3] - words[:, 0:1], 1e-9, None) + y_cov = np.clip(inter_h, 0, None) / np.clip(words[:, 3:4] - words[:, 1:2], 1e-9, None) + uncovered = np.flatnonzero(covering < 0) + if uncovered.size > 0: + areas = (regions[:, 2] - regions[:, 0]) * (regions[:, 3] - regions[:, 1]) + holding = (x_cov[uncovered] >= 0.9) & (y_cov[uncovered] >= 0.3) + smallest = np.where(holding, areas[None, :], np.inf).argmin(axis=1) + covering[uncovered] = np.where(holding.any(axis=1), smallest, -1) + coverage = np.zeros((words.shape[0], regions.shape[0])) + np.maximum.at(coverage.T, groups, (x_cov * y_cov).T) # coverage by a group: the max over its regions + return np.where(covering >= 0, groups[covering], -1), coverage + + @staticmethod + def _resolve_rows(boxes: np.ndarray, idxs: np.ndarray, y_med: float) -> list[list[int]]: + """Group straight boxes into visual rows + + A box can join a row if its vertical center is close enough (less than half a median box height) to the median + center of its horizontally closest boxes in that row. Comparing a box with its neighbors rather than with the + mean of the whole row keeps a slanted line in a single row, and prevents a word from being drawn into the row + of another column whose baseline is slightly offset. When several rows are level with the box, the best + aligned one wins, the horizontal distance to the row being added as a small penalty. + + Args: + boxes: straight bounding boxes of shape (N, 4) + idxs: the order in which the boxes are visited (cf. `_sort_boxes`: top to bottom, left to right) + y_med: the median box height + + Returns: + the rows (lists of box indices), from top to bottom + """ + if not y_med > 0: # degenerate heights: no box is close enough to another one + return [[idx] for idx in idxs.tolist()] + lefts, rights = boxes[:, 0].tolist(), boxes[:, 2].tolist() + centers = ((boxes[:, 1] + boxes[:, 3]) / 2).tolist() + rows: list[list[int]] = [] + row_lefts: list[list[float]] = [] # left edges of the boxes of each row, sorted + row_boxes: list[list[int]] = [] # boxes of each row, sorted by left edge + highest: list[float] = [] # highest box center of each row + lowest: list[float] = [] # lowest box center of each row + active: list[int] = [] # rows which can still receive a box + for idx in idxs.tolist(): + center, left, right = centers[idx], lefts[idx], rights[idx] + # Boxes come roughly top to bottom: a row whose lowest box is far above can no longer be joined + active = [row for row in active if lowest[row] > center - 2 * y_med] + best, best_key = -1, (float("inf"), float("inf")) + for row in active: + # The neighbors of the box can only be level with it if the row spans its center + if highest[row] >= center + y_med / 2 or lowest[row] <= center - y_med / 2: + continue + # The horizontally closest boxes of the row lie around the position of the box in that row. A few of + # them rather than a single one, so that a small mark (accent, punctuation) can't steer the box. + pos = bisect.bisect_left(row_lefts[row], left) + nearest = sorted( + (max(lefts[other] - right, left - rights[other], 0.0), other) + for other in row_boxes[row][max(pos - 3, 0) : pos + 3] + )[:3] + offset = abs(statistics.median(centers[other] for _, other in nearest) - center) + # On a slanted line, the next line can start before the end of the current one: the offset alone + # would let a distant row of similar height take the box + key = (offset + _ROW_DISTANCE_WEIGHT * nearest[0][0], nearest[0][0]) + if offset < y_med / 2 and key < best_key: + best, best_key = row, key + if best < 0: + best = len(rows) + rows.append([]) + row_lefts.append([]) + row_boxes.append([]) + highest.append(center) + lowest.append(center) + active.append(best) + rows[best].append(idx) + pos = bisect.bisect_right(row_lefts[best], left) + row_lefts[best].insert(pos, left) + row_boxes[best].insert(pos, idx) + highest[best] = min(highest[best], center) + lowest[best] = max(lowest[best], center) + + return rows + + def _resolve_lines( + self, + boxes: np.ndarray, + shape: tuple[int, int] | None = None, + word_regions: tuple[np.ndarray, np.ndarray] | None = None, + ) -> list[list[int]]: """Order boxes to group them in lines Args: boxes: bounding boxes of shape (N, 4) or (N, 4, 2) in case of rotated bbox shape: the page dimensions (height, width), used to de-skew rotated pages exactly + word_regions: optional layout region (group) of each box and coverage by each group (cf. `_word_regions`): + a line never spans two distinct regions Returns: nested list of box indices @@ -163,31 +337,7 @@ def _resolve_lines(self, boxes: np.ndarray, shape: tuple[int, int] | None = None y_med = np.median(boxes[:, 3] - boxes[:, 1]) # Group the words into visual rows - rows = [] - words = [idxs[0]] # Assign the top-left word to the first line - # Define a mean y-center for the line - y_center_sum = boxes[idxs[0]][[1, 3]].mean() - - for idx in idxs[1:]: - vert_break = True - - # Compute y_dist - y_dist = abs(boxes[idx][[1, 3]].mean() - y_center_sum / len(words)) - # If y-center of the box is close enough to mean y-center of the line, same line - if y_dist < y_med / 2: - vert_break = False - - if vert_break: - rows.append(words) - words = [] - y_center_sum = 0 - - words.append(idx) - y_center_sum += boxes[idx][[1, 3]].mean() - - # Use the remaining words to form the last row - if len(words) > 0: - rows.append(words) + rows = self._resolve_rows(boxes, idxs, float(y_med)) # Split the rows horizontally, using a break distance adapted to the page spacing. # The gaps are collected with one vectorized pass per row (the loop is over rows, not word pairs). @@ -199,22 +349,25 @@ def _resolve_lines(self, boxes: np.ndarray, shape: tuple[int, int] | None = None row_idcs = np.asarray(row) row_idcs = row_idcs[np.argsort(boxes[row_idcs, 0], kind="stable")] n_pairs += row_idcs.shape[0] - 1 - gap_chunks.append(boxes[row_idcs[1:], 0] - boxes[row_idcs[:-1], 2]) + # Same gap definition as `_resolve_sub_lines`: distance to the right edge of everything on the left + gap_chunks.append(boxes[row_idcs[1:], 0] - np.maximum.accumulate(boxes[row_idcs, 2])[:-1]) all_gaps = np.concatenate(gap_chunks) if gap_chunks else np.empty(0, dtype=boxes.dtype) pos_gaps = all_gaps[all_gaps > 0] # Median word height, converted to width units aspect = (shape[0] / shape[1]) if shape is not None else 1.0 floor = float(y_med) * aspect + # The width-relative paragraph break grows with the type size (large type on a narrow page or crop) + cap = self.paragraph_break * max(1.0, floor / _BODY_WORD_HEIGHT) if pos_gaps.shape[0] >= 5 and pos_gaps.shape[0] >= 0.5 * n_pairs: - break_dist = min(self.paragraph_break, max(3.0 * float(np.median(pos_gaps)), floor)) + break_dist = min(cap, max(3.0 * float(np.median(pos_gaps)), floor)) elif n_pairs >= 5: - break_dist = min(self.paragraph_break, floor) + break_dist = min(cap, floor) else: break_dist = self.paragraph_break lines = [] for row in rows: - lines.extend(self._resolve_sub_lines(boxes, row, break_dist)) + lines.extend(self._resolve_sub_lines(boxes, row, break_dist, word_regions)) return lines @@ -587,6 +740,7 @@ def _build_blocks( word_preds: list[tuple[str, float]], crop_orientations: list[dict[str, Any]], shape: tuple[int, int] | None = None, + layout: list[LayoutElement] | None = None, ) -> list[Block]: """Gather independent words in structured blocks @@ -598,6 +752,7 @@ def _build_blocks( the general orientation (orientations + confidences) of the crops shape: the page dimensions (height, width), used to de-skew rotated pages exactly when resolving lines + layout: optional layout regions of the page: a line never spans two of them Returns: list of block elements @@ -611,7 +766,12 @@ def _build_blocks( # Decide whether we try to form lines _boxes = boxes if self.resolve_lines: - lines = self._resolve_lines(_boxes if _boxes.ndim == 3 else _boxes[:, :4], shape) + word_regions = ( + self._word_regions(_boxes, [elt.geometry for elt in layout], [elt.type for elt in layout], shape) + if layout + else None + ) + lines = self._resolve_lines(_boxes if _boxes.ndim == 3 else _boxes[:, :4], shape, word_regions) # Decide whether we try to form blocks if self.resolve_blocks and len(lines) > 1: _blocks = self._resolve_blocks(_boxes if _boxes.ndim == 3 else _boxes[:, :4], lines) @@ -673,7 +833,11 @@ def _apply_reading_order(self, page: Page) -> None: words = [word for block in page.blocks for line in block.lines for word in line.words] if len(words) < 2: return - groups = self._resolve_lines(self._words_to_boxes(words), page.dimensions) + word_boxes = self._words_to_boxes(words) + word_regions = self._word_regions( + word_boxes, [elt.geometry for elt in page.layout], [elt.type for elt in page.layout], page.dimensions + ) + groups = self._resolve_lines(word_boxes, page.dimensions, word_regions) page.blocks = [Block([Line([words[idx] for idx in group]) for group in groups])] # With the figures: the result is stored back in the cache read by the exports @@ -788,12 +952,14 @@ def __call__( word_preds = [wp for wp, k in zip(word_preds, keep) if k] word_crop_orientations = [co for co, k in zip(word_crop_orientations, keep) if k] + page_layout = self._build_layout_elements(page_regions) page_blocks = self._build_blocks( page_boxes, loc_scores, word_preds, word_crop_orientations, shape, + page_layout, ) _page = Page( page, @@ -802,7 +968,7 @@ def __call__( shape, orientation, language, - self._build_layout_elements(page_regions), + page_layout, page_tables, ) if self.keep_reading_order: @@ -819,7 +985,9 @@ class KIEDocumentBuilder(DocumentBuilder): Args: resolve_lines: whether words should be automatically grouped into lines resolve_blocks: whether lines should be automatically grouped into blocks - paragraph_break: relative length of the minimum space separating paragraphs + paragraph_break: relative length of the minimum space separating paragraphs. It is scaled up on pages with + large type (e.g. a crop of a page), and lines can be split at a smaller distance adapted to the word + spacing of the page. export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. keep_reading_order: if True, arrange the content of every page in reading order (cf. diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index 7e27fe610d..b87e311e6e 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -44,6 +44,10 @@ _FURNITURE_ROLES = ("header", "footer", "footnote") +# Label suffix of the text lying in a float region without float element (e.g. 'Table text' for the text of a table +# region without table structure): it is ordered and grouped like body text, but moved after a region like a float +_FLOAT_TEXT_SUFFIX = " text" + def normalize_layout_label(label: str | None) -> str: """Normalize a layout label to lower snake case (e.g. 'Page-header' -> 'page_header') @@ -82,6 +86,17 @@ def layout_label_role(label: str | None) -> str: return "body" +def _float_text_label(label: str) -> str: + """Return the label of a text element lying in a float region (e.g. 'Table' -> 'Table text')""" + return f"{label}{_FLOAT_TEXT_SUFFIX}" + + +def _is_float_text(label: str | None) -> bool: + """Check whether a label is the one of a text element lying in a float region (cf. `_float_text_label`)""" + norm = normalize_layout_label(label) + return any(norm == normalize_layout_label(_float_text_label(float_label)) for float_label in _FLOAT_LABELS) + + def detect_text_direction(texts: Iterable[str], language: str | None = None) -> str: """Infer the horizontal base direction of a document from its text content. @@ -215,6 +230,16 @@ def _topological_order( same_row = y_overlap > y_overlap_threshold same_row &= ~x_linked same_row &= x_rank[:, None] < x_rank[None, :] # i is on the left of j + # An element belongs to the row it overlaps the most: an element shifted between two rows (e.g. an amount on a + # crumpled receipt) must not wait for the element on the left of the next row. Only the elements of another row + # than its best partner are dropped: a slightly lower cell of its own row is still read before it. + best_left = np.max(y_overlap, axis=0, where=same_row, initial=0.0) + weak_i, weak_j = np.nonzero(same_row & (y_overlap < best_left[None, :] - 0.1)) + if weak_i.size > 0: + cols = np.unique(weak_j) + best = np.argmax(np.where(same_row[:, cols], y_overlap[:, cols], -1.0), axis=0) + other_row = y_overlap[weak_i, best[np.searchsorted(cols, weak_j)]] <= y_overlap_threshold + same_row[weak_i[other_row], weak_j[other_row]] = False edges |= same_row del same_row np.fill_diagonal(edges, False) @@ -375,7 +400,7 @@ def _defer_floats( inner = {idx for idcs in members.values() for idx in idcs} def _is_float(idx: int) -> bool: - return labels[idx] in _FLOAT_LABELS or idx in inner + return labels[idx] in _FLOAT_LABELS or _is_float_text(labels[idx]) or idx in inner out = list(order) pos = 0 @@ -656,12 +681,14 @@ def resolve_reading_segments( angle_geoms: Sequence[Any] | np.ndarray | None = None, column_voters: Sequence[bool] | None = None, region_groups: Sequence[int] | None = None, + row_gap: float = 1.5, ) -> list[list[int]]: """Order elements in reading order and group consecutive ones into segments (paragraphs or regions). This is typically used to regroup individual lines into paragraphs once they have been put in reading order: two consecutive elements are merged into the same segment when they share the same (possibly None) - layout label, belong to the same column and are vertically close enough. + layout label and layout region, and either belong to the same column and are vertically close enough, or + continue the same visual row (e.g. the fragments of a line split at a wide word gap). Args: geoms: sequence of geometries in any docTR format (cf. `sort_reading_order`) @@ -682,7 +709,12 @@ def resolve_reading_segments( estimate the page angle on rotated pages (cf. `deskew_reading_geometries`) column_voters: optional mask of the elements used to detect a multi-column page (all by default) region_groups: optional layout region index of each element (-1 for none), used to attach the captions of - a region as a unit and to read a float standing beside a region after it + a region as a unit, to read a float standing beside a region after it, and to keep distinct regions in + distinct segments + row_gap: maximum horizontal gap between two consecutive elements of the same visual row to belong to the + same segment, as a multiple of the median element height (pass `page_shape` along with relative + coordinates of a non-square page, so that both are measured in the same unit). It does not apply to two + elements of body text in the same layout region. Returns: a partition of the input indices into reading-ordered segments (each segment being itself in @@ -704,24 +736,52 @@ def resolve_reading_segments( return [] canonical = _to_canonical_ltr(_to_boxes(geoms), direction) median_height = float(np.median(canonical[:, 3] - canonical[:, 1])) + # Relative coordinates distort the horizontal distances on non-square pages: scale them to the height unit + aspect = 1.0 + if page_shape is not None and page_shape[0] > 0 and page_shape[1] > 0: + aspect = page_shape[0] / page_shape[1] if direction.startswith("ttb") else page_shape[1] / page_shape[0] norm_labels = [normalize_layout_label(labels[idx] if labels is not None else None) for idx in range(len(order))] segments: list[list[int]] = [[order[0]]] + # Extent of the trailing visual row of the current segment (the fragments of a line split at wide word gaps) + row = canonical[order[0]].copy() for prev, cur in zip(order[:-1], order[1:]): - x_gap = min(canonical[prev, 2], canonical[cur, 2]) - max(canonical[prev, 0], canonical[cur, 0]) - x_overlap = x_gap / max( - min(canonical[prev, 2] - canonical[prev, 0], canonical[cur, 2] - canonical[cur, 0]), 1e-9 + box = canonical[cur] + x_gap = min(row[2], box[2]) - max(row[0], box[0]) + x_overlap = x_gap / max(min(row[2] - row[0], box[2] - box[0]), 1e-9) + y_inter = min(canonical[prev, 3], box[3]) - max(canonical[prev, 1], box[1]) + y_overlap = y_inter / max(min(canonical[prev, 3] - canonical[prev, 1], box[3] - box[1]), 1e-9) + stacked = x_overlap > x_overlap_threshold and box[1] - row[3] <= paragraph_gap * median_height + same_region = region_groups is None or region_groups[prev] == region_groups[cur] + # The next element on the right in the same visual row, e.g. the remainder of a line split at a wide word gap + # (whatever the gap within a region of body text: the layout model already grouped both) + same_row = ( + y_overlap > y_overlap_threshold + and box[0] > canonical[prev, 0] + and ( + (box[0] - canonical[prev, 2]) * aspect <= row_gap * median_height + or ( + region_groups is not None + and region_groups[cur] >= 0 + and same_region + and layout_label_role(norm_labels[cur]) == "body" + ) + ) ) same_segment = ( norm_labels[prev] == norm_labels[cur] and layout_label_role(norm_labels[cur]) != "float" - and x_overlap > x_overlap_threshold - and canonical[cur, 1] - canonical[prev, 3] <= paragraph_gap * median_height + and same_region + and (stacked or same_row) ) if same_segment: segments[-1].append(cur) else: segments.append([cur]) + if same_segment and same_row and not stacked: + row = np.concatenate((np.minimum(row[:2], box[:2]), np.maximum(row[2:], box[2:]))) + else: + row = box.copy() return segments @@ -767,6 +827,42 @@ def assign_layout_labels( ] +def _region_groups(regions: np.ndarray, labels: Sequence[str | None], min_overlap: float = 0.1) -> list[int]: + """Group the overlapping layout regions of the same label + + Layout models can return overlapping regions for a single paragraph (e.g. a second region over the part of a + paragraph wrapping around a figure), which must not split it. + + Args: + regions: the (R, 4) boxes of the layout regions + labels: the label of each region + min_overlap: minimum intersection of two regions, relative to the area of the smaller one, to group them + + Returns: + the group of each region (the index of one of the regions of the group) + """ + norm = [normalize_layout_label(label) for label in labels] + areas = (regions[:, 2] - regions[:, 0]) * (regions[:, 3] - regions[:, 1]) + parent = list(range(regions.shape[0])) + + def _find(idx: int) -> int: + while parent[idx] != idx: + parent[idx] = parent[parent[idx]] + idx = parent[idx] + return idx + + for i in range(regions.shape[0]): + for j in range(i + 1, regions.shape[0]): + if norm[i] != norm[j]: + continue + inter = max(min(regions[i, 2], regions[j, 2]) - max(regions[i, 0], regions[j, 0]), 0.0) * max( + min(regions[i, 3], regions[j, 3]) - max(regions[i, 1], regions[j, 1]), 0.0 + ) + if inter > min_overlap * min(areas[i], areas[j]): + parent[_find(i)] = _find(j) + return [_find(idx) for idx in range(regions.shape[0])] + + def _covering_regions( boxes: np.ndarray, regions: np.ndarray, diff --git a/tests/common/test_io_exporters.py b/tests/common/test_io_exporters.py index ac0b94faef..f7adb77998 100644 --- a/tests/common/test_io_exporters.py +++ b/tests/common/test_io_exporters.py @@ -128,6 +128,113 @@ def test_page_export_as_markdown_rtl(): assert page.export_as_markdown(direction="ltr").startswith("الأيسر") +def test_page_reading_order_text_in_a_table_region(): + # A table region from the layout model, without table structure (`detect_tables=False`): its lines keep the + # region label, but they are not floats themselves, so they are neither nested into each other nor isolated + lines = [ + _line_at("Pos Datum Einzelpreis Gesamt", 0.10, 0.40, 0.90, 0.44), # header line merged across the cells + _line_at("pro UE", 0.55, 0.425, 0.75, 0.46), # second header row, under "Einzelpreis" + _line_at("EUR", 0.80, 0.42, 0.88, 0.45), # second header row, under "Gesamt": mostly inside the first line + _line_at("1 Deutsch 22,50 22,50", 0.10, 0.48, 0.90, 0.52), + ] + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), + [elements.Block(lines=lines)], + 0, + (1000, 800), + layout=[elements.LayoutElement("Table", 0.9, ((0.08, 0.38), (0.92, 0.54)))], + ) + items, labels, _ = page_reading_order(page) + texts = [" ".join(word.value for word in line.words) for item in items for line in item.lines] + assert texts == ["Pos Datum Einzelpreis Gesamt", "pro UE", "EUR", "1 Deutsch 22,50 22,50"] + assert labels == ["Table"] * len(items) + assert len(items) < len(lines) + + +def _texts(items): + return [[" ".join(word.value for word in line.words) for line in item.lines] for item in items] + + +def test_page_reading_order_table_region_beside_a_paragraph(): + # The rows of a table region (without table structure) level with the lines of the paragraph beside it are read + # after that paragraph, not interleaved with its lines + lines = [_line_at(f"intro line {i} of the body", 0.08, 0.05 + 0.04 * i, 0.92, 0.08 + 0.04 * i) for i in range(6)] + lines += [ + _line_at(f"para line {i} beside the table", 0.08, 0.30 + 0.04 * i, 0.55, 0.33 + 0.04 * i) for i in range(4) + ] + lines += [_line_at(f"Row{i} {i}0", 0.62, 0.30 + 0.04 * i, 0.92, 0.33 + 0.04 * i) for i in range(4)] + lines += [_line_at(f"outro line {i} of the body", 0.08, 0.50 + 0.04 * i, 0.92, 0.53 + 0.04 * i) for i in range(6)] + layout = [ + elements.LayoutElement("Text", 0.9, ((0.06, 0.04), (0.94, 0.28))), + elements.LayoutElement("Text", 0.9, ((0.06, 0.29), (0.57, 0.46))), + elements.LayoutElement("Table", 0.9, ((0.60, 0.29), (0.94, 0.46))), + elements.LayoutElement("Text", 0.9, ((0.06, 0.49), (0.94, 0.74))), + ] + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), layout=layout + ) + items, labels, _ = page_reading_order(page) + texts = [text for item in _texts(items) for text in item] + paragraph = [texts.index(f"para line {i} beside the table") for i in range(4)] + rows = [texts.index(f"Row{i} {i}0") for i in range(4)] + assert paragraph == list(range(paragraph[0], paragraph[0] + 4)) + assert rows == list(range(paragraph[-1] + 1, paragraph[-1] + 5)) + assert {label for item, label in zip(_texts(items), labels) if item[0].startswith("Row")} == {"Table"} + + +def test_page_reading_order_overlapping_text_regions(): + # The layout model returns a second, overlapping region for the lines of a paragraph wrapping around a figure: + # the paragraph is not split where the lines switch from one region to the other + lines = [ + _line_at("one two three four five six", 0.05, 0.10, 0.90, 0.13), + _line_at("seven eight nine ten eleven", 0.05, 0.14, 0.90, 0.17), + _line_at("twelve thirteen fourteen fifteen", 0.05, 0.18, 0.90, 0.21), + _line_at("sixteen seventeen", 0.05, 0.22, 0.50, 0.25), + _line_at("eighteen nineteen", 0.05, 0.26, 0.50, 0.29), + _line_at("twenty end.", 0.05, 0.30, 0.50, 0.33), + ] + layout = [ + elements.LayoutElement("Text", 0.90, ((0.04, 0.09), (0.91, 0.215))), + elements.LayoutElement("Text", 0.59, ((0.04, 0.16), (0.51, 0.34))), + elements.LayoutElement("Picture", 0.86, ((0.55, 0.22), (0.90, 0.34))), + ] + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1810, 1280), layout=layout + ) + items, _, _ = page_reading_order(page) + assert len(items) == 1 and len(items[0].lines) == 6 + + +def test_page_export_hides_figure_text_sticking_out_as_a_block(): + # Each line of the figure text lies in the figure, but the box of the block they form sticks out of it + lines = [ + _line_at("Body text above the chart spanning the page", 0.05, 0.05, 0.95, 0.09), + _line_at("AXIS", 0.30, 0.185, 0.40, 0.225), + _line_at("LABEL", 0.45, 0.20, 0.45 + 0.299 / 0.9, 0.24), + _line_at("Body text below the chart spanning the page", 0.05, 0.70, 0.95, 0.74), + ] + page = elements.Page( + np.full((100, 80, 3), 200, dtype=np.uint8), + [elements.Block(lines=lines)], + 0, + (1000, 800), + layout=[elements.LayoutElement("Picture", 0.9, ((0.20, 0.20), (0.60, 0.60)))], + ) + markdown = page.export_as_markdown(images="embedded") + assert "![" in markdown + assert "AXIS" not in markdown and "LABEL" not in markdown + + +def test_page_reading_order_degenerate_dimensions(): + lines = [_line_at("first line here", 0.1, 0.1, 0.5, 0.13), _line_at("second part", 0.6, 0.1, 0.9, 0.13)] + for dims in [(0, 0), (0, 100), (100, 0)]: + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, dims) + assert [text for item in _texts(page_reading_order(page)[0]) for text in item] == [ + "first line here", + "second part", + ] + + def test_page_export_with_tables(): cells = [ elements.TableCell("Name", 0.9, ((0.1, 0.55), (0.4, 0.6)), 0, 0, 0, 0), diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index b9636f47b4..641a47e4c6 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -449,6 +449,163 @@ def test_resolve_lines(input_boxes, lines): assert doc_builder._resolve_lines(np.asarray(input_boxes)) == lines +def test_resolve_lines_offset_columns(): + # Two columns whose baselines are offset by half a line (scanned page): the right column line must not lose a + # word to the next line of the left column + boxes = [ + [0.1422, 0.5908, 0.2472, 0.6045], # 0-4: left column, "Nichttragende Wände bestehen aus Kalk-" + [0.261, 0.5898, 0.3176, 0.6035], + [0.3314, 0.5898, 0.3978, 0.6035], + [0.413, 0.5928, 0.4406, 0.6035], + [0.4558, 0.5889, 0.4972, 0.6035], + [0.1394, 0.6055, 0.2099, 0.6182], # 5-10: left column, "sandstein oder Vollwanddielen bzw. aus stabilen" + [0.2126, 0.6064, 0.2486, 0.6172], + [0.2513, 0.6055, 0.366, 0.6191], + [0.3674, 0.6045, 0.4047, 0.6191], + [0.4074, 0.6084, 0.4364, 0.6191], + [0.4378, 0.6055, 0.4972, 0.6191], + [0.5539, 0.5811, 0.5953, 0.5957], # 11-13: right column, "Holz- oder Stahlkonstruktion." + [0.5953, 0.5811, 0.634, 0.5947], + [0.634, 0.5811, 0.7694, 0.5977], + [0.5539, 0.5967, 0.6755, 0.6133], # 14-18: right column, "Hauseingangstür als Holz-/Glaselement mit Si-/" + [0.6782, 0.5977, 0.6989, 0.6104], + [0.7045, 0.5977, 0.8426, 0.6113], + [0.844, 0.5977, 0.8758, 0.6133], + [0.8772, 0.5977, 0.9089, 0.6123], + ] + lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), (2480, 1753)) + assert sorted(lines) == [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 10], [11, 12, 13], [14, 15, 16, 17, 18]] + + +def test_resolve_lines_slanted_line(): + # A line rising to the right (perspective) by twice its height stays a single line, in reading order + boxes = [[0.05 + 0.11 * i, 0.485 - 0.007 * i, 0.15 + 0.11 * i, 0.515 - 0.007 * i] for i in range(8)] + boxes += [[0.05 + 0.11 * i, 0.70, 0.15 + 0.11 * i, 0.73] for i in range(8)] + assert builder.DocumentBuilder()._resolve_lines(np.asarray(boxes)) == [list(range(8)), list(range(8, 16))] + + +@pytest.mark.parametrize("slope", [0.0045, -0.0045]) # down / up to the right (about 2 degrees of skew) +def test_resolve_lines_slanted_lines_with_tight_leading(slope): + # The next line starts before the end of the current one: its first word must not take the last words of the + # current line, although it is better aligned with them than their own (lagging) neighbors + boxes = [ + [0.05 + 0.11 * i, y + slope * i, 0.15 + 0.11 * i, y + 0.03 + slope * i] for y in (0.30, 0.339) for i in range(8) + ] + lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes)) + assert sorted(lines) == [list(range(8)), list(range(8, 16))] + + +def test_resolve_lines_level_rows_of_other_columns(): + # 3 columns whose baselines are slightly offset: the words of the last column stay together, although each of + # them is level with the row of another column + boxes = [ + [0.05, 0.190, 0.12, 0.210], # 0-2: column A + [0.13, 0.190, 0.20, 0.210], + [0.21, 0.190, 0.26, 0.210], + [0.36, 0.2005, 0.44, 0.2205], # 3-5: column B + [0.45, 0.2005, 0.50, 0.2205], + [0.51, 0.2005, 0.56, 0.2205], + [0.66, 0.194, 0.74, 0.214], # 6-7: column C + [0.75, 0.197, 0.83, 0.217], + ] + lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes)) + assert sorted(lines) == [[0, 1, 2], [3, 4, 5], [6, 7]] + + +def test_resolve_lines_layout_regions(): + # 3 columns separated by gutters as narrow as the word gaps of the justified lines + boxes, regions = [], [] + for col in range(3): + x0 = 0.02 + 0.33 * col + regions.append(((x0 - 0.01, 0.05), (x0 + 0.31, 0.95))) + for row in range(2): + for word in range(3): + boxes.append([x0 + 0.105 * word, 0.4 + 0.05 * row, x0 + 0.095 + 0.105 * word, 0.43 + 0.05 * row]) + boxes = np.asarray(boxes) + doc_builder = builder.DocumentBuilder() + # Without layout, each row spans the 3 columns + assert len(doc_builder._resolve_lines(boxes)) == 2 + # With a layout region per column, a line never spans two of them + word_regions = doc_builder._word_regions(boxes, regions, ["Text"] * 3) + lines = doc_builder._resolve_lines(boxes, word_regions=word_regions) + assert sorted(lines) == [[6 * col + 3 * row + word for word in range(3)] for col in range(3) for row in range(2)] + # Overlapping regions of the same label (duplicates of a layout model) are one region + overlapping = [((0.0, 0.0), (0.6, 1.0)), ((0.3, 0.0), (1.0, 1.0))] + assert ( + len(doc_builder._resolve_lines(boxes, word_regions=doc_builder._word_regions(boxes, overlapping, ["Text"] * 2))) + == 2 + ) + # Words outside any region do not split a line + assert doc_builder._word_regions(boxes, [((0.5, 0.0), (0.6, 0.1))], ["Text"])[0].tolist() == [-1] * 18 + # End to end + doc = doc_builder( + [np.zeros((100, 100, 3), dtype=np.uint8)], + [boxes], + [np.ones(18)], + [[("word", 0.9)] * 18], + [(100, 100)], + [[{"value": 0, "confidence": None}] * 18], + regions=[{"boxes": np.asarray([[*p0, *p1] for p0, p1 in regions]), "class_names": ["Text"] * 3}], + ) + assert [len(line.words) for block in doc.pages[0].blocks for line in block.lines] == [3] * 6 + + +def test_resolve_lines_nested_and_short_regions(): + doc_builder = builder.DocumentBuilder() + # A line of 8 words in a Text region, with a Picture box over its last word (or an inline Formula box) + boxes = np.asarray([[0.10 + 0.1 * i, 0.20, 0.18 + 0.1 * i, 0.23] for i in range(8)]) + text = ((0.08, 0.18), (0.92, 0.40)) + for nested in [((0.78, 0.15), (0.95, 0.24)), ((0.38, 0.19), (0.59, 0.24))]: + word_regions = doc_builder._word_regions(boxes, [text, nested], ["Text", "Picture"]) + assert doc_builder._resolve_lines(boxes, word_regions=word_regions) == [list(range(8))] + # Two columns whose regions start slightly below the first line of the right column (its words are 46% covered) + boxes = np.asarray([[0.05 + 0.1 * i, 0.30, 0.13 + 0.1 * i, 0.33] for i in range(4)]) + boxes = np.concatenate([boxes, boxes + [0.42, 0.0, 0.42, 0.0]]) + regions = [((0.03, 0.28), (0.45, 0.9)), ((0.46, 0.3162), (0.89, 0.9))] + word_regions = doc_builder._word_regions(boxes, regions, ["Text", "Text"]) + assert sorted(doc_builder._resolve_lines(boxes, word_regions=word_regions)) == [[0, 1, 2, 3], [4, 5, 6, 7]] + + +def test_resolve_lines_small_paragraph_break(): + # A table whose cells are separated by one word height: a small `paragraph_break` still splits the cells + boxes = [] + for row in range(3): + x, y = 0.1, 0.3 + 0.03 * row + for _ in range(4): + for _ in range(2): + boxes.append([x, y, x + 0.05, y + 0.015]) + x += 0.054 + x += 0.011 + shape = (2339, 1654) + assert len(builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), shape)) == 3 + assert len(builder.DocumentBuilder(paragraph_break=0.01)._resolve_lines(np.asarray(boxes), shape)) == 12 + + +def test_resolve_lines_large_type_on_a_narrow_page(): + # Justified lines in large type on a narrow crop: a word gap wider than `paragraph_break` (relative to the page + # width) but smaller than a word height does not split the line + def _line(y, gaps): + x, boxes = 0.05, [] + for gap in [0.0, *gaps]: + x += gap + boxes.append([x, y, x + 0.12, y + 0.04]) + x += 0.12 + return boxes + + boxes = _line(0.10, [0.019, 0.019, 0.040, 0.019, 0.019]) + _line(0.16, [0.019] * 5) + _line(0.22, [0.019] * 5) + lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), (643, 468)) + assert lines == [list(range(6)), list(range(6, 12)), list(range(12, 18))] + # The page-relative paragraph break still applies on a page with regular type + assert len(builder.DocumentBuilder()._resolve_lines(np.asarray(boxes) * [1, 0.2, 1, 0.2], (643, 468))) == 4 + + +def test_resolve_sub_lines_nested_box(): + # A box nested in a wider one (e.g. a punctuation mark detected twice) must not inflate the gap to the next word + boxes = np.asarray([[0.10, 0.1, 0.30, 0.14], [0.25, 0.11, 0.26, 0.13], [0.32, 0.1, 0.45, 0.14]]) + assert builder.DocumentBuilder()._resolve_sub_lines(boxes, [0, 1, 2], 0.05) == [[0, 1, 2]] + assert builder.DocumentBuilder()._resolve_sub_lines(boxes, [0, 1, 2], 0.01) == [[0, 1], [2]] + + def test_points_in_polygons(): polys = np.array( [ diff --git a/tests/common/test_models_reading_order.py b/tests/common/test_models_reading_order.py index 916217c962..e1df962bf1 100644 --- a/tests/common/test_models_reading_order.py +++ b/tests/common/test_models_reading_order.py @@ -11,6 +11,7 @@ resolve_reading_segments, sort_reading_order, ) +from doctr.models.reading_order.base import _region_groups @pytest.mark.parametrize( @@ -194,6 +195,81 @@ def test_resolve_reading_segments(): assert resolve_reading_segments([]) == [] +def test_resolve_reading_segments_rows_and_regions(): + # The fragments of a line split at a wide word gap belong to the paragraph of that line + geoms = [((0.10, 0.10), (0.40, 0.13)), ((0.43, 0.10), (0.90, 0.13)), ((0.10, 0.14), (0.90, 0.17))] + assert resolve_reading_segments(geoms) == [[0, 1, 2]] + # but not two elements separated by a gutter + assert resolve_reading_segments([((0.10, 0.10), (0.40, 0.13)), ((0.48, 0.10), (0.90, 0.13))]) == [[0], [1]] + # The gap is measured in the unit of the element heights, relative coordinates of a wide page being stretched + geoms = [((0.10, 0.10), (0.40, 0.13)), ((0.44, 0.10), (0.90, 0.13))] + assert resolve_reading_segments(geoms) == [[0, 1]] + assert resolve_reading_segments(geoms, page_shape=(500, 1000)) == [[0], [1]] + # Distinct layout regions are distinct segments + geoms = [((0.1, 0.1), (0.9, 0.13)), ((0.1, 0.14), (0.9, 0.17))] + assert resolve_reading_segments(geoms, region_groups=[0, 0]) == [[0, 1]] + assert resolve_reading_segments(geoms, region_groups=[0, 1]) == [[0], [1]] + + +def test_sort_reading_order_values_shifted_between_rows(): + # Receipt totals on a crumpled page (pixels): each amount sits lower than its label and also overlaps the label of + # the next row, but it overlaps its own label the most + geoms = [((60, 1200 + 60 * i), (600, 1240 + 60 * i)) for i in range(3)] # 0-2: full-width lines + geoms += [ + ((114, 1376), (282, 1420)), # 3: label + ((420, 1388), (520, 1431)), # 4: its amount + ((112, 1408), (226, 1451)), # 5: label + ((420, 1422), (520, 1465)), # 6: its amount + ((110, 1445), (162, 1483)), # 7: label + ((410, 1456), (520, 1502)), # 8: its amount + ((110, 1481), (162, 1518)), # 9: last label + ] + geoms += [((60, 1560 + 60 * i), (600, 1600 + 60 * i)) for i in range(2)] # 10-11: full-width lines + assert sort_reading_order(geoms) == list(range(12)) + # The cells of a single row are still read left to right when the middle one sits a little lower + row = [((100, 100), (200, 120)), ((300, 103), (400, 123)), ((500, 100), (600, 120))] + assert sort_reading_order(row) == [0, 1, 2] + assert sort_reading_order(row + [((100, 140), (200, 160)), ((300, 140), (400, 160))]) == [0, 1, 2, 3, 4] + + +def test_region_groups(): + regions = np.asarray([ + [0.0, 0.0, 0.5, 0.5], # 0 + [0.3, 0.3, 0.6, 0.6], # 1: overlaps 0 by 44% of its area + [0.7, 0.0, 0.9, 0.2], # 2: disjoint + [0.0, 0.0, 0.5, 0.5], # 3: same box as 0, other label + ]) + groups = _region_groups(regions, ["Text", "Text", "Text", "Picture"]) + assert groups[0] == groups[1] and len({groups[0], groups[2], groups[3]}) == 3 + # A small overlap does not group + assert len(set(_region_groups(regions[:2], ["Text"] * 2, min_overlap=0.5))) == 2 + + +def test_resolve_reading_segments_row_then_next_line(): + # A line split into 3 fragments (words missed by the detection), then the short last line of its paragraph, + # then 2 more paragraphs (pixels) + boxes = [ + [46, 30, 731, 70], + [769, 36, 819, 70], + [857, 31, 1233, 68], + [41, 89, 125, 121], + [42, 245, 1199, 288], + [39, 288, 1208, 338], + [41, 459, 774, 502], + [825, 457, 1236, 500], + [41, 518, 1208, 558], + [39, 576, 1019, 622], + ] + geoms = [((x0, y0), (x1, y1)) for x0, y0, x1, y1 in boxes] + assert resolve_reading_segments(geoms) == [[0, 1, 2, 3], [4, 5], [6, 7, 8, 9]] + # Within a layout region of body text, the fragments of a row belong together whatever the gap + geoms = [((0.05, 0.10), (0.30, 0.13)), ((0.70, 0.10), (0.95, 0.13)), ((0.05, 0.14), (0.95, 0.17))] + assert resolve_reading_segments(geoms) == [[0], [1, 2]] + assert resolve_reading_segments(geoms, labels=["Text"] * 3, region_groups=[0, 0, 0]) == [[0, 1, 2]] + # A degenerate page shape is ignored + assert resolve_reading_segments(geoms, page_shape=(0, 0)) == resolve_reading_segments(geoms) + + def test_assign_layout_labels(): geoms = [((0.1, 0.1), (0.4, 0.2)), ((0.6, 0.1), (0.9, 0.2)), ((0.1, 0.5), (0.4, 0.6))] regions = [((0.05, 0.05), (0.45, 0.25)), ((0.55, 0.05), (0.95, 0.25))] From 4fca20bd680df542772d9dd2fc07017d1013da47 Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 13 Oct 2026 13:11:00 +0200 Subject: [PATCH 2/4] Update --- doctr/io/exporters.py | 16 ++- doctr/models/builder.py | 128 +++++++++---------- doctr/models/reading_order/base.py | 54 ++++---- tests/common/test_io_exporters.py | 129 +++++++++---------- tests/common/test_models_builder.py | 146 ++++++++++------------ tests/common/test_models_reading_order.py | 75 ++++++----- 6 files changed, 254 insertions(+), 294 deletions(-) diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py index 9a90535d22..0942ce8ccc 100644 --- a/doctr/io/exporters.py +++ b/doctr/io/exporters.py @@ -90,7 +90,7 @@ def _covering_region_indices( def _merge_overlapping_regions(elt_regions: list[int], region_geoms: list[Any], region_labels: list[str]) -> list[int]: - """Map the layout region of each element to a group shared by the overlapping regions of the same label + """Map the layout region of each element to its group of overlapping regions of the same label Args: elt_regions: the layout region index of each element (-1 for none) @@ -225,14 +225,14 @@ def page_reading_order( elt_labels[idx] = "Table" elif not isinstance(elt, LayoutElement) and elt_regions[idx] in furniture_labels: elt_labels[idx] = furniture_labels[elt_regions[idx]] - # Page-wide figures do not vote in the multi-column detection + # Page-wide figures and tables do not vote in the multi-column detection boxes = _to_boxes(elt_geoms) span = float(boxes[:, 2].max() - boxes[:, 0].min()) or 1.0 column_voters = [ - not (isinstance(elt, LayoutElement) and box[2] - box[0] > 0.5 * span) for elt, box in zip(elements, boxes) + not (isinstance(elt, (LayoutElement, Table)) and box[2] - box[0] > 0.5 * span) + for elt, box in zip(elements, boxes) ] - # A text line lying in a table or a figure region keeps the label of its region, but it is not a float itself: - # it is grouped into paragraphs like regular text (and still read right after its figure) + # A text line in a table or figure region keeps its label, but is grouped into paragraphs like regular text order_labels = [ _float_text_label(str(label)) if not isinstance(elt, (Table, LayoutElement)) and layout_label_role(label) == "float" @@ -681,7 +681,11 @@ def _flush_list() -> None: if len(item_lines) == 0: continue norm_label = normalize_layout_label(label) - if norm_label in self.headings: + if layout_label_role(label) == "float": + # Text of a table or figure region without structure: one paragraph per row + _flush_list() + parts.extend(self.render_paragraph([line], escape) for line in item_lines) + elif norm_label in self.headings: _flush_list() parts.append(self.render_heading(norm_label, item_lines, escape)) elif norm_label in _LIST_LABELS: diff --git a/doctr/models/builder.py b/doctr/models/builder.py index 118b5ac199..d8218d6540 100644 --- a/doctr/models/builder.py +++ b/doctr/models/builder.py @@ -36,11 +36,9 @@ __all__ = ["DocumentBuilder"] -# Median word height of body text on an A4 page, relative to the page width: `paragraph_break` is scaled up on pages -# whose type is larger than that (e.g. a crop of a page) +# Median word height of body text relative to the page width, above which `paragraph_break` is scaled up _BODY_WORD_HEIGHT = 0.02 -# Cost of the horizontal distance between a box and a row, relative to their vertical offset, when several rows are -# level with the box: 50 units of horizontal distance weigh as much as 1 unit of vertical offset +# Weight of the horizontal distance to a row, relative to the vertical offset, when several rows are level with a box _ROW_DISTANCE_WEIGHT = 0.02 @@ -50,9 +48,7 @@ class DocumentBuilder(NestedObject): Args: resolve_lines: whether words should be automatically grouped into lines resolve_blocks: whether lines should be automatically grouped into blocks - paragraph_break: relative length of the minimum space separating paragraphs. It is scaled up on pages with - large type (e.g. a crop of a page), and lines can be split at a smaller distance adapted to the word - spacing of the page. + paragraph_break: relative length of the minimum space separating paragraphs (scaled up for large type) export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. keep_reading_order: if True, arrange the content of every page in reading order (cf. @@ -128,10 +124,8 @@ def _resolve_sub_lines( boxes: bounding boxes of shape (N, 4) word_idcs: list of indexes for the words of the line break_dist: horizontal distance above which the line is split (defaults to `self.paragraph_break`) - word_regions: optional layout region (group) of each box (-1 for none) and coverage of each box by each - region group (cf. `_word_regions`): the line is also split between two boxes of distinct regions (e.g. - two columns separated by a gutter as narrow as a word gap), unless one also lies in the region of the - other (nested or overlapping regions) + word_regions: optional region group of each box and coverage by each group (cf. `_word_regions`): the + line is also split between two distinct regions, unless one box lies in the region of the other Returns: A list of (sub-)lines computed from the original line (words) @@ -146,9 +140,7 @@ def _resolve_sub_lines( lines.append(word_idcs) else: sub_line = [word_idcs[0]] - # Right edge of the current sub-line: a box can end before an earlier, wider one (e.g. a punctuation box - # nested in a word box), so the running max is tracked instead of the right edge of the last box - right = boxes[word_idcs[0], 2] + right = boxes[word_idcs[0], 2] # a box can be nested in an earlier one: keep the max right edge groups, coverage = word_regions if word_regions is not None else (None, None) region = -1 if groups is None else int(groups[word_idcs[0]]) # region of the sub-line prev = word_idcs[0] @@ -160,8 +152,7 @@ def _resolve_sub_lines( # If distance between boxes is lower than the break distance, same sub-line if dist < break_dist: horiz_break = False - # Unless the box lies in another layout region than the sub-line, and neither the box lies in the - # region of the sub-line nor the previous box in the region of the box + # Unless the box and the sub-line lie in distinct regions (e.g. columns separated by a narrow gutter) box_region = -1 if groups is None else int(groups[i]) if ( coverage is not None @@ -195,10 +186,7 @@ def _word_regions( region_labels: list[str], shape: tuple[int, int] | None = None, ) -> tuple[np.ndarray, np.ndarray] | None: - """Assign each box to the layout region covering it, overlapping regions of the same label being grouped - - A box covered by no region by half (e.g. a word sticking out of a region slightly too short) is assigned to - the smallest region holding it horizontally and covering it at least partly vertically. + """Assign each box to the group of the layout region covering it (or holding it horizontally) Args: boxes: bounding boxes of shape (N, 4) or (N, 4, 2) @@ -207,8 +195,7 @@ def _word_regions( shape: the page dimensions (height, width), used to de-skew rotated pages exactly Returns: - the region group of each box (-1 for none) and the coverage of each box by each region group (indexed by - group), None without layout regions + the region group of each box (-1 for none) and the coverage of each box by each group, None without regions """ from doctr.models.reading_order.base import ( _covering_regions, @@ -227,7 +214,6 @@ def _word_regions( regions, words = _to_boxes(region_geoms), _to_boxes(word_geoms) groups = np.asarray(_region_groups(regions, region_labels)) covering = np.asarray(_covering_regions(words, regions, 0.5, region_labels)) - # Coverage of each box by each region, on both axes and in area inter_w = np.minimum(words[:, None, 2], regions[None, :, 2]) - np.maximum( words[:, None, 0], regions[None, :, 0] ) @@ -237,81 +223,90 @@ def _word_regions( x_cov = np.clip(inter_w, 0, None) / np.clip(words[:, 2:3] - words[:, 0:1], 1e-9, None) y_cov = np.clip(inter_h, 0, None) / np.clip(words[:, 3:4] - words[:, 1:2], 1e-9, None) uncovered = np.flatnonzero(covering < 0) - if uncovered.size > 0: + if uncovered.size > 0: # e.g. a word sticking out of a region slightly too short areas = (regions[:, 2] - regions[:, 0]) * (regions[:, 3] - regions[:, 1]) holding = (x_cov[uncovered] >= 0.9) & (y_cov[uncovered] >= 0.3) smallest = np.where(holding, areas[None, :], np.inf).argmin(axis=1) covering[uncovered] = np.where(holding.any(axis=1), smallest, -1) coverage = np.zeros((words.shape[0], regions.shape[0])) - np.maximum.at(coverage.T, groups, (x_cov * y_cov).T) # coverage by a group: the max over its regions + np.maximum.at(coverage.T, groups, (x_cov * y_cov).T) return np.where(covering >= 0, groups[covering], -1), coverage @staticmethod def _resolve_rows(boxes: np.ndarray, idxs: np.ndarray, y_med: float) -> list[list[int]]: """Group straight boxes into visual rows - A box can join a row if its vertical center is close enough (less than half a median box height) to the median - center of its horizontally closest boxes in that row. Comparing a box with its neighbors rather than with the - mean of the whole row keeps a slanted line in a single row, and prevents a word from being drawn into the row - of another column whose baseline is slightly offset. When several rows are level with the box, the best - aligned one wins, the horizontal distance to the row being added as a small penalty. + A box joins the row whose horizontally closest boxes are the best aligned with it (less than half a median + height apart), which follows slanted lines and keeps the rows of offset columns apart. Args: boxes: straight bounding boxes of shape (N, 4) - idxs: the order in which the boxes are visited (cf. `_sort_boxes`: top to bottom, left to right) + idxs: the order in which the boxes are visited (cf. `_sort_boxes`) y_med: the median box height Returns: the rows (lists of box indices), from top to bottom """ - if not y_med > 0: # degenerate heights: no box is close enough to another one + if not y_med > 0: return [[idx] for idx in idxs.tolist()] lefts, rights = boxes[:, 0].tolist(), boxes[:, 2].tolist() centers = ((boxes[:, 1] + boxes[:, 3]) / 2).tolist() - rows: list[list[int]] = [] - row_lefts: list[list[float]] = [] # left edges of the boxes of each row, sorted + row_lefts: list[list[float]] = [] # sorted left edges of the boxes of each row row_boxes: list[list[int]] = [] # boxes of each row, sorted by left edge - highest: list[float] = [] # highest box center of each row - lowest: list[float] = [] # lowest box center of each row - active: list[int] = [] # rows which can still receive a box - for idx in idxs.tolist(): + spans: list[list[float]] = [] # highest and lowest box center of each row + + def _best_row(idx: int, rows: list[int]) -> int: center, left, right = centers[idx], lefts[idx], rights[idx] - # Boxes come roughly top to bottom: a row whose lowest box is far above can no longer be joined - active = [row for row in active if lowest[row] > center - 2 * y_med] best, best_key = -1, (float("inf"), float("inf")) - for row in active: - # The neighbors of the box can only be level with it if the row spans its center - if highest[row] >= center + y_med / 2 or lowest[row] <= center - y_med / 2: + for row in rows: + if spans[row][0] >= center + y_med / 2 or spans[row][1] <= center - y_med / 2 or not row_boxes[row]: continue - # The horizontally closest boxes of the row lie around the position of the box in that row. A few of - # them rather than a single one, so that a small mark (accent, punctuation) can't steer the box. + # A few neighbors rather than one, so that a small mark (accent, punctuation) can't steer the box pos = bisect.bisect_left(row_lefts[row], left) nearest = sorted( (max(lefts[other] - right, left - rights[other], 0.0), other) for other in row_boxes[row][max(pos - 3, 0) : pos + 3] )[:3] offset = abs(statistics.median(centers[other] for _, other in nearest) - center) - # On a slanted line, the next line can start before the end of the current one: the offset alone - # would let a distant row of similar height take the box + # The distance breaks ties on slanted lines, where the next line can start before this one ends key = (offset + _ROW_DISTANCE_WEIGHT * nearest[0][0], nearest[0][0]) if offset < y_med / 2 and key < best_key: best, best_key = row, key - if best < 0: - best = len(rows) - rows.append([]) + return best + + def _add(row: int, idx: int) -> None: + if row == len(row_boxes): row_lefts.append([]) row_boxes.append([]) - highest.append(center) - lowest.append(center) - active.append(best) - rows[best].append(idx) - pos = bisect.bisect_right(row_lefts[best], left) - row_lefts[best].insert(pos, left) - row_boxes[best].insert(pos, idx) - highest[best] = min(highest[best], center) - lowest[best] = max(lowest[best], center) - - return rows + spans.append([centers[idx], centers[idx]]) + pos = bisect.bisect_right(row_lefts[row], lefts[idx]) + row_lefts[row].insert(pos, lefts[idx]) + row_boxes[row].insert(pos, idx) + spans[row] = [min(spans[row][0], centers[idx]), max(spans[row][1], centers[idx])] + + row_of: dict[int, int] = {} + active: list[int] = [] # rows which can still receive a box: the boxes come roughly top to bottom + for idx in idxs.tolist(): + active = [row for row in active if spans[row][1] > centers[idx] - 2 * y_med] + row = _best_row(idx, active) + if row < 0: + row = len(row_boxes) + active.append(row) + _add(row, idx) + row_of[idx] = row + # Assign each box again once all rows exist: a box visited before the rest of its line (e.g. a short word + # ending higher) may have joined a level row of another column + highest, lowest = np.asarray(spans).T + for idx in idxs.tolist(): + own = row_of[idx] + pos = row_boxes[own].index(idx) + del row_boxes[own][pos], row_lefts[own][pos] + candidates = np.flatnonzero((highest < centers[idx] + y_med / 2) & (lowest > centers[idx] - y_med / 2)) + row = _best_row(idx, candidates.tolist()) + row_of[idx] = own if row < 0 else row + _add(row_of[idx], idx) + + return [row for row in row_boxes if row] def _resolve_lines( self, @@ -324,8 +319,7 @@ def _resolve_lines( Args: boxes: bounding boxes of shape (N, 4) or (N, 4, 2) in case of rotated bbox shape: the page dimensions (height, width), used to de-skew rotated pages exactly - word_regions: optional layout region (group) of each box and coverage by each group (cf. `_word_regions`): - a line never spans two distinct regions + word_regions: optional region group of each box and coverage by each group (cf. `_word_regions`) Returns: nested list of box indices @@ -349,14 +343,12 @@ def _resolve_lines( row_idcs = np.asarray(row) row_idcs = row_idcs[np.argsort(boxes[row_idcs, 0], kind="stable")] n_pairs += row_idcs.shape[0] - 1 - # Same gap definition as `_resolve_sub_lines`: distance to the right edge of everything on the left gap_chunks.append(boxes[row_idcs[1:], 0] - np.maximum.accumulate(boxes[row_idcs, 2])[:-1]) all_gaps = np.concatenate(gap_chunks) if gap_chunks else np.empty(0, dtype=boxes.dtype) pos_gaps = all_gaps[all_gaps > 0] # Median word height, converted to width units aspect = (shape[0] / shape[1]) if shape is not None else 1.0 floor = float(y_med) * aspect - # The width-relative paragraph break grows with the type size (large type on a narrow page or crop) cap = self.paragraph_break * max(1.0, floor / _BODY_WORD_HEIGHT) if pos_gaps.shape[0] >= 5 and pos_gaps.shape[0] >= 0.5 * n_pairs: break_dist = min(cap, max(3.0 * float(np.median(pos_gaps)), floor)) @@ -752,7 +744,7 @@ def _build_blocks( the general orientation (orientations + confidences) of the crops shape: the page dimensions (height, width), used to de-skew rotated pages exactly when resolving lines - layout: optional layout regions of the page: a line never spans two of them + layout: optional layout regions of the page Returns: list of block elements @@ -985,9 +977,7 @@ class KIEDocumentBuilder(DocumentBuilder): Args: resolve_lines: whether words should be automatically grouped into lines resolve_blocks: whether lines should be automatically grouped into blocks - paragraph_break: relative length of the minimum space separating paragraphs. It is scaled up on pages with - large type (e.g. a crop of a page), and lines can be split at a smaller distance adapted to the word - spacing of the page. + paragraph_break: relative length of the minimum space separating paragraphs (scaled up for large type) export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. keep_reading_order: if True, arrange the content of every page in reading order (cf. diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index b87e311e6e..a2b029c020 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -44,8 +44,7 @@ _FURNITURE_ROLES = ("header", "footer", "footnote") -# Label suffix of the text lying in a float region without float element (e.g. 'Table text' for the text of a table -# region without table structure): it is ordered and grouped like body text, but moved after a region like a float +# Label suffix of the text lying in a float region without float element (e.g. a table without structure) _FLOAT_TEXT_SUFFIX = " text" @@ -92,7 +91,7 @@ def _float_text_label(label: str) -> str: def _is_float_text(label: str | None) -> bool: - """Check whether a label is the one of a text element lying in a float region (cf. `_float_text_label`)""" + """Check whether a label is the one of a text element lying in a float region""" norm = normalize_layout_label(label) return any(norm == normalize_layout_label(_float_text_label(float_label)) for float_label in _FLOAT_LABELS) @@ -230,9 +229,7 @@ def _topological_order( same_row = y_overlap > y_overlap_threshold same_row &= ~x_linked same_row &= x_rank[:, None] < x_rank[None, :] # i is on the left of j - # An element belongs to the row it overlaps the most: an element shifted between two rows (e.g. an amount on a - # crumpled receipt) must not wait for the element on the left of the next row. Only the elements of another row - # than its best partner are dropped: a slightly lower cell of its own row is still read before it. + # An element shifted between two rows (e.g. an amount on a crumpled receipt) only waits for the row it overlaps most best_left = np.max(y_overlap, axis=0, where=same_row, initial=0.0) weak_i, weak_j = np.nonzero(same_row & (y_overlap < best_left[None, :] - 0.1)) if weak_i.size > 0: @@ -249,23 +246,14 @@ def _topological_order( order: list[int] = [] last = -1 - # Column components: elements connected through horizontal overlap, ignoring page-spanning lines (e.g. - # titles) which would otherwise merge separate columns. This keeps the traversal inside a column even when - # the direct top-to-bottom continuation is momentarily broken by a fragmented or misaligned OCR line. - page_width = float(x1.max() - x0.min()) or 1.0 - spanning = (x1 - x0) > 0.5 * page_width - - # Reuse the horizontal-overlap matrix; keep the upper triangle only, so each pair is visited once - col_edges = np.triu(x_linked, 1) - col_edges &= ~spanning[:, None] - col_edges &= ~spanning[None, :] - component = connected_components(csr_matrix(col_edges), directed=False, return_labels=True)[1] - # Detect whether the page is multi-column: if a vertical line can be drawn that separates the boxes into two # groups with a small number of crossing boxes, the page is considered multi-column. This is used to # favor the continuation of the current column when traversing the graph, which keeps multi-column bodies intact # even for non-Manhattan layouts where recursive XY-cuts fail to find a valid split. multi_column = False + page_width = float(x1.max() - x0.min()) or 1.0 + spanning = (x1 - x0) > 0.5 * page_width + bridging = spanning.copy() # elements which must not merge the columns: page-wide or crossing the gutter if voters is None or np.count_nonzero(voters) < 3: voters = np.ones(num_boxes, dtype=bool) vx0, vx1, num_voters = x0[voters], x1[voters], int(np.count_nonzero(voters)) @@ -281,8 +269,16 @@ def _topological_order( left = int(np.count_nonzero(centers <= split)) if crossing <= tolerance and left >= 0.25 * num_voters and num_voters - left >= 0.25 * num_voters: multi_column = True + bridging |= np.minimum(x1 - split, split - x0) > 0.02 * span break + # Column components: elements connected through horizontal overlap. This keeps the traversal inside a column even + # when the direct top-to-bottom continuation is momentarily broken by a fragmented or misaligned OCR line. + col_edges = np.triu(x_linked, 1) # each pair once + col_edges &= ~bridging[:, None] + col_edges &= ~bridging[None, :] + component = connected_components(csr_matrix(col_edges), directed=False, return_labels=True)[1] + while len(order) < num_boxes: ready = np.flatnonzero((in_degree == 0) & ~emitted) if ready.size == 0: # cycle safety net (can only happen with degenerate overlapping geometries) @@ -688,7 +684,7 @@ def resolve_reading_segments( This is typically used to regroup individual lines into paragraphs once they have been put in reading order: two consecutive elements are merged into the same segment when they share the same (possibly None) layout label and layout region, and either belong to the same column and are vertically close enough, or - continue the same visual row (e.g. the fragments of a line split at a wide word gap). + continue the same visual row. Args: geoms: sequence of geometries in any docTR format (cf. `sort_reading_order`) @@ -711,10 +707,8 @@ def resolve_reading_segments( region_groups: optional layout region index of each element (-1 for none), used to attach the captions of a region as a unit, to read a float standing beside a region after it, and to keep distinct regions in distinct segments - row_gap: maximum horizontal gap between two consecutive elements of the same visual row to belong to the - same segment, as a multiple of the median element height (pass `page_shape` along with relative - coordinates of a non-square page, so that both are measured in the same unit). It does not apply to two - elements of body text in the same layout region. + row_gap: maximum horizontal gap between two consecutive elements of a visual row to join them, as a multiple + of the median element height (not applied to body text of the same layout region) Returns: a partition of the input indices into reading-ordered segments (each segment being itself in @@ -736,15 +730,13 @@ def resolve_reading_segments( return [] canonical = _to_canonical_ltr(_to_boxes(geoms), direction) median_height = float(np.median(canonical[:, 3] - canonical[:, 1])) - # Relative coordinates distort the horizontal distances on non-square pages: scale them to the height unit - aspect = 1.0 + aspect = 1.0 # horizontal distances in the height unit if page_shape is not None and page_shape[0] > 0 and page_shape[1] > 0: aspect = page_shape[0] / page_shape[1] if direction.startswith("ttb") else page_shape[1] / page_shape[0] norm_labels = [normalize_layout_label(labels[idx] if labels is not None else None) for idx in range(len(order))] segments: list[list[int]] = [[order[0]]] - # Extent of the trailing visual row of the current segment (the fragments of a line split at wide word gaps) - row = canonical[order[0]].copy() + row = canonical[order[0]].copy() # extent of the last visual row of the current segment for prev, cur in zip(order[:-1], order[1:]): box = canonical[cur] x_gap = min(row[2], box[2]) - max(row[0], box[0]) @@ -753,8 +745,7 @@ def resolve_reading_segments( y_overlap = y_inter / max(min(canonical[prev, 3] - canonical[prev, 1], box[3] - box[1]), 1e-9) stacked = x_overlap > x_overlap_threshold and box[1] - row[3] <= paragraph_gap * median_height same_region = region_groups is None or region_groups[prev] == region_groups[cur] - # The next element on the right in the same visual row, e.g. the remainder of a line split at a wide word gap - # (whatever the gap within a region of body text: the layout model already grouped both) + # The rest of a line split at a wide word gap (any gap within a region of body text) same_row = ( y_overlap > y_overlap_threshold and box[0] > canonical[prev, 0] @@ -828,10 +819,7 @@ def assign_layout_labels( def _region_groups(regions: np.ndarray, labels: Sequence[str | None], min_overlap: float = 0.1) -> list[int]: - """Group the overlapping layout regions of the same label - - Layout models can return overlapping regions for a single paragraph (e.g. a second region over the part of a - paragraph wrapping around a figure), which must not split it. + """Group the overlapping layout regions of the same label (e.g. two regions for a paragraph around a figure) Args: regions: the (R, 4) boxes of the layout regions diff --git a/tests/common/test_io_exporters.py b/tests/common/test_io_exporters.py index f7adb77998..bb9e57c47c 100644 --- a/tests/common/test_io_exporters.py +++ b/tests/common/test_io_exporters.py @@ -35,6 +35,10 @@ def _line_at(text, x0, y0, x1, y1, rtl=False): ]) +def _texts(items): + return [[" ".join(word.value for word in line.words) for line in item.lines] for item in items] + + def _reading_order_page(): """A page in the default builder configuration (single block) with a title, 2 columns & a footer""" lines = [_line_at("A Two Column Study", 0.2, 0.05, 0.8, 0.09)] @@ -70,6 +74,14 @@ def test_page_items_in_reading_order(): bottom = elements.Block([_line_at("last words", 0.1, 0.5, 0.9, 0.55)]) page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [bottom, top], 0, (1000, 800)) assert [block.render() for block in page.items_in_reading_order()] == ["first words", "last words"] + # Degenerate page dimensions + lines = [_line_at("first line here", 0.1, 0.1, 0.5, 0.13), _line_at("second part", 0.6, 0.1, 0.9, 0.13)] + for dims in [(0, 0), (0, 100), (100, 0)]: + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, dims) + assert [text for item in _texts(page_reading_order(page)[0]) for text in item] == [ + "first line here", + "second part", + ] def test_page_export_as_markdown(): @@ -128,7 +140,7 @@ def test_page_export_as_markdown_rtl(): assert page.export_as_markdown(direction="ltr").startswith("الأيسر") -def test_page_reading_order_text_in_a_table_region(): +def test_page_reading_order_table_region_text(): # A table region from the layout model, without table structure (`detect_tables=False`): its lines keep the # region label, but they are not floats themselves, so they are neither nested into each other nor isolated lines = [ @@ -149,15 +161,9 @@ def test_page_reading_order_text_in_a_table_region(): assert texts == ["Pos Datum Einzelpreis Gesamt", "pro UE", "EUR", "1 Deutsch 22,50 22,50"] assert labels == ["Table"] * len(items) assert len(items) < len(lines) - - -def _texts(items): - return [[" ".join(word.value for word in line.words) for line in item.lines] for item in items] - - -def test_page_reading_order_table_region_beside_a_paragraph(): - # The rows of a table region (without table structure) level with the lines of the paragraph beside it are read - # after that paragraph, not interleaved with its lines + # Rendered like the rows of a table: one paragraph per line + assert page.export_as_markdown().split("\n\n") == texts + # The rows of a table region level with the paragraph beside it are read after that paragraph lines = [_line_at(f"intro line {i} of the body", 0.08, 0.05 + 0.04 * i, 0.92, 0.08 + 0.04 * i) for i in range(6)] lines += [ _line_at(f"para line {i} beside the table", 0.08, 0.30 + 0.04 * i, 0.55, 0.33 + 0.04 * i) for i in range(4) @@ -182,59 +188,6 @@ def test_page_reading_order_table_region_beside_a_paragraph(): assert {label for item, label in zip(_texts(items), labels) if item[0].startswith("Row")} == {"Table"} -def test_page_reading_order_overlapping_text_regions(): - # The layout model returns a second, overlapping region for the lines of a paragraph wrapping around a figure: - # the paragraph is not split where the lines switch from one region to the other - lines = [ - _line_at("one two three four five six", 0.05, 0.10, 0.90, 0.13), - _line_at("seven eight nine ten eleven", 0.05, 0.14, 0.90, 0.17), - _line_at("twelve thirteen fourteen fifteen", 0.05, 0.18, 0.90, 0.21), - _line_at("sixteen seventeen", 0.05, 0.22, 0.50, 0.25), - _line_at("eighteen nineteen", 0.05, 0.26, 0.50, 0.29), - _line_at("twenty end.", 0.05, 0.30, 0.50, 0.33), - ] - layout = [ - elements.LayoutElement("Text", 0.90, ((0.04, 0.09), (0.91, 0.215))), - elements.LayoutElement("Text", 0.59, ((0.04, 0.16), (0.51, 0.34))), - elements.LayoutElement("Picture", 0.86, ((0.55, 0.22), (0.90, 0.34))), - ] - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1810, 1280), layout=layout - ) - items, _, _ = page_reading_order(page) - assert len(items) == 1 and len(items[0].lines) == 6 - - -def test_page_export_hides_figure_text_sticking_out_as_a_block(): - # Each line of the figure text lies in the figure, but the box of the block they form sticks out of it - lines = [ - _line_at("Body text above the chart spanning the page", 0.05, 0.05, 0.95, 0.09), - _line_at("AXIS", 0.30, 0.185, 0.40, 0.225), - _line_at("LABEL", 0.45, 0.20, 0.45 + 0.299 / 0.9, 0.24), - _line_at("Body text below the chart spanning the page", 0.05, 0.70, 0.95, 0.74), - ] - page = elements.Page( - np.full((100, 80, 3), 200, dtype=np.uint8), - [elements.Block(lines=lines)], - 0, - (1000, 800), - layout=[elements.LayoutElement("Picture", 0.9, ((0.20, 0.20), (0.60, 0.60)))], - ) - markdown = page.export_as_markdown(images="embedded") - assert "![" in markdown - assert "AXIS" not in markdown and "LABEL" not in markdown - - -def test_page_reading_order_degenerate_dimensions(): - lines = [_line_at("first line here", 0.1, 0.1, 0.5, 0.13), _line_at("second part", 0.6, 0.1, 0.9, 0.13)] - for dims in [(0, 0), (0, 100), (100, 0)]: - page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, dims) - assert [text for item in _texts(page_reading_order(page)[0]) for text in item] == [ - "first line here", - "second part", - ] - - def test_page_export_with_tables(): cells = [ elements.TableCell("Name", 0.9, ((0.1, 0.55), (0.4, 0.6)), 0, 0, 0, 0), @@ -259,6 +212,20 @@ def test_page_export_with_tables(): asciidoc = page.export_as_asciidoc() assert "|===\n|Name |Qty\n\n|Bolt |12\\|3\n|===" in asciidoc assert asciidoc.index("before the table") < asciidoc.index("|===") < asciidoc.index("after the table") + # A page-wide table does not prevent the columns above and below it from being detected + lines = [_line_at("A title over both columns", 0.2, 0.02, 0.8, 0.06)] + for col, x in (("left", 0.05), ("right", 0.55)): + lines += [_line_at(f"{col} {i}", x, 0.1 + 0.05 * i, x + 0.4, 0.13 + 0.05 * i) for i in range(3)] + lines += [_line_at(f"{col} {i}", x, 0.3 + 0.05 * i, x + 0.4, 0.33 + 0.05 * i) for i in range(3, 6)] + table = elements.Table(cells, 2, 2, ((0.05, 0.25), (0.95, 0.28)), 0.95) + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), tables=[table] + ) + blocks = [item for item in page.items_in_reading_order() if isinstance(item, elements.Block)] + texts = [text for item in _texts(blocks) for text in item] + # Each section is read column after column + sections = [[f"{col} {i}" for col in ("left", "right") for i in idcs] for idcs in (range(3), range(3, 6))] + assert texts == ["A title over both columns", *sections[0], *sections[1]] def test_page_export_as_asciidoc(): @@ -1072,6 +1039,23 @@ def test_page_export_markdown_images(): ) jpeg = FigureEncoder("embedded", image_format="jpeg") assert "data:image/jpeg;base64," in page.export_as_markdown(images=jpeg) + # Also when each line lies in the figure, but not the box of the block they form + lines = [ + _line_at("Body text above the chart spanning the page", 0.05, 0.05, 0.95, 0.09), + _line_at("AXIS", 0.30, 0.185, 0.40, 0.225), + _line_at("LABEL", 0.45, 0.20, 0.45 + 0.299 / 0.9, 0.24), + _line_at("Body text below the chart spanning the page", 0.05, 0.70, 0.95, 0.74), + ] + page = elements.Page( + np.full((100, 80, 3), 200, dtype=np.uint8), + [elements.Block(lines=lines)], + 0, + (1000, 800), + layout=[elements.LayoutElement("Picture", 0.9, ((0.20, 0.20), (0.60, 0.60)))], + ) + markdown = page.export_as_markdown(images="embedded") + assert "![" in markdown + assert "AXIS" not in markdown and "LABEL" not in markdown def test_page_export_referenced_images(tmp_path): @@ -1488,6 +1472,25 @@ def test_page_export_figure_does_not_split_a_paragraph(): < markdown.index("Next section") ) assert "body line 11\nlast line of the paragraph" in markdown + # A second, overlapping region over the lines wrapping around the figure does not split the paragraph + lines = [ + _line_at("one two three four five six", 0.05, 0.10, 0.90, 0.13), + _line_at("seven eight nine ten eleven", 0.05, 0.14, 0.90, 0.17), + _line_at("twelve thirteen fourteen fifteen", 0.05, 0.18, 0.90, 0.21), + _line_at("sixteen seventeen", 0.05, 0.22, 0.50, 0.25), + _line_at("eighteen nineteen", 0.05, 0.26, 0.50, 0.29), + _line_at("twenty end.", 0.05, 0.30, 0.50, 0.33), + ] + layout = [ + elements.LayoutElement("Text", 0.90, ((0.04, 0.09), (0.91, 0.215))), + elements.LayoutElement("Text", 0.59, ((0.04, 0.16), (0.51, 0.34))), + elements.LayoutElement("Picture", 0.86, ((0.55, 0.22), (0.90, 0.34))), + ] + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1810, 1280), layout=layout + ) + items, _, _ = page_reading_order(page) + assert len(items) == 1 and len(items[0].lines) == 6 def test_page_export_figure_stacked_in_a_loose_region(): diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index 641a47e4c6..16bd7c05b1 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -426,90 +426,79 @@ def test_sort_boxes(input_boxes, sorted_idxs): assert doc_builder._sort_boxes(np.asarray(input_boxes))[0].tolist() == sorted_idxs +# Two columns whose baselines are offset by half a line (scanned page) +_OFFSET_COLUMNS = [ + [0.1422, 0.5908, 0.2472, 0.6045], # 0-4: left column, "Nichttragende Wände bestehen aus Kalk-" + [0.261, 0.5898, 0.3176, 0.6035], + [0.3314, 0.5898, 0.3978, 0.6035], + [0.413, 0.5928, 0.4406, 0.6035], + [0.4558, 0.5889, 0.4972, 0.6035], + [0.1394, 0.6055, 0.2099, 0.6182], # 5-10: left column, "sandstein oder Vollwanddielen bzw. aus stabilen" + [0.2126, 0.6064, 0.2486, 0.6172], + [0.2513, 0.6055, 0.366, 0.6191], + [0.3674, 0.6045, 0.4047, 0.6191], + [0.4074, 0.6084, 0.4364, 0.6191], + [0.4378, 0.6055, 0.4972, 0.6191], + [0.5539, 0.5811, 0.5953, 0.5957], # 11-13: right column, "Holz- oder Stahlkonstruktion." + [0.5953, 0.5811, 0.634, 0.5947], + [0.634, 0.5811, 0.7694, 0.5977], + [0.5539, 0.5967, 0.6755, 0.6133], # 14-18: right column, "Hauseingangstür als Holz-/Glaselement mit Si-/" + [0.6782, 0.5977, 0.6989, 0.6104], + [0.7045, 0.5977, 0.8426, 0.6113], + [0.844, 0.5977, 0.8758, 0.6133], + [0.8772, 0.5977, 0.9089, 0.6123], +] + + +def _slanted_lines(slope, ys): + return [[0.05 + 0.11 * i, y + slope * i, 0.15 + 0.11 * i, y + 0.03 + slope * i] for y in ys for i in range(8)] + + @pytest.mark.parametrize( - "input_boxes, lines", + "input_boxes, shape, lines", [ - [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], [[2], [1], [0]]], # vertical - [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], [[2], [1], [0]]], # diagonal - [[[0, 0.5, 0.14, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], [[0, 1], [2]]], # same line, 2p - [[[0, 0.5, 0.18, 0.6], [0.2, 0.48, 0.35, 0.58], [0.8, 0.52, 0.9, 0.63]], [[0, 1], [2]]], # ~same line - [[[0, 0.3, 0.48, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], [[0, 1], [2]]], # 2 lines - [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], [[0], [1], [2]]], # 2 lines + [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], None, [[2], [1], [0]]], # vertical + [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], None, [[2], [1], [0]]], # diagonal + [[[0, 0.5, 0.14, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], None, [[0, 1], [2]]], # same line, 2p + [[[0, 0.5, 0.18, 0.6], [0.2, 0.48, 0.35, 0.58], [0.8, 0.52, 0.9, 0.63]], None, [[0, 1], [2]]], # ~same line + [[[0, 0.3, 0.48, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], None, [[0, 1], [2]]], # 2 lines + [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], None, [[0], [1], [2]]], # 2 lines [ [ [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]], [[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]], ], + None, [[0], [1]], ], # rot + # Offset columns: a line does not lose a word to the next line of the other column + [_OFFSET_COLUMNS, (2480, 1753), [[11, 12, 13], [0, 1, 2, 3, 4], [14, 15, 16, 17, 18], [5, 6, 7, 8, 9, 10]]], + # A line rising by twice its height (perspective) + [_slanted_lines(-0.007, [0.485]) + _slanted_lines(0, [0.7]), None, [list(range(8)), list(range(8, 16))]], + # About 2 degrees of skew: the next line starts before the end of the current one + [_slanted_lines(0.0045, [0.3, 0.339]), None, [list(range(8)), list(range(8, 16))]], + [_slanted_lines(-0.0045, [0.3, 0.339]), None, [list(range(8)), list(range(8, 16))]], + # 3 columns whose baselines are slightly offset + [ + [[0.05, 0.19, 0.12, 0.21], [0.13, 0.19, 0.2, 0.21], [0.21, 0.19, 0.26, 0.21]] + + [[0.36, 0.2005, 0.44, 0.2205], [0.45, 0.2005, 0.5, 0.2205], [0.51, 0.2005, 0.56, 0.2205]] + + [[0.66, 0.194, 0.74, 0.214], [0.75, 0.197, 0.83, 0.217]], + None, + [[0, 1, 2], [6, 7], [3, 4, 5]], + ], + # A short word ending higher than its line is level with the label of the previous column (form) + [ + [[0.1179, 0.2285, 0.1956, 0.2461], [0.1956, 0.2295, 0.2241, 0.2451]] # "source OF" + + [[0.3122, 0.2383, 0.3549, 0.2529], [0.3575, 0.2383, 0.4327, 0.252], [0.4378, 0.2363, 0.5518, 0.2529]] + + [[0.557, 0.2373, 0.579, 0.249]], # "HALL RODECAP, NOBLESVILLE, IN" + (1000, 754), + [[0, 1], [2, 3, 4, 5]], + ], ], ) -def test_resolve_lines(input_boxes, lines): +def test_resolve_lines(input_boxes, shape, lines): doc_builder = builder.DocumentBuilder() - assert doc_builder._resolve_lines(np.asarray(input_boxes)) == lines - - -def test_resolve_lines_offset_columns(): - # Two columns whose baselines are offset by half a line (scanned page): the right column line must not lose a - # word to the next line of the left column - boxes = [ - [0.1422, 0.5908, 0.2472, 0.6045], # 0-4: left column, "Nichttragende Wände bestehen aus Kalk-" - [0.261, 0.5898, 0.3176, 0.6035], - [0.3314, 0.5898, 0.3978, 0.6035], - [0.413, 0.5928, 0.4406, 0.6035], - [0.4558, 0.5889, 0.4972, 0.6035], - [0.1394, 0.6055, 0.2099, 0.6182], # 5-10: left column, "sandstein oder Vollwanddielen bzw. aus stabilen" - [0.2126, 0.6064, 0.2486, 0.6172], - [0.2513, 0.6055, 0.366, 0.6191], - [0.3674, 0.6045, 0.4047, 0.6191], - [0.4074, 0.6084, 0.4364, 0.6191], - [0.4378, 0.6055, 0.4972, 0.6191], - [0.5539, 0.5811, 0.5953, 0.5957], # 11-13: right column, "Holz- oder Stahlkonstruktion." - [0.5953, 0.5811, 0.634, 0.5947], - [0.634, 0.5811, 0.7694, 0.5977], - [0.5539, 0.5967, 0.6755, 0.6133], # 14-18: right column, "Hauseingangstür als Holz-/Glaselement mit Si-/" - [0.6782, 0.5977, 0.6989, 0.6104], - [0.7045, 0.5977, 0.8426, 0.6113], - [0.844, 0.5977, 0.8758, 0.6133], - [0.8772, 0.5977, 0.9089, 0.6123], - ] - lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), (2480, 1753)) - assert sorted(lines) == [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 10], [11, 12, 13], [14, 15, 16, 17, 18]] - - -def test_resolve_lines_slanted_line(): - # A line rising to the right (perspective) by twice its height stays a single line, in reading order - boxes = [[0.05 + 0.11 * i, 0.485 - 0.007 * i, 0.15 + 0.11 * i, 0.515 - 0.007 * i] for i in range(8)] - boxes += [[0.05 + 0.11 * i, 0.70, 0.15 + 0.11 * i, 0.73] for i in range(8)] - assert builder.DocumentBuilder()._resolve_lines(np.asarray(boxes)) == [list(range(8)), list(range(8, 16))] - - -@pytest.mark.parametrize("slope", [0.0045, -0.0045]) # down / up to the right (about 2 degrees of skew) -def test_resolve_lines_slanted_lines_with_tight_leading(slope): - # The next line starts before the end of the current one: its first word must not take the last words of the - # current line, although it is better aligned with them than their own (lagging) neighbors - boxes = [ - [0.05 + 0.11 * i, y + slope * i, 0.15 + 0.11 * i, y + 0.03 + slope * i] for y in (0.30, 0.339) for i in range(8) - ] - lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes)) - assert sorted(lines) == [list(range(8)), list(range(8, 16))] - - -def test_resolve_lines_level_rows_of_other_columns(): - # 3 columns whose baselines are slightly offset: the words of the last column stay together, although each of - # them is level with the row of another column - boxes = [ - [0.05, 0.190, 0.12, 0.210], # 0-2: column A - [0.13, 0.190, 0.20, 0.210], - [0.21, 0.190, 0.26, 0.210], - [0.36, 0.2005, 0.44, 0.2205], # 3-5: column B - [0.45, 0.2005, 0.50, 0.2205], - [0.51, 0.2005, 0.56, 0.2205], - [0.66, 0.194, 0.74, 0.214], # 6-7: column C - [0.75, 0.197, 0.83, 0.217], - ] - lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes)) - assert sorted(lines) == [[0, 1, 2], [3, 4, 5], [6, 7]] + assert doc_builder._resolve_lines(np.asarray(input_boxes), shape) == lines def test_resolve_lines_layout_regions(): @@ -548,10 +537,6 @@ def test_resolve_lines_layout_regions(): regions=[{"boxes": np.asarray([[*p0, *p1] for p0, p1 in regions]), "class_names": ["Text"] * 3}], ) assert [len(line.words) for block in doc.pages[0].blocks for line in block.lines] == [3] * 6 - - -def test_resolve_lines_nested_and_short_regions(): - doc_builder = builder.DocumentBuilder() # A line of 8 words in a Text region, with a Picture box over its last word (or an inline Formula box) boxes = np.asarray([[0.10 + 0.1 * i, 0.20, 0.18 + 0.1 * i, 0.23] for i in range(8)]) text = ((0.08, 0.18), (0.92, 0.40)) @@ -566,7 +551,7 @@ def test_resolve_lines_nested_and_short_regions(): assert sorted(doc_builder._resolve_lines(boxes, word_regions=word_regions)) == [[0, 1, 2, 3], [4, 5, 6, 7]] -def test_resolve_lines_small_paragraph_break(): +def test_resolve_lines_break_distance(): # A table whose cells are separated by one word height: a small `paragraph_break` still splits the cells boxes = [] for row in range(3): @@ -580,8 +565,6 @@ def test_resolve_lines_small_paragraph_break(): assert len(builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), shape)) == 3 assert len(builder.DocumentBuilder(paragraph_break=0.01)._resolve_lines(np.asarray(boxes), shape)) == 12 - -def test_resolve_lines_large_type_on_a_narrow_page(): # Justified lines in large type on a narrow crop: a word gap wider than `paragraph_break` (relative to the page # width) but smaller than a word height does not split the line def _line(y, gaps): @@ -597,10 +580,7 @@ def _line(y, gaps): assert lines == [list(range(6)), list(range(6, 12)), list(range(12, 18))] # The page-relative paragraph break still applies on a page with regular type assert len(builder.DocumentBuilder()._resolve_lines(np.asarray(boxes) * [1, 0.2, 1, 0.2], (643, 468))) == 4 - - -def test_resolve_sub_lines_nested_box(): - # A box nested in a wider one (e.g. a punctuation mark detected twice) must not inflate the gap to the next word + # A box nested in a wider one (e.g. a punctuation mark detected twice) does not inflate the gap to the next word boxes = np.asarray([[0.10, 0.1, 0.30, 0.14], [0.25, 0.11, 0.26, 0.13], [0.32, 0.1, 0.45, 0.14]]) assert builder.DocumentBuilder()._resolve_sub_lines(boxes, [0, 1, 2], 0.05) == [[0, 1, 2]] assert builder.DocumentBuilder()._resolve_sub_lines(boxes, [0, 1, 2], 0.01) == [[0, 1], [2]] diff --git a/tests/common/test_models_reading_order.py b/tests/common/test_models_reading_order.py index e1df962bf1..f8a5a1e055 100644 --- a/tests/common/test_models_reading_order.py +++ b/tests/common/test_models_reading_order.py @@ -96,6 +96,12 @@ def test_sort_reading_order_columns(): geoms += [((0.1, 0.7 + 0.05 * i), (0.4, 0.73 + 0.05 * i)) for i in range(3)] geoms += [((0.6, 0.7 + 0.05 * i), (0.9, 0.73 + 0.05 * i)) for i in range(3)] assert sort_reading_order(geoms, column_voters=[True, False] + [True] * 6) == list(range(8)) + # A centered title crossing the gutter does not merge the columns: the column of a figure is read to its end + geoms = [((0.3, 0.02), (0.7, 0.06)), ((0.1, 0.1), (0.45, 0.13))] # title, line above the figure + geoms += [((0.1, 0.16), (0.13, 0.19)), ((0.25, 0.22), (0.3, 0.25))] # figure labels on 2 sides + geoms += [((0.1, 0.28 + 0.05 * i), (0.45, 0.31 + 0.05 * i)) for i in range(2)] # left column, below the figure + geoms += [((0.55, 0.1 + 0.07 * i), (0.9, 0.13 + 0.07 * i)) for i in range(4)] # right column + assert sort_reading_order(geoms) == list(range(10)) def test_sort_reading_order_float_members(): @@ -193,9 +199,6 @@ def test_resolve_reading_segments(): boxes = _two_columns() assert resolve_reading_segments(boxes) == [[0, 1, 2, 3], [4, 5, 6, 7]] assert resolve_reading_segments([]) == [] - - -def test_resolve_reading_segments_rows_and_regions(): # The fragments of a line split at a wide word gap belong to the paragraph of that line geoms = [((0.10, 0.10), (0.40, 0.13)), ((0.43, 0.10), (0.90, 0.13)), ((0.10, 0.14), (0.90, 0.17))] assert resolve_reading_segments(geoms) == [[0, 1, 2]] @@ -209,43 +212,6 @@ def test_resolve_reading_segments_rows_and_regions(): geoms = [((0.1, 0.1), (0.9, 0.13)), ((0.1, 0.14), (0.9, 0.17))] assert resolve_reading_segments(geoms, region_groups=[0, 0]) == [[0, 1]] assert resolve_reading_segments(geoms, region_groups=[0, 1]) == [[0], [1]] - - -def test_sort_reading_order_values_shifted_between_rows(): - # Receipt totals on a crumpled page (pixels): each amount sits lower than its label and also overlaps the label of - # the next row, but it overlaps its own label the most - geoms = [((60, 1200 + 60 * i), (600, 1240 + 60 * i)) for i in range(3)] # 0-2: full-width lines - geoms += [ - ((114, 1376), (282, 1420)), # 3: label - ((420, 1388), (520, 1431)), # 4: its amount - ((112, 1408), (226, 1451)), # 5: label - ((420, 1422), (520, 1465)), # 6: its amount - ((110, 1445), (162, 1483)), # 7: label - ((410, 1456), (520, 1502)), # 8: its amount - ((110, 1481), (162, 1518)), # 9: last label - ] - geoms += [((60, 1560 + 60 * i), (600, 1600 + 60 * i)) for i in range(2)] # 10-11: full-width lines - assert sort_reading_order(geoms) == list(range(12)) - # The cells of a single row are still read left to right when the middle one sits a little lower - row = [((100, 100), (200, 120)), ((300, 103), (400, 123)), ((500, 100), (600, 120))] - assert sort_reading_order(row) == [0, 1, 2] - assert sort_reading_order(row + [((100, 140), (200, 160)), ((300, 140), (400, 160))]) == [0, 1, 2, 3, 4] - - -def test_region_groups(): - regions = np.asarray([ - [0.0, 0.0, 0.5, 0.5], # 0 - [0.3, 0.3, 0.6, 0.6], # 1: overlaps 0 by 44% of its area - [0.7, 0.0, 0.9, 0.2], # 2: disjoint - [0.0, 0.0, 0.5, 0.5], # 3: same box as 0, other label - ]) - groups = _region_groups(regions, ["Text", "Text", "Text", "Picture"]) - assert groups[0] == groups[1] and len({groups[0], groups[2], groups[3]}) == 3 - # A small overlap does not group - assert len(set(_region_groups(regions[:2], ["Text"] * 2, min_overlap=0.5))) == 2 - - -def test_resolve_reading_segments_row_then_next_line(): # A line split into 3 fragments (words missed by the detection), then the short last line of its paragraph, # then 2 more paragraphs (pixels) boxes = [ @@ -426,6 +392,24 @@ def test_sort_reading_order_keeps_key_value_rows_together(): _box(0.05, 0.44, 0.95, 0.48), # 10 full-width ] assert sort_reading_order(geoms) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + # Receipt totals on a crumpled page (pixels): each amount sits lower than its label and also overlaps the label of + # the next row, but it overlaps its own label the most + geoms = [((60, 1200 + 60 * i), (600, 1240 + 60 * i)) for i in range(3)] # 0-2: full-width lines + geoms += [ + ((114, 1376), (282, 1420)), # 3: label + ((420, 1388), (520, 1431)), # 4: its amount + ((112, 1408), (226, 1451)), # 5: label + ((420, 1422), (520, 1465)), # 6: its amount + ((110, 1445), (162, 1483)), # 7: label + ((410, 1456), (520, 1502)), # 8: its amount + ((110, 1481), (162, 1518)), # 9: last label + ] + geoms += [((60, 1560 + 60 * i), (600, 1600 + 60 * i)) for i in range(2)] # 10-11: full-width lines + assert sort_reading_order(geoms) == list(range(12)) + # The cells of a single row are still read left to right when the middle one sits a little lower + row = [((100, 100), (200, 120)), ((300, 103), (400, 123)), ((500, 100), (600, 120))] + assert sort_reading_order(row) == [0, 1, 2] + assert sort_reading_order(row + [((100, 140), (200, 160)), ((300, 140), (400, 160))]) == [0, 1, 2, 3, 4] def test_sort_reading_order_follows_columns_when_a_gutter_exists(): @@ -461,6 +445,17 @@ def test_sort_reading_order_region_groups(): segments = resolve_reading_segments(geoms, labels=labels, region_groups=[-1, -1, 7, 7, 7]) flat = [idx for segment in segments for idx in segment] assert flat[flat.index(0) + 1 : flat.index(0) + 4] == [2, 3, 4] + # Overlapping regions of the same label are grouped + regions = np.asarray([ + [0.0, 0.0, 0.5, 0.5], # 0 + [0.3, 0.3, 0.6, 0.6], # 1: overlaps 0 by 44% of its area + [0.7, 0.0, 0.9, 0.2], # 2: disjoint + [0.0, 0.0, 0.5, 0.5], # 3: same box as 0, other label + ]) + groups = _region_groups(regions, ["Text", "Text", "Text", "Picture"]) + assert groups[0] == groups[1] and len({groups[0], groups[2], groups[3]}) == 3 + # A small overlap does not group + assert len(set(_region_groups(regions[:2], ["Text"] * 2, min_overlap=0.5))) == 2 def _paragraph_beside_a_figure(): From 8d792f377b6563f5f40d6f140ab2586010e2746d Mon Sep 17 00:00:00 2001 From: felix Date: Wed, 21 Oct 2026 11:14:44 +0200 Subject: [PATCH 3/4] Update --- doctr/io/exporters.py | 11 +- doctr/models/builder.py | 24 +++- doctr/models/reading_order/base.py | 26 +--- tests/common/test_io_exporters.py | 81 ++++++----- tests/common/test_models_builder.py | 202 ++++++++++++++++------------ 5 files changed, 184 insertions(+), 160 deletions(-) diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py index 0942ce8ccc..a2a7a6fa96 100644 --- a/doctr/io/exporters.py +++ b/doctr/io/exporters.py @@ -90,16 +90,7 @@ def _covering_region_indices( def _merge_overlapping_regions(elt_regions: list[int], region_geoms: list[Any], region_labels: list[str]) -> list[int]: - """Map the layout region of each element to its group of overlapping regions of the same label - - Args: - elt_regions: the layout region index of each element (-1 for none) - region_geoms: the geometries of the layout regions - region_labels: the labels of the layout regions - - Returns: - the region group of each element (-1 for none) - """ + """Map the layout region of each element (-1 for none) to its group of overlapping regions of the same label""" from doctr.models.reading_order.base import _region_groups, _to_boxes if len(region_geoms) < 2: diff --git a/doctr/models/builder.py b/doctr/models/builder.py index d8218d6540..ec9c31da0f 100644 --- a/doctr/models/builder.py +++ b/doctr/models/builder.py @@ -149,7 +149,7 @@ def _resolve_sub_lines( # Compute distance between the sub-line and the box dist = boxes[i, 0] - right - # If distance between boxes is lower than the break distance, same sub-line + # If the distance is lower than the break distance, same sub-line if dist < break_dist: horiz_break = False # Unless the box and the sub-line lie in distinct regions (e.g. columns separated by a narrow gutter) @@ -237,7 +237,8 @@ def _resolve_rows(boxes: np.ndarray, idxs: np.ndarray, y_med: float) -> list[lis """Group straight boxes into visual rows A box joins the row whose horizontally closest boxes are the best aligned with it (less than half a median - height apart), which follows slanted lines and keeps the rows of offset columns apart. + height apart), which follows slanted lines and keeps the rows of offset columns apart. A box right under or over + one of them, away from the others (e.g. stacked figure labels), does not join the row. Args: boxes: straight bounding boxes of shape (N, 4) @@ -245,16 +246,27 @@ def _resolve_rows(boxes: np.ndarray, idxs: np.ndarray, y_med: float) -> list[lis y_med: the median box height Returns: - the rows (lists of box indices), from top to bottom + the rows (lists of box indices) """ if not y_med > 0: return [[idx] for idx in idxs.tolist()] lefts, rights = boxes[:, 0].tolist(), boxes[:, 2].tolist() centers = ((boxes[:, 1] + boxes[:, 3]) / 2).tolist() + heights = (boxes[:, 3] - boxes[:, 1]).tolist() row_lefts: list[list[float]] = [] # sorted left edges of the boxes of each row row_boxes: list[list[int]] = [] # boxes of each row, sorted by left edge spans: list[list[float]] = [] # highest and lowest box center of each row + def _stacked(idx: int, other: int, median: float) -> bool: + """Whether two full-height boxes lie one over the other, `other` being level with its row""" + overlap = min(rights[idx], rights[other]) - max(lefts[idx], lefts[other]) + return ( + overlap >= 0.5 * min(rights[idx] - lefts[idx], rights[other] - lefts[other]) + and abs(centers[idx] - centers[other]) >= y_med / 2 + and abs(centers[other] - median) < y_med / 2 + and min(heights[idx], heights[other]) >= 0.75 * y_med + ) + def _best_row(idx: int, rows: list[int]) -> int: center, left, right = centers[idx], lefts[idx], rights[idx] best, best_key = -1, (float("inf"), float("inf")) @@ -267,7 +279,11 @@ def _best_row(idx: int, rows: list[int]) -> int: (max(lefts[other] - right, left - rights[other], 0.0), other) for other in row_boxes[row][max(pos - 3, 0) : pos + 3] )[:3] - offset = abs(statistics.median(centers[other] for _, other in nearest) - center) + median = statistics.median(centers[other] for _, other in nearest) + stacked = [_stacked(idx, other, median) for _, other in nearest] + if any(stacked) and all(dist >= y_med for (dist, _), flag in zip(nearest, stacked) if not flag): + continue + offset = abs(median - center) # The distance breaks ties on slanted lines, where the next line can start before this one ends key = (offset + _ROW_DISTANCE_WEIGHT * nearest[0][0], nearest[0][0]) if offset < y_med / 2 and key < best_key: diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index a2b029c020..56061646e5 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -44,8 +44,8 @@ _FURNITURE_ROLES = ("header", "footer", "footnote") -# Label suffix of the text lying in a float region without float element (e.g. a table without structure) -_FLOAT_TEXT_SUFFIX = " text" +# Normalized labels of the text lying in a float region without float element (e.g. a table without structure) +_FLOAT_TEXT_LABELS = {f"{label}_text" for label in _FLOAT_LABELS} def normalize_layout_label(label: str | None) -> str: @@ -87,13 +87,7 @@ def layout_label_role(label: str | None) -> str: def _float_text_label(label: str) -> str: """Return the label of a text element lying in a float region (e.g. 'Table' -> 'Table text')""" - return f"{label}{_FLOAT_TEXT_SUFFIX}" - - -def _is_float_text(label: str | None) -> bool: - """Check whether a label is the one of a text element lying in a float region""" - norm = normalize_layout_label(label) - return any(norm == normalize_layout_label(_float_text_label(float_label)) for float_label in _FLOAT_LABELS) + return f"{label} text" def detect_text_direction(texts: Iterable[str], language: str | None = None) -> str: @@ -396,7 +390,7 @@ def _defer_floats( inner = {idx for idcs in members.values() for idx in idcs} def _is_float(idx: int) -> bool: - return labels[idx] in _FLOAT_LABELS or _is_float_text(labels[idx]) or idx in inner + return labels[idx] in _FLOAT_LABELS or labels[idx] in _FLOAT_TEXT_LABELS or idx in inner out = list(order) pos = 0 @@ -819,15 +813,9 @@ def assign_layout_labels( def _region_groups(regions: np.ndarray, labels: Sequence[str | None], min_overlap: float = 0.1) -> list[int]: - """Group the overlapping layout regions of the same label (e.g. two regions for a paragraph around a figure) - - Args: - regions: the (R, 4) boxes of the layout regions - labels: the label of each region - min_overlap: minimum intersection of two regions, relative to the area of the smaller one, to group them - - Returns: - the group of each region (the index of one of the regions of the group) + """Return the group of each (R, 4) layout region (the index of one of its members), grouping the regions of the + same label which overlap by more than `min_overlap` of the smaller one (e.g. two regions for a paragraph around a + figure) """ norm = [normalize_layout_label(label) for label in labels] areas = (regions[:, 2] - regions[:, 0]) * (regions[:, 3] - regions[:, 1]) diff --git a/tests/common/test_io_exporters.py b/tests/common/test_io_exporters.py index bb9e57c47c..a3ff6d7011 100644 --- a/tests/common/test_io_exporters.py +++ b/tests/common/test_io_exporters.py @@ -140,7 +140,44 @@ def test_page_export_as_markdown_rtl(): assert page.export_as_markdown(direction="ltr").startswith("الأيسر") -def test_page_reading_order_table_region_text(): +def test_page_export_with_tables(): + cells = [ + elements.TableCell("Name", 0.9, ((0.1, 0.55), (0.4, 0.6)), 0, 0, 0, 0), + elements.TableCell("Qty", 0.9, ((0.4, 0.55), (0.7, 0.6)), 0, 0, 1, 1), + elements.TableCell("Bolt", 0.9, ((0.1, 0.6), (0.4, 0.65)), 1, 1, 0, 0), + elements.TableCell("12|3", 0.9, ((0.4, 0.6), (0.7, 0.65)), 1, 1, 1, 1), + ] + table = elements.Table(cells, 2, 2, ((0.1, 0.55), (0.7, 0.65)), 0.95) + lines = [ + _line_at("before the table", 0.1, 0.1, 0.9, 0.14), + _line_at("after the table", 0.1, 0.7, 0.9, 0.74), + ] + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), tables=[table] + ) + markdown = page.export_as_markdown() + assert markdown.split("\n\n") == [ + "before the table", + "| Name | Qty |\n| --- | --- |\n| Bolt | 12\\|3 |", + "after the table", + ] + asciidoc = page.export_as_asciidoc() + assert "|===\n|Name |Qty\n\n|Bolt |12\\|3\n|===" in asciidoc + assert asciidoc.index("before the table") < asciidoc.index("|===") < asciidoc.index("after the table") + # A page-wide table does not prevent the columns above and below it from being detected + lines = [_line_at("A title over both columns", 0.2, 0.02, 0.8, 0.06)] + for col, x in (("left", 0.05), ("right", 0.55)): + lines += [_line_at(f"{col} {i}", x, 0.1 + 0.05 * i, x + 0.4, 0.13 + 0.05 * i) for i in range(3)] + lines += [_line_at(f"{col} {i}", x, 0.3 + 0.05 * i, x + 0.4, 0.33 + 0.05 * i) for i in range(3, 6)] + table = elements.Table(cells, 2, 2, ((0.05, 0.25), (0.95, 0.28)), 0.95) + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), tables=[table] + ) + blocks = [item for item in page.items_in_reading_order() if isinstance(item, elements.Block)] + texts = [text for item in _texts(blocks) for text in item] + # Each section is read column after column + sections = [[f"{col} {i}" for col in ("left", "right") for i in idcs] for idcs in (range(3), range(3, 6))] + assert texts == ["A title over both columns", *sections[0], *sections[1]] # A table region from the layout model, without table structure (`detect_tables=False`): its lines keep the # region label, but they are not floats themselves, so they are neither nested into each other nor isolated lines = [ @@ -157,7 +194,7 @@ def test_page_reading_order_table_region_text(): layout=[elements.LayoutElement("Table", 0.9, ((0.08, 0.38), (0.92, 0.54)))], ) items, labels, _ = page_reading_order(page) - texts = [" ".join(word.value for word in line.words) for item in items for line in item.lines] + texts = [text for item in _texts(items) for text in item] assert texts == ["Pos Datum Einzelpreis Gesamt", "pro UE", "EUR", "1 Deutsch 22,50 22,50"] assert labels == ["Table"] * len(items) assert len(items) < len(lines) @@ -188,46 +225,6 @@ def test_page_reading_order_table_region_text(): assert {label for item, label in zip(_texts(items), labels) if item[0].startswith("Row")} == {"Table"} -def test_page_export_with_tables(): - cells = [ - elements.TableCell("Name", 0.9, ((0.1, 0.55), (0.4, 0.6)), 0, 0, 0, 0), - elements.TableCell("Qty", 0.9, ((0.4, 0.55), (0.7, 0.6)), 0, 0, 1, 1), - elements.TableCell("Bolt", 0.9, ((0.1, 0.6), (0.4, 0.65)), 1, 1, 0, 0), - elements.TableCell("12|3", 0.9, ((0.4, 0.6), (0.7, 0.65)), 1, 1, 1, 1), - ] - table = elements.Table(cells, 2, 2, ((0.1, 0.55), (0.7, 0.65)), 0.95) - lines = [ - _line_at("before the table", 0.1, 0.1, 0.9, 0.14), - _line_at("after the table", 0.1, 0.7, 0.9, 0.74), - ] - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), tables=[table] - ) - markdown = page.export_as_markdown() - assert markdown.split("\n\n") == [ - "before the table", - "| Name | Qty |\n| --- | --- |\n| Bolt | 12\\|3 |", - "after the table", - ] - asciidoc = page.export_as_asciidoc() - assert "|===\n|Name |Qty\n\n|Bolt |12\\|3\n|===" in asciidoc - assert asciidoc.index("before the table") < asciidoc.index("|===") < asciidoc.index("after the table") - # A page-wide table does not prevent the columns above and below it from being detected - lines = [_line_at("A title over both columns", 0.2, 0.02, 0.8, 0.06)] - for col, x in (("left", 0.05), ("right", 0.55)): - lines += [_line_at(f"{col} {i}", x, 0.1 + 0.05 * i, x + 0.4, 0.13 + 0.05 * i) for i in range(3)] - lines += [_line_at(f"{col} {i}", x, 0.3 + 0.05 * i, x + 0.4, 0.33 + 0.05 * i) for i in range(3, 6)] - table = elements.Table(cells, 2, 2, ((0.05, 0.25), (0.95, 0.28)), 0.95) - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), tables=[table] - ) - blocks = [item for item in page.items_in_reading_order() if isinstance(item, elements.Block)] - texts = [text for item in _texts(blocks) for text in item] - # Each section is read column after column - sections = [[f"{col} {i}" for col in ("left", "right") for i in idcs] for idcs in (range(3), range(3, 6))] - assert texts == ["A title over both columns", *sections[0], *sections[1]] - - def test_page_export_as_asciidoc(): page = _reading_order_page() asciidoc = page.export_as_asciidoc() diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index 16bd7c05b1..0c3b801929 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -95,6 +95,10 @@ def test_documentbuilder(): assert out.pages[0].blocks[0].lines[0].words[-1].geometry == ((0.45, 0.5), (0.6, 0.65)) assert out.pages[0].blocks[0].lines[0].words[-1].objectness_score == 0.99 + # A small `paragraph_break` still splits table cells separated by one word height + small_break = builder.DocumentBuilder(paragraph_break=0.01) + assert len(small_break._resolve_lines(np.asarray(_table_rows()), (2339, 1654))) == 12 + # Repr assert ( repr(doc_builder) == "DocumentBuilder(resolve_lines=True, " @@ -254,6 +258,53 @@ def test_documentbuilder_layout(): assert region.type == "Table" assert isinstance(region.geometry, tuple) and len(region.geometry) == 4 + # 3 columns separated by gutters as narrow as the word gaps of the justified lines + boxes, columns = [], [] + for col in range(3): + x0 = 0.02 + 0.33 * col + columns.append(((x0 - 0.01, 0.05), (x0 + 0.31, 0.95))) + for row in range(2): + for word in range(3): + boxes.append([x0 + 0.105 * word, 0.4 + 0.05 * row, x0 + 0.095 + 0.105 * word, 0.43 + 0.05 * row]) + boxes = np.asarray(boxes) + # Without layout, each row spans the 3 columns + assert len(doc_builder._resolve_lines(boxes)) == 2 + # With a layout region per column, a line never spans two of them + word_regions = doc_builder._word_regions(boxes, columns, ["Text"] * 3) + lines = doc_builder._resolve_lines(boxes, word_regions=word_regions) + assert sorted(lines) == [[6 * col + 3 * row + word for word in range(3)] for col in range(3) for row in range(2)] + # Overlapping regions of the same label (duplicates of a layout model) are one region + overlapping = [((0.0, 0.0), (0.6, 1.0)), ((0.3, 0.0), (1.0, 1.0))] + assert ( + len(doc_builder._resolve_lines(boxes, word_regions=doc_builder._word_regions(boxes, overlapping, ["Text"] * 2))) + == 2 + ) + # Words outside any region do not split a line + assert doc_builder._word_regions(boxes, [((0.5, 0.0), (0.6, 0.1))], ["Text"])[0].tolist() == [-1] * 18 + # End to end + doc = doc_builder( + [np.zeros((100, 100, 3), dtype=np.uint8)], + [boxes], + [np.ones(18)], + [[("word", 0.9)] * 18], + [(100, 100)], + [[{"value": 0, "confidence": None}] * 18], + regions=[{"boxes": np.asarray([[*p0, *p1] for p0, p1 in columns]), "class_names": ["Text"] * 3}], + ) + assert [len(line.words) for block in doc.pages[0].blocks for line in block.lines] == [3] * 6 + # A line of 8 words in a Text region, with a Picture box over its last word (or an inline Formula box) + boxes = np.asarray([[0.10 + 0.1 * i, 0.20, 0.18 + 0.1 * i, 0.23] for i in range(8)]) + text = ((0.08, 0.18), (0.92, 0.40)) + for nested, label in [(((0.78, 0.15), (0.95, 0.24)), "Picture"), (((0.38, 0.19), (0.59, 0.24)), "Formula")]: + word_regions = doc_builder._word_regions(boxes, [text, nested], ["Text", label]) + assert doc_builder._resolve_lines(boxes, word_regions=word_regions) == [list(range(8))] + # Two columns whose regions start slightly below the first line of the right column (its words are 46% covered) + boxes = np.asarray([[0.05 + 0.1 * i, 0.30, 0.13 + 0.1 * i, 0.33] for i in range(4)]) + boxes = np.concatenate([boxes, boxes + [0.42, 0.0, 0.42, 0.0]]) + columns = [((0.03, 0.28), (0.45, 0.9)), ((0.46, 0.3162), (0.89, 0.9))] + word_regions = doc_builder._word_regions(boxes, columns, ["Text", "Text"]) + assert sorted(doc_builder._resolve_lines(boxes, word_regions=word_regions)) == [[0, 1, 2, 3], [4, 5, 6, 7]] + def _table_cell(x0, y0, x1, y1, rs, re, cs, ce, score=0.9): return { @@ -454,6 +505,30 @@ def _slanted_lines(slope, ys): return [[0.05 + 0.11 * i, y + slope * i, 0.15 + 0.11 * i, y + 0.03 + slope * i] for y in ys for i in range(8)] +def _table_rows(): + # 3 rows of 4 cells (2 words each), the cells being separated by one word height + boxes = [] + for row in range(3): + x, y = 0.1, 0.3 + 0.03 * row + for _ in range(4): + for _ in range(2): + boxes.append([x, y, x + 0.05, y + 0.015]) + x += 0.054 + x += 0.011 + return boxes + + +def _justified_lines(scale=1.0): + boxes = [] + for y, gaps in [(0.10, [0.019, 0.019, 0.040, 0.019, 0.019]), (0.16, [0.019] * 5), (0.22, [0.019] * 5)]: + x = 0.05 + for gap in [0.0, *gaps]: + x += gap + boxes.append([x, y * scale, x + 0.12, (y + 0.04) * scale]) + x += 0.12 + return boxes + + @pytest.mark.parametrize( "input_boxes, shape, lines", [ @@ -494,6 +569,48 @@ def _slanted_lines(slope, ys): (1000, 754), [[0, 1], [2, 3, 4, 5]], ], + # Table cells separated by one word height: one line per row with the default `paragraph_break` + [_table_rows(), (2339, 1654), [list(range(8 * row, 8 * row + 8)) for row in range(3)]], + # Large type on a narrow crop: a word gap wider than `paragraph_break` but smaller than a word height + [_justified_lines(), (643, 468), [list(range(6)), list(range(6, 12)), list(range(12, 18))]], + # The page-relative paragraph break still applies to regular type + [_justified_lines(0.2), (643, 468), [[0, 1, 2], [3, 4, 5], list(range(6, 12)), list(range(12, 18))]], + # A box nested in a wider one (e.g. a punctuation mark detected twice) does not inflate the gap to the next word + [[[0.10, 0.1, 0.30, 0.14], [0.25, 0.11, 0.26, 0.13], [0.32, 0.1, 0.45, 0.14]], None, [[0, 1, 2]]], + # A figure label right below another one is not part of its row, even when level with a row further right + [ + [[0.6327, 0.2393, 0.6656, 0.2529], [0.6289, 0.2471, 0.6693, 0.2607]] # "Image" above "Encoder" + + [[0.7098, 0.2432, 0.7199, 0.2568], [0.7401, 0.2432, 0.7629, 0.2568]], # matrix cells + (1584, 1224), + [[0], [2, 3], [1]], + ], + # A word overlapping the previous one a little, a superscript or the parts of a fraction stay in their line + [ + [[0.1, 0.5, 0.2, 0.53], [0.24, 0.5, 0.34, 0.53], [0.38, 0.49, 0.48, 0.52], [0.47, 0.51, 0.57, 0.54]], + None, + [[0], [1], [2, 3]], + ], + [ + [[0.1, 0.5, 0.2, 0.53], [0.24, 0.5, 0.34, 0.53], [0.38, 0.51, 0.48, 0.54], [0.465, 0.499, 0.485, 0.511]], + None, + [[0], [1], [2, 3]], + ], + [ + [[0.1, 0.5, 0.2, 0.53], [0.205, 0.5, 0.3, 0.53], [0.305, 0.49, 0.335, 0.516], [0.31, 0.514, 0.34, 0.54]] + + [[0.345, 0.5, 0.44, 0.53]], + None, + [list(range(5))], + ], + # Nested boxes do not inflate the word gaps from which the break distance is estimated + [ + [[0.1 + 0.045 * i + 0.02 * (i >= 5), 0.5, 0.14 + 0.045 * i + 0.02 * (i >= 5), 0.51] for i in range(10)] + + [ + [0.11 + 0.045 * i + 0.02 * (i >= 5), 0.5, 0.12 + 0.045 * i + 0.02 * (i >= 5), 0.51] + for i in (0, 1, 2, 5, 6, 7) + ], + None, + [[0, 10, 1, 11, 2, 12, 3, 4], [5, 13, 6, 14, 7, 15, 8, 9]], + ], ], ) def test_resolve_lines(input_boxes, shape, lines): @@ -501,91 +618,6 @@ def test_resolve_lines(input_boxes, shape, lines): assert doc_builder._resolve_lines(np.asarray(input_boxes), shape) == lines -def test_resolve_lines_layout_regions(): - # 3 columns separated by gutters as narrow as the word gaps of the justified lines - boxes, regions = [], [] - for col in range(3): - x0 = 0.02 + 0.33 * col - regions.append(((x0 - 0.01, 0.05), (x0 + 0.31, 0.95))) - for row in range(2): - for word in range(3): - boxes.append([x0 + 0.105 * word, 0.4 + 0.05 * row, x0 + 0.095 + 0.105 * word, 0.43 + 0.05 * row]) - boxes = np.asarray(boxes) - doc_builder = builder.DocumentBuilder() - # Without layout, each row spans the 3 columns - assert len(doc_builder._resolve_lines(boxes)) == 2 - # With a layout region per column, a line never spans two of them - word_regions = doc_builder._word_regions(boxes, regions, ["Text"] * 3) - lines = doc_builder._resolve_lines(boxes, word_regions=word_regions) - assert sorted(lines) == [[6 * col + 3 * row + word for word in range(3)] for col in range(3) for row in range(2)] - # Overlapping regions of the same label (duplicates of a layout model) are one region - overlapping = [((0.0, 0.0), (0.6, 1.0)), ((0.3, 0.0), (1.0, 1.0))] - assert ( - len(doc_builder._resolve_lines(boxes, word_regions=doc_builder._word_regions(boxes, overlapping, ["Text"] * 2))) - == 2 - ) - # Words outside any region do not split a line - assert doc_builder._word_regions(boxes, [((0.5, 0.0), (0.6, 0.1))], ["Text"])[0].tolist() == [-1] * 18 - # End to end - doc = doc_builder( - [np.zeros((100, 100, 3), dtype=np.uint8)], - [boxes], - [np.ones(18)], - [[("word", 0.9)] * 18], - [(100, 100)], - [[{"value": 0, "confidence": None}] * 18], - regions=[{"boxes": np.asarray([[*p0, *p1] for p0, p1 in regions]), "class_names": ["Text"] * 3}], - ) - assert [len(line.words) for block in doc.pages[0].blocks for line in block.lines] == [3] * 6 - # A line of 8 words in a Text region, with a Picture box over its last word (or an inline Formula box) - boxes = np.asarray([[0.10 + 0.1 * i, 0.20, 0.18 + 0.1 * i, 0.23] for i in range(8)]) - text = ((0.08, 0.18), (0.92, 0.40)) - for nested in [((0.78, 0.15), (0.95, 0.24)), ((0.38, 0.19), (0.59, 0.24))]: - word_regions = doc_builder._word_regions(boxes, [text, nested], ["Text", "Picture"]) - assert doc_builder._resolve_lines(boxes, word_regions=word_regions) == [list(range(8))] - # Two columns whose regions start slightly below the first line of the right column (its words are 46% covered) - boxes = np.asarray([[0.05 + 0.1 * i, 0.30, 0.13 + 0.1 * i, 0.33] for i in range(4)]) - boxes = np.concatenate([boxes, boxes + [0.42, 0.0, 0.42, 0.0]]) - regions = [((0.03, 0.28), (0.45, 0.9)), ((0.46, 0.3162), (0.89, 0.9))] - word_regions = doc_builder._word_regions(boxes, regions, ["Text", "Text"]) - assert sorted(doc_builder._resolve_lines(boxes, word_regions=word_regions)) == [[0, 1, 2, 3], [4, 5, 6, 7]] - - -def test_resolve_lines_break_distance(): - # A table whose cells are separated by one word height: a small `paragraph_break` still splits the cells - boxes = [] - for row in range(3): - x, y = 0.1, 0.3 + 0.03 * row - for _ in range(4): - for _ in range(2): - boxes.append([x, y, x + 0.05, y + 0.015]) - x += 0.054 - x += 0.011 - shape = (2339, 1654) - assert len(builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), shape)) == 3 - assert len(builder.DocumentBuilder(paragraph_break=0.01)._resolve_lines(np.asarray(boxes), shape)) == 12 - - # Justified lines in large type on a narrow crop: a word gap wider than `paragraph_break` (relative to the page - # width) but smaller than a word height does not split the line - def _line(y, gaps): - x, boxes = 0.05, [] - for gap in [0.0, *gaps]: - x += gap - boxes.append([x, y, x + 0.12, y + 0.04]) - x += 0.12 - return boxes - - boxes = _line(0.10, [0.019, 0.019, 0.040, 0.019, 0.019]) + _line(0.16, [0.019] * 5) + _line(0.22, [0.019] * 5) - lines = builder.DocumentBuilder()._resolve_lines(np.asarray(boxes), (643, 468)) - assert lines == [list(range(6)), list(range(6, 12)), list(range(12, 18))] - # The page-relative paragraph break still applies on a page with regular type - assert len(builder.DocumentBuilder()._resolve_lines(np.asarray(boxes) * [1, 0.2, 1, 0.2], (643, 468))) == 4 - # A box nested in a wider one (e.g. a punctuation mark detected twice) does not inflate the gap to the next word - boxes = np.asarray([[0.10, 0.1, 0.30, 0.14], [0.25, 0.11, 0.26, 0.13], [0.32, 0.1, 0.45, 0.14]]) - assert builder.DocumentBuilder()._resolve_sub_lines(boxes, [0, 1, 2], 0.05) == [[0, 1, 2]] - assert builder.DocumentBuilder()._resolve_sub_lines(boxes, [0, 1, 2], 0.01) == [[0, 1], [2]] - - def test_points_in_polygons(): polys = np.array( [ From e25d4871ffad9cf2665752447090664cd9ee672d Mon Sep 17 00:00:00 2001 From: felix Date: Mon, 26 Oct 2026 10:21:21 +0100 Subject: [PATCH 4/4] Update --- doctr/io/exporters.py | 12 ++++- doctr/models/reading_order/base.py | 60 ++++++++++++++++++++++- tests/common/test_io_exporters.py | 31 ++++++++++++ tests/common/test_models_reading_order.py | 32 ++++++++++++ 4 files changed, 132 insertions(+), 3 deletions(-) diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py index a2a7a6fa96..c78409aadc 100644 --- a/doctr/io/exporters.py +++ b/doctr/io/exporters.py @@ -216,12 +216,20 @@ def page_reading_order( elt_labels[idx] = "Table" elif not isinstance(elt, LayoutElement) and elt_regions[idx] in furniture_labels: elt_labels[idx] = furniture_labels[elt_regions[idx]] - # Page-wide figures and tables do not vote in the multi-column detection + # Page-wide figures and tables do not vote in the multi-column detection, nor do the text lines of a page-wide table + # region: its cells (e.g. the amounts of a receipt detected as a table) would read as columns boxes = _to_boxes(elt_geoms) span = float(boxes[:, 2].max() - boxes[:, 0].min()) or 1.0 + region_boxes = _to_boxes(region_geoms) if len(region_geoms) > 0 else np.zeros((0, 4)) + wide_tables = { + reg + for reg, (label, box) in enumerate(zip(region_labels, region_boxes)) + if normalize_layout_label(label) == "table" and box[2] - box[0] > 0.5 * span + } column_voters = [ not (isinstance(elt, (LayoutElement, Table)) and box[2] - box[0] > 0.5 * span) - for elt, box in zip(elements, boxes) + and (isinstance(elt, (LayoutElement, Table)) or reg not in wide_tables) + for elt, box, reg in zip(elements, boxes, elt_regions) ] # A text line in a table or figure region keeps its label, but is grouped into paragraphs like regular text order_labels = [ diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index 56061646e5..cb10422ce6 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -47,6 +47,14 @@ # Normalized labels of the text lying in a float region without float element (e.g. a table without structure) _FLOAT_TEXT_LABELS = {f"{label}_text" for label in _FLOAT_LABELS} +# Key-value rows (receipts, invoices, forms): a "column" of short values, each level with the text it belongs to and +# far from it, is not a text column. Below this number of column voters, such a split does not make a multi-column page +_KEY_VALUE_MAX_VOTERS = 60 +# Maximum median aspect ratio (width / height) of the values: a number or a short label, not a line of running text +_KEY_VALUE_MAX_ASPECT = 8.0 +# Minimum median gap between a value and the text level with it, in median element heights (a gutter is narrower) +_KEY_VALUE_MIN_GAP = 4.0 + def normalize_layout_label(label: str | None) -> str: """Normalize a layout label to lower snake case (e.g. 'Page-header' -> 'page_header') @@ -187,6 +195,52 @@ def _strict_rank(primary: np.ndarray, secondary: np.ndarray) -> np.ndarray: return ranks +def _key_value_rows( + boxes: np.ndarray, voters: np.ndarray, split: float, y_overlap: np.ndarray, y_overlap_threshold: float, span: float +) -> bool: + """Check whether the voters on the right of a column split are the values of key-value rows rather than a column + + On a receipt or an invoice, the amounts form a narrow right-aligned "column": each amount is level with the item + it belongs to, far from it, and the rows must be read one after the other. In a text column, the lines are long + and the gutter separating them from the lines of the previous column is narrow. + + Args: + boxes: the (N, 4) boxes in canonical LTR space + voters: the mask of the boxes used to detect the columns + split: the abscissa of the column split + y_overlap: the (N, N) vertical overlap ratios of the boxes + y_overlap_threshold: the vertical overlap above which two boxes lie on the same row + span: the horizontal extent of the voters + + Returns: + whether the split separates the keys from the values of key-value rows + """ + num_voters = int(np.count_nonzero(voters)) + if num_voters > _KEY_VALUE_MAX_VOTERS: + return False + x0, y0, x1, y1 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] + left_side = voters & ((x0 + x1) / 2 <= split) + right = np.flatnonzero(voters & ~left_side) + left = np.flatnonzero(left_side) + if right.size == 0 or left.size == 0: + return False + height = float(np.median(y1[voters] - y0[voters])) + if not height > 0: + return False + aspect = float(np.median((x1[right] - x0[right]) / np.maximum(y1[right] - y0[right], 1e-9))) + if aspect >= _KEY_VALUE_MAX_ASPECT: + return False + # Gap between each value and the text level with it on its left + level = y_overlap[np.ix_(right, left)] > y_overlap_threshold + level &= x1[left][None, :] <= x0[right][:, None] + 0.02 * span + paired = level.any(axis=1) + if np.count_nonzero(paired) < 0.5 * right.size: + return False + ends = np.where(level, x1[left][None, :], -np.inf).max(axis=1) + gaps = (x0[right][paired] - ends[paired]) / height + return float(np.median(gaps)) >= _KEY_VALUE_MIN_GAP + + def _topological_order( boxes: np.ndarray, x_overlap_threshold: float, y_overlap_threshold: float, voters: np.ndarray | None = None ) -> list[int]: @@ -201,7 +255,9 @@ def _topological_order( The relations are resolved with Kahn's algorithm; among the available elements, the traversal favors the continuation of the current column (the closest element below the last emitted one with a horizontal overlap), which keeps multi-column bodies intact even for non-Manhattan layouts where recursive XY-cuts - fail to find a valid split. `voters` optionally restricts the boxes used to detect a multi-column page. + fail to find a valid split. `voters` optionally restricts the boxes used to detect a multi-column page. A split + separating the values of key-value rows from their keys (e.g. the amounts of a receipt, cf. `_key_value_rows`) + does not make the page multi-column. """ num_boxes = boxes.shape[0] if num_boxes <= 1: @@ -262,6 +318,8 @@ def _topological_order( crossing = int(np.count_nonzero(np.minimum(vx1 - split, split - vx0) > 0.02 * span)) left = int(np.count_nonzero(centers <= split)) if crossing <= tolerance and left >= 0.25 * num_voters and num_voters - left >= 0.25 * num_voters: + if _key_value_rows(boxes, voters, float(split), y_overlap, y_overlap_threshold, span): + break # rows of keys and values (e.g. a receipt): read row after row multi_column = True bridging |= np.minimum(x1 - split, split - x0) > 0.02 * span break diff --git a/tests/common/test_io_exporters.py b/tests/common/test_io_exporters.py index a3ff6d7011..945ea32c8e 100644 --- a/tests/common/test_io_exporters.py +++ b/tests/common/test_io_exporters.py @@ -225,6 +225,37 @@ def test_page_export_with_tables(): assert {label for item, label in zip(_texts(items), labels) if item[0].startswith("Row")} == {"Table"} +def test_page_reading_order_page_wide_table_region(): + # A receipt detected as one page-wide table under its header: the cells left over by a small recognized table are + # rows of that table, not two text columns, whether or not the page-wide table element votes in the column detection + lines = [ + _line_at("Corner Shop", 0.05, 0.010, 0.25, 0.030), + _line_at("12 Main Street, Springfield", 0.05, 0.032, 0.40, 0.052), + _line_at("2026-10-26 09:13", 0.05, 0.054, 0.30, 0.074), + ] + for i in range(8): + lines.append(_line_at(f"Item {i} description", 0.05, 0.20 + 0.05 * i, 0.40, 0.23 + 0.05 * i)) + lines.append(_line_at(f"Remark for item {i} written out", 0.55, 0.20 + 0.05 * i, 0.90, 0.23 + 0.05 * i)) + layout = [elements.LayoutElement("Table", 0.9, ((0.03, 0.09), (0.97, 0.95)))] + cells = [ + elements.TableCell("Header", 0.9, ((0.04, 0.10), (0.96, 0.13)), 0, 0, 0, 0), + elements.TableCell("Subheader", 0.9, ((0.04, 0.13), (0.96, 0.16)), 1, 1, 0, 0), + ] + expected = [text for i in range(8) for text in (f"Item {i} description", f"Remark for item {i} written out")] + for tables in ([], [elements.Table(cells, 2, 1, ((0.04, 0.10), (0.96, 0.16)), 0.9)]): + page = elements.Page( + np.zeros((10, 10, 3), dtype=np.uint8), + [elements.Block(lines=lines)], + 0, + (1000, 800), + layout=layout, + tables=tables, + ) + items, _, _ = page_reading_order(page) + texts = [text for item in _texts(item for item in items if isinstance(item, elements.Block)) for text in item] + assert texts == ["Corner Shop", "12 Main Street, Springfield", "2026-10-26 09:13", *expected] + + def test_page_export_as_asciidoc(): page = _reading_order_page() asciidoc = page.export_as_asciidoc() diff --git a/tests/common/test_models_reading_order.py b/tests/common/test_models_reading_order.py index f8a5a1e055..f0ad29258d 100644 --- a/tests/common/test_models_reading_order.py +++ b/tests/common/test_models_reading_order.py @@ -412,6 +412,38 @@ def test_sort_reading_order_keeps_key_value_rows_together(): assert sort_reading_order(row + [((100, 140), (200, 160)), ((300, 140), (400, 160))]) == [0, 1, 2, 3, 4] +def test_sort_reading_order_receipt_amounts_are_not_a_column(): + # Lines of a real receipt (CORD): the right-aligned amounts are level with their item, far from it. Once the item + # fragments are merged into lines, the amounts are a quarter of the elements, but they are not a text column + geoms = [ + ((0.096, 0.409), (0.478, 0.439)), # 0: Enting2 Dua Bandeng Ori + ((0.650, 0.407), (0.753, 0.431)), # 1: 31.000 + ((0.094, 0.444), (0.178, 0.467)), # 2: 01069 + ((0.243, 0.440), (0.411, 0.469)), # 3: 1 31.000 + ((0.094, 0.471), (0.561, 0.500)), # 4: TAHU BAKSO Grg Ayam VC [biji + ((0.651, 0.470), (0.754, 0.492)), # 5: 68.000 + ((0.091, 0.506), (0.174, 0.529)), # 6: 07165 + ((0.243, 0.504), (0.411, 0.529)), # 7: 10 6.800 + ((0.087, 0.564), (0.430, 0.593)), # 8: QTY 11 ITEMS 2 + ((0.088, 0.599), (0.170, 0.622)), # 9: TOTAL + ((0.635, 0.597), (0.741, 0.619)), # 10: 99.000 + ((0.086, 0.629), (0.172, 0.655)), # 11: KARTU + ((0.637, 0.629), (0.742, 0.652)), # 12: 99.000 + ] + assert sort_reading_order(geoms) == list(range(13)) + # Right-to-left receipts: the amounts stand on the left of the items + mirrored = [((1 - x1, y0), (1 - x0, y1)) for (x0, y0), (x1, y1) in geoms] + assert sort_reading_order(mirrored, direction="rtl") == list(range(13)) + # Text columns with the same gutter are still read column after column: their lines are long + geoms = [((0.05, 0.10 + 0.03 * i), (0.40, 0.12 + 0.03 * i)) for i in range(4)] + geoms += [((0.55, 0.10 + 0.03 * i), (0.90, 0.12 + 0.03 * i)) for i in range(4)] + assert sort_reading_order(geoms) == list(range(8)) + # So is a column of short lines on a page with many lines (e.g. a list beside a text column) + geoms = [((0.05, 0.02 + 0.012 * i), (0.40, 0.03 + 0.012 * i)) for i in range(40)] + geoms += [((0.55, 0.02 + 0.012 * i), (0.60, 0.03 + 0.012 * i)) for i in range(40)] + assert sort_reading_order(geoms) == list(range(80)) + + def test_sort_reading_order_follows_columns_when_a_gutter_exists(): geoms = [ _box(0.05, 0.02, 0.95, 0.06), # 0 full-width header