diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py index c78409aadc..c626c41409 100644 --- a/doctr/io/exporters.py +++ b/doctr/io/exporters.py @@ -157,160 +157,203 @@ def page_reading_order( Returns: the ordered items, their layout label (None without layout) and the effective reading direction """ - from doctr.io.elements import Block, LayoutElement, Table - from doctr.models.reading_order import ( - ReadingOrderPredictor, - deskew_reading_geometries, - layout_label_role, - normalize_layout_label, - resolve_reading_segments, - ) - from doctr.models.reading_order.base import _FURNITURE_ROLES, _float_text_label, _to_boxes + from doctr.models.reading_order import ReadingOrderPredictor, deskew_reading_geometries, resolve_reading_segments signature = _reading_order_signature(page, direction) cached = getattr(page, "_reading_order_cache", None) if cached is not None and cached[0] == signature: - items, labels, resolved = cached[1] - return _select_items(list(items), list(labels), resolved, include_figures) - - texts = [word.value for block in page.blocks for line in block.lines for word in line.words] - language = page.language.get("value") if isinstance(page.language, dict) else None - direction = ReadingOrderPredictor(direction=direction).resolve_direction(texts, language=language) - region_geoms = [region.geometry for region in page.layout] - region_labels = [region.type for region in page.layout] + cached_items, cached_labels, resolved = cached[1] + return _select_items(list(cached_items), list(cached_labels), resolved, include_figures) lines = [line for block in page.blocks for line in block.lines] - # Figures take part in the ordering as floats - figures = picture_regions(page) - elements: list[Any] = [*lines, *page.tables, *figures] - if len(elements) == 0: - _store_reading_order(page, signature, ([], [], direction)) - return [], [], direction - # De-skew once so labeling, ordering and region grouping share the same upright frame; the page angle is - # estimated from the word polygons, which carry the detection model's true orientation - elt_geoms, region_geoms = deskew_reading_geometries( - [elt.geometry for elt in elements], - region_geoms, - page_shape=page.dimensions, - angle_geoms=[word.geometry for line in lines for word in line.words], + words = [word for line in lines for word in line.words] + language = page.language.get("value") if isinstance(page.language, dict) else None + direction = ReadingOrderPredictor(direction=direction).resolve_direction( + [word.value for word in words], language=language ) - elt_labels: list[str | None] = [None] * len(elements) - elt_regions = _covering_region_indices(elt_geoms, region_geoms, region_labels=region_labels) - if len(region_geoms) > 0: - elt_labels = [str(region_labels[reg]) if reg >= 0 else None for reg in elt_regions] + # Lines, tables and figures (as floats) are ordered together + elements: list[Any] = [*lines, *page.tables, *picture_regions(page)] + items: list[Any] = [] + labels: list[str | None] = [] + if len(elements) > 0: + # De-skew once so labeling, ordering and region grouping share the same upright frame; the page angle is + # estimated from the word polygons, which carry the detection model's true orientation + elt_geoms, region_geoms = deskew_reading_geometries( + [elt.geometry for elt in elements], + [region.geometry for region in page.layout], + page_shape=page.dimensions, + angle_geoms=[word.geometry for word in words], + ) + region_labels = [region.type for region in page.layout] + elt_regions = _covering_region_indices(elt_geoms, region_geoms, region_labels=region_labels) + elt_labels = _element_labels(page, elements, elt_geoms, elt_regions, region_geoms) + segments = resolve_reading_segments( + elt_geoms, + direction=direction, + labels=[_order_label(elt, label) for elt, label in zip(elements, elt_labels)], + page_shape=page.dimensions, + column_voters=_column_voters(elements, elt_geoms, elt_regions, region_geoms, region_labels), + region_groups=_merge_overlapping_regions(elt_regions, region_geoms, region_labels), + ) + items, labels = _segments_to_items(page, elements, segments, elt_labels, elt_regions) + _store_reading_order(page, signature, (items, labels, direction)) + return _select_items(list(items), list(labels), direction, include_figures) + + +def _element_labels( + page: "Page", elements: list[Any], elt_geoms: list[Any], elt_regions: list[int], region_geoms: list[Any] +) -> list[str | None]: + """Return the layout label of each element: the label of its layout region (None for none), 'Table' for a table + + A figure takes its own type, unless it lies in page furniture (e.g. a logo in the page header): it is then read with + the furniture, along with its text. + """ + from doctr.io.elements import LayoutElement, Table + from doctr.models.reading_order import layout_label_role + from doctr.models.reading_order.base import _FURNITURE_ROLES - # A figure lying in page furniture (e.g. a logo in the page header) is read with it, along with its text + region_labels = [str(region.type) for region in page.layout] + labels: list[str | None] = [region_labels[reg] if reg >= 0 else None for reg in elt_regions] furniture = [reg for reg, label in enumerate(region_labels) if layout_label_role(label) in _FURNITURE_ROLES] - fig_idcs = [idx for idx, elt in enumerate(elements) if isinstance(elt, LayoutElement)] - fig_furniture = _covering_region_indices( - [elt_geoms[idx] for idx in fig_idcs], [region_geoms[reg] for reg in furniture] + figures = [idx for idx, elt in enumerate(elements) if isinstance(elt, LayoutElement)] + in_furniture = _covering_region_indices( + [elt_geoms[idx] for idx in figures], [region_geoms[reg] for reg in furniture] ) region_index = {id(region): reg for reg, region in enumerate(page.layout)} - furniture_labels: dict[int, str] = {} # figure region index -> furniture label - for idx, reg in zip(fig_idcs, fig_furniture): - elt_labels[idx] = str(region_labels[furniture[reg]]) if reg >= 0 else elements[idx].type + furniture_of_figure: dict[int, str] = {} # figure region index -> label of the furniture holding it + for idx, reg in zip(figures, in_furniture): + labels[idx] = region_labels[furniture[reg]] if reg >= 0 else elements[idx].type if reg >= 0: - furniture_labels[region_index[id(elements[idx])]] = str(region_labels[furniture[reg]]) + furniture_of_figure[region_index[id(elements[idx])]] = region_labels[furniture[reg]] for idx, elt in enumerate(elements): if isinstance(elt, Table): - 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, 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 + labels[idx] = "Table" + elif not isinstance(elt, LayoutElement) and elt_regions[idx] in furniture_of_figure: + labels[idx] = furniture_of_figure[elt_regions[idx]] + return labels + + +def _order_label(element: Any, label: str | None) -> str | None: + """Return the label used to order an element: a text line in a table or figure region keeps the label of its + region, but is grouped into paragraphs like regular text""" + from doctr.io.elements import LayoutElement, Table + from doctr.models.reading_order import layout_label_role + from doctr.models.reading_order.base import _float_text_label + + if not isinstance(element, (Table, LayoutElement)) and layout_label_role(label) == "float": + return _float_text_label(str(label)) + return label + + +def _column_voters( + elements: list[Any], elt_geoms: list[Any], elt_regions: list[int], region_geoms: list[Any], region_labels: list[str] +) -> list[bool]: + """Return the elements voting in the multi-column detection + + Page-wide figures and tables do not vote, 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. + """ + from doctr.io.elements import LayoutElement, Table + from doctr.models.reading_order import normalize_layout_label + from doctr.models.reading_order.base import _to_boxes + 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)) + region_boxes = _to_boxes(region_geoms) 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) - and (isinstance(elt, (LayoutElement, Table)) or reg not in wide_tables) + return [ + not (box[2] - box[0] > 0.5 * span) if isinstance(elt, (LayoutElement, Table)) else 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 = [ - _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=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: - return isinstance(elements[idx], (Table, LayoutElement)) - def _split_floats(segment: list[int]) -> list[list[int]]: - """Split the tables and figures out of a segment (a figure in page furniture can share one)""" - runs: list[list[int]] = [] - for idx in segment: - if runs and not _is_float_item(idx) and not _is_float_item(runs[-1][-1]): - runs[-1].append(idx) - else: - runs.append([idx]) - return runs - - segments = [run for segment in segments for run in _split_floats(segment)] - - items = [] - labels = [] - line_owner = {id(line): idx for idx, block in enumerate(page.blocks) for line in block.lines} - pending_artefacts = {idx: list(block.artefacts) for idx, block in enumerate(page.blocks) if block.artefacts} - - def _claim_artefacts(block_lines: list[Any]) -> list[Any]: - claimed: list[Any] = [] - for line in block_lines: - owner = line_owner.get(id(line)) - if owner is not None and owner in pending_artefacts: - claimed.extend(pending_artefacts.pop(owner)) - return claimed - - # Region index covering each element, used to group the lines of a wrapped list item under a single bullet - region_idx = elt_regions - open_list_region: int | None = None # region of the list bullet currently being built (None outside a list) - for segment in segments: - first = elements[segment[0]] - seg_label = elt_labels[segment[0]] +def _segments_to_items( + page: "Page", elements: list[Any], segments: list[list[int]], elt_labels: list[str | None], elt_regions: list[int] +) -> tuple[list[Any], list[str | None]]: + """Turn reading-ordered segments of elements into page items: one block per paragraph, one block per list item + (a list item wrapped over several lines is one block), tables and figures as they are + + The artefacts of a block follow its first line; those of blocks without lines go to the last block. + """ + from doctr.io.elements import LayoutElement, Table + from doctr.models.reading_order import normalize_layout_label + + artefacts = _ArtefactClaims(page.blocks) + items: list[Any] = [] + labels: list[str | None] = [] + list_region: int | None = None # region of the list item being built + for run in (run for segment in segments for run in _float_runs(elements, segment)): + first, label = elements[run[0]], elt_labels[run[0]] if isinstance(first, (Table, LayoutElement)): items.append(first) - labels.append("Table" if isinstance(first, Table) else seg_label) - open_list_region = None - continue - if normalize_layout_label(seg_label) in _LIST_LABELS: - # One bullet per list-item region: consecutive lines sharing the same region are one bullet, so a - # list item wrapped over several visual lines renders as a single bullet point. - for idx in segment: - region = region_idx[idx] - if open_list_region is not None and region == open_list_region and region != -1: - merged = [*items[-1].lines, elements[idx]] - items[-1] = Block(lines=merged, artefacts=[*items[-1].artefacts, *_claim_artefacts([merged[-1]])]) + labels.append("Table" if isinstance(first, Table) else label) + list_region = None + elif normalize_layout_label(label) in _LIST_LABELS: + # One bullet per list-item region: consecutive lines of the same region are a single bullet + for idx in run: + region = elt_regions[idx] + if region != -1 and region == list_region: + items[-1] = artefacts.block([*items[-1].lines, elements[idx]], items[-1].artefacts) else: - items.append(Block(lines=[elements[idx]], artefacts=_claim_artefacts([elements[idx]]))) - labels.append(seg_label) - open_list_region = region + items.append(artefacts.block([elements[idx]])) + labels.append(label) + list_region = region else: - block_lines = [elements[idx] for idx in segment] - items.append(Block(lines=block_lines, artefacts=_claim_artefacts(block_lines))) - labels.append(seg_label) - open_list_region = None - # Artefacts of blocks without any line stay attached to the page - leftover = [artefact for artefacts in pending_artefacts.values() for artefact in artefacts] - if leftover: + items.append(artefacts.block([elements[idx] for idx in run])) + labels.append(label) + list_region = None + artefacts.attach_leftovers(items) + return items, labels + + +class _ArtefactClaims: + """Hand the artefacts of the page blocks over to the blocks built in reading order + + Args: + blocks: the blocks of the page + """ + + def __init__(self, blocks: list[Any]) -> None: + self.owner = {id(line): idx for idx, block in enumerate(blocks) for line in block.lines} + self.pending = {idx: list(block.artefacts) for idx, block in enumerate(blocks) if block.artefacts} + + def block(self, lines: list[Any], artefacts: list[Any] | None = None) -> Any: + """Build a block, which claims the artefacts of the blocks its lines come from (unless already claimed)""" + from doctr.io.elements import Block + + claimed = list(artefacts or []) + for line in lines: + owner = self.owner.get(id(line)) + if owner is not None and owner in self.pending: + claimed.extend(self.pending.pop(owner)) + return Block(lines=lines, artefacts=claimed) + + def attach_leftovers(self, items: list[Any]) -> None: + """Attach the artefacts of blocks without any line to the last block, so they stay on the page""" + from doctr.io.elements import Block + + leftover = [artefact for artefacts in self.pending.values() for artefact in artefacts] last_block = next((item for item in reversed(items) if isinstance(item, Block)), None) - if last_block is not None: + if leftover and last_block is not None: last_block.artefacts = [*last_block.artefacts, *leftover] - _store_reading_order(page, signature, (items, labels, direction)) - return _select_items(list(items), list(labels), direction, include_figures) + + +def _float_runs(elements: list[Any], segment: list[int]) -> list[list[int]]: + """Split the tables and figures out of a segment (a figure in page furniture can share one)""" + from doctr.io.elements import LayoutElement, Table + + runs: list[list[int]] = [] + for idx in segment: + is_float = isinstance(elements[idx], (Table, LayoutElement)) + if runs and not is_float and not isinstance(elements[runs[-1][-1]], (Table, LayoutElement)): + runs[-1].append(idx) + else: + runs.append([idx]) + return runs def _select_items( @@ -392,6 +435,145 @@ class _FigurePlan(NamedTuple): hidden: set[int] # text left out of the export +class _PageFigures: + """The figures of a linearized page, in the upright frame of the reading order + + Args: + page: the page + items: the linearized page content, figures included + labels: the layout label of each item + """ + + def __init__(self, page: "Page", items: list[Any], labels: list[str | None]) -> None: + from doctr.io.elements import Block + from doctr.models.reading_order import deskew_reading_geometries + + self.items, self.labels = items, labels + # Same upright frame as the reading order, for the items and the lines of the blocks + block_lines = [(idx, line) for idx, item in enumerate(items) if isinstance(item, Block) for line in item.lines] + upright, upright_regions = deskew_reading_geometries( + [item.geometry for item in items] + [line.geometry for _, line in block_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], + ) + self.boxes = [_xyxy(geom) for geom in upright[: len(items)]] + self.line_boxes: dict[int, list[tuple[float, ...]]] = {} + for (owner, _), geom in zip(block_lines, upright[len(items) :]): + self.line_boxes.setdefault(owner, []).append(_xyxy(geom)) + self.region_of = _covering_region_indices( + upright[: len(items)], upright_regions, region_labels=[region.type for region in page.layout] + ) + + def in_furniture(self, fig: int) -> bool: + """Whether a figure is read with the page furniture""" + from doctr.models.reading_order import layout_label_role + from doctr.models.reading_order.base import _FURNITURE_ROLES + + return layout_label_role(self.labels[fig]) in _FURNITURE_ROLES + + def without_furniture(self, figs: list[int]) -> tuple[list[int], set[int]]: + """Leave out the figures read with the page furniture + + Returns: + the other figures, and the text of the figures left out + """ + furniture = [fig for fig in figs if self.in_furniture(fig)] + return [fig for fig in figs if fig not in furniture], {text for fig in furniture for text in self.text_of(fig)} + + def _inside(self, idx: int, fig: int) -> bool: + """Whether each line of a block lies in a figure, by at least half of its area""" + target = self.boxes[fig] + + def _coverage(box: tuple[float, ...]) -> float: + inter = max(min(box[2], target[2]) - max(box[0], target[0]), 0.0) * max( + min(box[3], target[3]) - max(box[1], target[1]), 0.0 + ) + return inter / max((box[2] - box[0]) * (box[3] - box[1]), 1e-9) + + return all(_coverage(box) >= 0.5 for box in self.line_boxes.get(idx, [self.boxes[idx]])) + + def _is_block(self, idx: int) -> bool: + from doctr.io.elements import Block + + return 0 <= idx < len(self.items) and isinstance(self.items[idx], Block) + + def _is_caption(self, idx: int) -> bool: + from doctr.models.reading_order import normalize_layout_label + + return self._is_block(idx) and normalize_layout_label(self.labels[idx]) == "caption" + + def text_of(self, fig: int) -> set[int]: + """The blocks inside a figure and labeled as part of it (or as furniture, for a figure in furniture)""" + from doctr.models.reading_order import layout_label_role + from doctr.models.reading_order.base import _FURNITURE_ROLES + + in_furniture = self.in_furniture(fig) + return { + idx + for idx in range(len(self.items)) + if self._is_block(idx) + and ( + is_picture_label(self.labels[idx]) + or (in_furniture and layout_label_role(self.labels[idx]) in _FURNITURE_ROLES) + ) + and self._inside(idx, fig) + } + + def _caption_run(self, start: int) -> tuple[int, ...]: + """The caption blocks of the caption region of a block (a caption region can span several blocks)""" + if self.region_of[start] == -1: + return (start,) + return tuple( + idx + for idx in range(len(self.items)) + if self._is_caption(idx) and self.region_of[idx] == self.region_of[start] + ) + + def _distance(self, run: tuple[int, ...], target: int) -> float: + from doctr.models.reading_order.base import _caption_distance + + run_box = tuple(func(self.boxes[idx][k] for idx in run) for k, func in enumerate((min, min, max, max))) + return _caption_distance(run_box, self.boxes[target]) + + def _table_without_caption(self, idx: int, run: tuple[int, ...]) -> bool: + """Whether a table stands next to a caption run, without its own caption on its other side""" + from doctr.io.elements import Table + + if not (0 <= idx < len(self.items) and isinstance(self.items[idx], Table)): + return False + return not self._is_caption(idx + 1 if idx > run[-1] else idx - 1) + + def captions(self, figures: list[int]) -> dict[int, tuple[int, ...]]: + """Pair figures and captions: each caption read right before a figure, or right after the text inside it, goes + to its closest float (a table without caption competes for it), and each figure keeps its closest caption + + Args: + figures: the figures to caption + + Returns: + the caption blocks of each captioned figure + """ + claims: dict[tuple[int, ...], list[int]] = {} + for fig in figures: + after = fig + 1 + while self._is_block(after) and not self._is_caption(after) and self._inside(after, fig): + after += 1 + for cap in (fig - 1, after): + if self._is_caption(cap): + claims.setdefault(self._caption_run(cap), []).append(fig) + captions: dict[int, tuple[int, ...]] = {} + for run, candidates in claims.items(): + tables = [idx for idx in (run[0] - 1, run[-1] + 1) if self._table_without_caption(idx, run)] + best = min([*candidates, *tables], key=lambda target: self._distance(run, target)) + if best in tables or ( + best in captions and self._distance(captions[best], best) <= self._distance(run, best) + ): + continue + captions[best] = run + return captions + + class _PageTextExporter: """Shared logic of the reading-order-aware text exporters @@ -489,136 +671,40 @@ def _plan_figures( Returns: the figure markup, the captions rendered with their figure and the text to leave out """ - from doctr.io.elements import Block, LayoutElement, Table - from doctr.models.reading_order import deskew_reading_geometries, layout_label_role, normalize_layout_label - from doctr.models.reading_order.base import _FURNITURE_ROLES, _caption_distance + from doctr.io.elements import LayoutElement figure_idcs = [idx for idx, item in enumerate(items) if isinstance(item, LayoutElement)] render = self.supports_figures and encoder.enabled if not figure_idcs or not (render or not include_furniture): return _FigurePlan({}, set(), set()) - # 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] - + [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[: 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`""" - inter = max(min(box[2], target[2]) - max(box[0], target[0]), 0.0) * max( - min(box[3], target[3]) - max(box[1], target[1]), 0.0 - ) - 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 - return { - idx - for idx, item in enumerate(items) - if isinstance(item, Block) - and ( - is_picture_label(labels[idx]) - or (in_furniture and layout_label_role(labels[idx]) in _FURNITURE_ROLES) - ) - and _inside(idx, fig_idx) - } - + figures = _PageFigures(page, items, labels) hidden: set[int] = set() if not include_furniture: - # Figures read with the page furniture are left out with it - furniture = {idx for idx in figure_idcs if layout_label_role(labels[idx]) in _FURNITURE_ROLES} - for idx in furniture: - hidden |= _figure_text(idx) - figure_idcs = [idx for idx in figure_idcs if idx not in furniture] + # Figures read with the page furniture are left out with it, along with their text + figure_idcs, hidden = figures.without_furniture(figure_idcs) if not render or not figure_idcs: return _FigurePlan({}, set(), hidden) sources = {idx: encoder.source(page, items[idx], rank) for rank, idx in enumerate(figure_idcs, start=1)} - # Only a figure exported with its pixels makes its text redundant - for idx in figure_idcs: - if sources[idx] is not None: - hidden |= _figure_text(idx) - 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 _inside(idx, fig_idx) - - def _caption_run(start: int) -> tuple[int, ...]: - # A caption region can span several blocks - if region_of[start] == -1: - return (start,) - return tuple(idx for idx in range(len(items)) if _is_caption(idx) and region_of[idx] == region_of[start]) - - def _run_box(run: tuple[int, ...]) -> tuple[float, ...]: - return tuple(func(boxes[idx][k] for idx in run) for k, func in enumerate((min, min, max, max))) - - # Candidate captions: right before the figure, or right after the text inside it - claims: dict[tuple[int, ...], list[int]] = {} - for idx in figure_idcs: - if sources[idx] is None: - continue - after = idx + 1 - while after < len(items) and _is_inner_text(after, idx): - after += 1 - for cap in (idx - 1, after): - if 0 <= cap < len(items) and _is_caption(cap): - claims.setdefault(_caption_run(cap), []).append(idx) - - def _table_without_caption(idx: int, run: tuple[int, ...]) -> bool: - # A table with its own caption on its other side does not compete for this one - if not (0 <= idx < len(items) and isinstance(items[idx], Table)): - return False - far = idx + 1 if idx > run[-1] else idx - 1 - return not (0 <= far < len(items) and _is_caption(far)) - - # Each caption goes to its closest float, and each figure keeps its closest caption - captions: dict[int, tuple[int, ...]] = {} - for run, candidates in claims.items(): - cap_box = _run_box(run) - tables = [idx for idx in (run[0] - 1, run[-1] + 1) if _table_without_caption(idx, run)] - best = min([*candidates, *tables], key=lambda elt: _caption_distance(cap_box, boxes[elt])) - if best in tables: - continue - dist = _caption_distance(cap_box, boxes[best]) - if best in captions and _caption_distance(_run_box(captions[best]), boxes[best]) <= dist: - continue - captions[best] = run - + # Only a figure exported with its pixels makes its text redundant, and takes a caption as alternative text + shown = [idx for idx in figure_idcs if sources[idx] is not None] + hidden |= {text for idx in shown for text in figures.text_of(idx)} + captions = figures.captions(shown) markup: dict[int, str] = {} consumed: set[int] = set() for idx in figure_idcs: - caption = None - if idx in captions: - lines = [ - line - for cap_idx in captions[idx] - for line in self._block_lines(items[cap_idx], direction, False, auto) - ] - caption = _join_caption_lines(lines) or None - if caption is not None: - consumed.update(captions[idx]) + caption = self._caption_text([items[cap] for cap in captions.get(idx, ())], direction, auto) + if caption is not None: + consumed.update(captions[idx]) markup[idx] = self.render_figure(sources[idx], caption, escape=escape) return _FigurePlan(markup, consumed, hidden) + def _caption_text(self, blocks: list[Any], direction: str, auto: bool) -> str | None: + """Join the lines of the caption blocks of a figure into one caption (None without any text)""" + lines = [line for block in blocks for line in self._block_lines(block, direction, False, auto)] + return _join_caption_lines(lines) or None + def export_page( self, page: "Page", diff --git a/doctr/io/figures.py b/doctr/io/figures.py index e30d0014e0..15aa91b382 100644 --- a/doctr/io/figures.py +++ b/doctr/io/figures.py @@ -69,16 +69,15 @@ def picture_regions(page: "Page") -> list["LayoutElement"]: Returns: the figure regions, in detection order """ - from doctr.models.reading_order.base import _to_boxes + from doctr.models.reading_order.base import _box_areas, _pairwise_overlaps, _to_boxes regions = [region for region in (getattr(page, "layout", None) or []) if is_picture_region(region)] if len(regions) < 2: return regions boxes = _to_boxes([region.geometry for region in regions]) - inter_w = np.minimum(boxes[:, None, 2], boxes[None, :, 2]) - np.maximum(boxes[:, None, 0], boxes[None, :, 0]) - inter_h = np.minimum(boxes[:, None, 3], boxes[None, :, 3]) - np.maximum(boxes[:, None, 1], boxes[None, :, 1]) - areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) - inside = np.clip(inter_w, 0, None) * np.clip(inter_h, 0, None) >= 0.5 * areas[:, None] + inter_w, inter_h = _pairwise_overlaps(boxes, boxes) + areas = _box_areas(boxes) + inside = inter_w * inter_h >= 0.5 * areas[:, None] # A region is nested when half of it lies in a larger region (or in an earlier one of the same area) rank = np.empty(len(regions), dtype=int) rank[np.argsort(-areas, kind="stable")] = np.arange(len(regions)) diff --git a/doctr/models/builder.py b/doctr/models/builder.py index ec9c31da0f..779161b92e 100644 --- a/doctr/models/builder.py +++ b/doctr/models/builder.py @@ -131,52 +131,31 @@ def _resolve_sub_lines( A list of (sub-)lines computed from the original line (words) """ break_dist = self.paragraph_break if break_dist is None else break_dist - lines = [] - # Sort words horizontally - word_idcs = [word_idcs[idx] for idx in boxes[word_idcs, 0].argsort().tolist()] - - # Eventually split line horizontally - if len(word_idcs) < 2: - lines.append(word_idcs) - else: - sub_line = [word_idcs[0]] - 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] - for i in word_idcs[1:]: - horiz_break = True - - # Compute distance between the sub-line and the box - dist = boxes[i, 0] - right - # 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) - box_region = -1 if groups is None else int(groups[i]) - if ( + groups, coverage = word_regions if word_regions is not None else (None, None) + lines: list[list[int]] = [] + right, region, prev = 0.0, -1, -1 # right edge, region and last word of the current sub-line + for idx in [word_idcs[i] for i in boxes[word_idcs, 0].argsort().tolist()]: + box_region = -1 if groups is None else int(groups[idx]) + # A new sub-line after a gap of at least `break_dist` from the sub-line (a box can be nested in an earlier + # one: its max right edge is used), or between two distinct regions (e.g. columns with a narrow gutter) + if ( + not lines + or not boxes[idx, 0] - right < break_dist + or ( coverage is not None and box_region >= 0 and region >= 0 and box_region != region - and coverage[i, region] < 0.5 + and coverage[idx, 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) - + ) + ): + lines.append([]) + right, region = boxes[idx, 2], -1 + lines[-1].append(idx) + right = max(right, boxes[idx, 2]) + region = box_region if region < 0 else region + prev = idx return lines @staticmethod @@ -198,7 +177,9 @@ def _word_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 ( + _box_areas, _covering_regions, + _pairwise_overlaps, _region_groups, _to_boxes, deskew_reading_geometries, @@ -214,19 +195,15 @@ 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)) - 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) + inter_w, inter_h = _pairwise_overlaps(words, regions) + x_cov = inter_w / np.clip(words[:, 2:3] - words[:, 0:1], 1e-9, None) + y_cov = inter_h / np.clip(words[:, 3:4] - words[:, 1:2], 1e-9, None) uncovered = np.flatnonzero(covering < 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]) + if ( + uncovered.size > 0 + ): # e.g. a word sticking out of a region slightly too short: the smallest region holding it holding = (x_cov[uncovered] >= 0.9) & (y_cov[uncovered] >= 0.3) - smallest = np.where(holding, areas[None, :], np.inf).argmin(axis=1) + smallest = np.where(holding, _box_areas(regions)[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) @@ -582,6 +559,88 @@ def _localize_logic(cells: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], norm.append(nc) return norm, max_row + 1, max_col + 1 + @classmethod + def _table_grids( + cls, page_table: dict[str, Any] | list[dict[str, Any]] | None + ) -> list[tuple[list[dict[str, Any]], int, int]]: + """Normalize the table structure prediction(s) of a page to non-empty grids with local logical coordinates + + Args: + page_table: a single grid, a list of grids or None (cf. `_build_tables`) + + Returns: + the cells, number of rows and number of columns of each non-empty grid + """ + raw_tables = [] if page_table is None else [page_table] if isinstance(page_table, dict) else list(page_table) + return [cls._localize_logic(raw["cells"]) for raw in raw_tables if raw and len(raw.get("cells", [])) > 0] + + @classmethod + def _assign_cell_words(cls, centers: np.ndarray, polys: np.ndarray, consumed: np.ndarray) -> list[list[int]]: + """Assign each word not yet consumed to at most one cell of a table, and flag it as consumed + + A word goes to the first cell containing its center. A word landing just outside every cell but inside the + table goes to the nearest cell, within half a median cell diagonal, so table text is not dropped into the body. + + Args: + centers: the (N, 2) word centers + polys: the (C, 4, 2) cell polygons, vertices ordered TL, TR, BR, BL + consumed: the (N,) mask of the words already assigned to a table, updated in place + + Returns: + the word indices of each cell + """ + cell_words: list[list[int]] = [[] for _ in polys] + free = np.flatnonzero(~consumed) + if free.size == 0: + return cell_words + inside = cls._points_in_polygons(centers[free], polys) # (F, C) + assigned = inside.any(axis=1) + for w_idx, c_idx in zip(free[assigned], inside[assigned].argmax(axis=1)): + cell_words[c_idx].append(int(w_idx)) + consumed[w_idx] = True + leftover = free[~assigned] + if leftover.size > 0: + (tx0, ty0), (tx1, ty1) = polys.reshape(-1, 2).min(axis=0), polys.reshape(-1, 2).max(axis=0) + cell_centers = polys.mean(axis=1) # (C, 2) + max_dist = 0.5 * float(np.median(np.linalg.norm(polys[:, 2] - polys[:, 0], axis=1))) # TL->BR diagonal + in_table = ( + (centers[leftover, 0] >= tx0) + & (centers[leftover, 0] <= tx1) + & (centers[leftover, 1] >= ty0) + & (centers[leftover, 1] <= ty1) + ) + for w_idx in leftover[in_table]: + dists = np.linalg.norm(cell_centers - centers[w_idx], axis=1) + nearest = int(dists.argmin()) + if dists[nearest] <= max_dist: + cell_words[nearest].append(int(w_idx)) + consumed[w_idx] = True + return cell_words + + @staticmethod + def _deskew_centers(centers: np.ndarray, polys: np.ndarray) -> np.ndarray: + """Rotate word centers by the angle of a table, to order the words along its own axes + + Args: + centers: the (N, 2) word centers + polys: the (C, 4, 2) cell polygons, vertices ordered TL, TR, BR, BL + + Returns: + the (N, 2) de-skewed centers + """ + top_edges = polys[:, 1] - polys[:, 0] # TR - TL + angle = float(np.median(np.arctan2(top_edges[:, 1], top_edges[:, 0]))) + cos_a, sin_a = np.cos(-angle), np.sin(-angle) + pivot = centers.mean(axis=0) + shifted = centers - pivot + return np.stack( + [ + pivot[0] + shifted[:, 0] * cos_a - shifted[:, 1] * sin_a, + pivot[1] + shifted[:, 0] * sin_a + shifted[:, 1] * cos_a, + ], + axis=1, + ) + def _build_tables( self, boxes: np.ndarray, @@ -601,117 +660,37 @@ def _build_tables( a tuple with the list of `Table` objects of the page (one per provided table) and a boolean mask of shape (N,) that is True for words assigned to a table (to be removed from `blocks`) """ - num_words = boxes.shape[0] - consumed = np.zeros(num_words, dtype=bool) - if page_table is None: + consumed = np.zeros(boxes.shape[0], dtype=bool) + grids = self._table_grids(page_table) + if len(grids) == 0: return [], consumed - - # Normalize the prediction(s) to a list of per-table grids with local 0-based logical coordinates - raw_tables = [page_table] if isinstance(page_table, dict) else list(page_table) - table_dicts: list[dict[str, Any]] = [] - for raw in raw_tables: - if not raw or len(raw.get("cells", [])) == 0: - continue - cells, n_rows, n_cols = self._localize_logic(raw["cells"]) - table_dicts.append({"cells": cells, "num_rows": n_rows, "num_cols": n_cols}) - if len(table_dicts) == 0: - return [], consumed - - centers = self._word_centers(boxes) if num_words > 0 else np.empty((0, 2)) - # Geometry format follows the page's word geometry: straight 2-point boxes when the word boxes are - # (N, 4), 4-point polygons when they are (N, 4, 2). + # The geometries follow the page's word geometry: straight 2-point boxes for (N, 4) word boxes, 4-point polygons + # for (N, 4, 2) ones. Word heights are rotation-invariant (left edge length for polygons). straight = boxes.ndim != 3 - # Rotation-invariant word heights (left edge length for polygons), used to cluster cell words into rows - if num_words == 0: - word_heights = np.empty(0) - elif straight: - word_heights = boxes[:, 3] - boxes[:, 1] + centers = self._word_centers(boxes) if boxes.shape[0] > 0 else np.empty((0, 2)) + if boxes.shape[0] == 0: + heights = np.empty(0) else: - word_heights = np.linalg.norm(boxes[:, 3] - boxes[:, 0], axis=1) - - tables_out: list[Table] = [] - for table_dict in table_dicts: - cells = table_dict["cells"] - if len(cells) == 0: - continue - cell_polys = [self._as_cell_polygon(cell["geometry"]) for cell in cells] - polys_arr = np.stack(cell_polys) # (C, 4, 2), vertices ordered TL, TR, BR, BL - - # Order along the table's own axes instead by de-skewing the centers with the table angle - order_centers = centers - if not straight and centers.shape[0] > 0: - top_edges = polys_arr[:, 1] - polys_arr[:, 0] # TR - TL - angle = float(np.median(np.arctan2(top_edges[:, 1], top_edges[:, 0]))) - cos_a, sin_a = np.cos(-angle), np.sin(-angle) - pivot = centers.mean(axis=0) - shifted = centers - pivot - order_centers = np.stack( - [ - pivot[0] + shifted[:, 0] * cos_a - shifted[:, 1] * sin_a, - pivot[1] + shifted[:, 0] * sin_a + shifted[:, 1] * cos_a, - ], - axis=1, - ) - - # Assign each (still unassigned) word to at most one cell of this table - cell_word_idcs: list[list[int]] = [[] for _ in cells] - free_idcs = np.flatnonzero(~consumed) - if free_idcs.size > 0: - # The first cell (in cell order) whose polygon contains the word center - inside = self._points_in_polygons(centers[free_idcs], polys_arr) # (F, C) - assigned = inside.any(axis=1) - first_cell = np.where(assigned, inside.argmax(axis=1), -1) - for w_idx, c_idx in zip(free_idcs, first_cell): - if c_idx >= 0: - cell_word_idcs[c_idx].append(int(w_idx)) - consumed[w_idx] = True - - # Words that landed just outside every cell but still inside the table region are attached - # to the nearest cell, so table text is not dropped into the body. - # The capture radius is bounded by the cell size to avoid pulling in body words. - leftover = free_idcs[~assigned] - if leftover.size > 0: - tx0, ty0 = polys_arr[..., 0].min(), polys_arr[..., 1].min() - tx1, ty1 = polys_arr[..., 0].max(), polys_arr[..., 1].max() - cell_centers = polys_arr.mean(axis=1) # (C, 2) - cell_sizes = np.linalg.norm(polys_arr[:, 2] - polys_arr[:, 0], axis=1) # TL->BR diagonal - max_dist = 0.5 * float(np.median(cell_sizes)) - in_region = ( - (centers[leftover, 0] >= tx0) - & (centers[leftover, 0] <= tx1) - & (centers[leftover, 1] >= ty0) - & (centers[leftover, 1] <= ty1) - ) - for w_idx in leftover[in_region]: - dists = np.linalg.norm(cell_centers - centers[w_idx], axis=1) - nearest = int(dists.argmin()) - if dists[nearest] <= max_dist: - cell_word_idcs[nearest].append(int(w_idx)) - consumed[w_idx] = True - - # Build the cells - table_cells: list[TableCell] = [] - for cell, poly, w_idcs in zip(cells, cell_polys, cell_word_idcs): + heights = boxes[:, 3] - boxes[:, 1] if straight else np.linalg.norm(boxes[:, 3] - boxes[:, 0], axis=1) + + tables: list[Table] = [] + for cells, num_rows, num_cols in grids: + polys = np.stack([self._as_cell_polygon(cell["geometry"]) for cell in cells]) # (C, 4, 2) + order_centers = centers if straight or centers.shape[0] == 0 else self._deskew_centers(centers, polys) + table_cells = [] + for cell, poly, w_idcs in zip(cells, polys, self._assign_cell_words(centers, polys, consumed)): if len(w_idcs) > 0: # Reading order inside the cell: rows top to bottom, words left to right (table axes) - ordered = self._order_cell_words(w_idcs, order_centers, word_heights) + ordered = self._order_cell_words(w_idcs, order_centers, heights) value = " ".join(word_preds[i][0] for i in ordered) confidence = float(np.mean([word_preds[i][1] for i in ordered])) else: value, confidence = "", float(cell["score"]) - if straight: - xs, ys = poly[:, 0], poly[:, 1] - geometry: Any = ( - (float(xs.min()), float(ys.min())), - (float(xs.max()), float(ys.max())), - ) - else: - geometry = tuple(tuple(float(c) for c in pt) for pt in poly.tolist()) table_cells.append( TableCell( value=value, confidence=confidence, - geometry=geometry, + geometry=self._export_geometry(poly, straight), row_start=int(cell["row_start"]), row_end=int(cell["row_end"]), col_start=int(cell["col_start"]), @@ -719,27 +698,29 @@ def _build_tables( objectness_score=float(cell["score"]), ) ) - # Enclosing geometry of the whole table - all_pts = np.concatenate(cell_polys, axis=0) - xmin, ymin = float(all_pts[:, 0].min()), float(all_pts[:, 1].min()) - xmax, ymax = float(all_pts[:, 0].max()), float(all_pts[:, 1].max()) - table_geometry: Any = ( - ((xmin, ymin), (xmax, ymax)) if straight else ((xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)) - ) - table_confidence = float(np.mean([float(cell["score"]) for cell in cells])) - - tables_out.append( + (xmin, ymin), (xmax, ymax) = polys.reshape(-1, 2).min(axis=0), polys.reshape(-1, 2).max(axis=0) + corners = np.asarray([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]) + tables.append( Table( cells=table_cells, - num_rows=int(table_dict["num_rows"]), - num_cols=int(table_dict["num_cols"]), - geometry=table_geometry, - confidence=table_confidence, + num_rows=num_rows, + num_cols=num_cols, + geometry=self._export_geometry(corners, straight), + confidence=float(np.mean([float(cell["score"]) for cell in cells])), ) ) + return tables, consumed - return tables_out, consumed + @staticmethod + def _export_geometry(poly: np.ndarray, straight: bool) -> Any: + """Export a (4, 2) polygon as a straight 2-point box or as a 4-point polygon, with Python floats""" + if straight: + return ( + (float(poly[:, 0].min()), float(poly[:, 1].min())), + (float(poly[:, 0].max()), float(poly[:, 1].max())), + ) + return tuple(tuple(float(c) for c in pt) for pt in poly.tolist()) def _build_blocks( self, @@ -771,49 +752,32 @@ def _build_blocks( if boxes.shape[0] == 0: return [] - # Decide whether we try to form lines - _boxes = boxes - if self.resolve_lines: + straight_or_rotated = boxes if boxes.ndim == 3 else boxes[:, :4] + if not self.resolve_lines: + # One line for all boxes, one block for the line + lines = [self._sort_boxes(straight_or_rotated, shape)[0].tolist()] + else: word_regions = ( - self._word_regions(_boxes, [elt.geometry for elt in layout], [elt.type for elt in layout], shape) + 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) - else: - _blocks = [lines] - else: - # Sort bounding boxes, one line for all boxes, one block for the line - lines = [self._sort_boxes(_boxes if _boxes.ndim == 3 else _boxes[:, :4], shape)[0]] # type: ignore[list-item] - _blocks = [lines] - - blocks = [ - Block([ - Line([ - Word( - *word_preds[idx], - tuple(tuple(pt) for pt in boxes[idx].tolist()), # type: ignore[arg-type] - float(objectness_scores[idx]), - crop_orientations[idx], - ) - if boxes.ndim == 3 - else Word( - *word_preds[idx], - ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3])), - float(objectness_scores[idx]), - crop_orientations[idx], - ) - for idx in line - ]) - for line in lines - ]) - for lines in _blocks - ] + lines = self._resolve_lines(straight_or_rotated, shape, word_regions) + blocks = ( + self._resolve_blocks(straight_or_rotated, lines) + if self.resolve_lines and self.resolve_blocks and len(lines) > 1 + else [lines] + ) - return blocks + def _word(idx: int) -> Word: + geometry = ( + tuple(tuple(pt) for pt in boxes[idx].tolist()) + if boxes.ndim == 3 + else ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3])) + ) + return Word(*word_preds[idx], geometry, float(objectness_scores[idx]), crop_orientations[idx]) # type: ignore[arg-type] + + return [Block([Line([_word(idx) for idx in line]) for line in block]) for block in blocks] @staticmethod def _words_to_boxes(words: list[Word]) -> np.ndarray: diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index cb10422ce6..ff1265ece9 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -35,15 +35,16 @@ # Unicode bidirectional categories of strong right-to-left characters (Hebrew-like & Arabic-like) _RTL_BIDI_CATEGORIES = {"R", "AL"} -# Normalized layout labels (DocLayNet-style) grouped by their role in the reading order -_HEADER_LABELS = {"page_header", "header", "running_header"} -_FOOTER_LABELS = {"page_footer", "footer", "running_footer"} -_FOOTNOTE_LABELS = {"footnote"} -_CAPTION_LABELS = {"caption"} -_FLOAT_LABELS = {"table", "picture", "figure", "image", "chart", "graphic"} - +# Role in the reading order of the normalized layout labels (DocLayNet-style); any other label is body text +_LABEL_ROLES = { + **dict.fromkeys(("page_header", "header", "running_header"), "header"), + **dict.fromkeys(("page_footer", "footer", "running_footer"), "footer"), + "footnote": "footnote", + "caption": "caption", + **dict.fromkeys(("table", "picture", "figure", "image", "chart", "graphic"), "float"), +} +_FLOAT_LABELS = {label for label, role in _LABEL_ROLES.items() if role == "float"} _FURNITURE_ROLES = ("header", "footer", "footnote") - # 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} @@ -79,18 +80,7 @@ def layout_label_role(label: str | None) -> str: Returns: one of 'header', 'footer', 'footnote', 'caption', 'float' or 'body' """ - norm = normalize_layout_label(label) - if norm in _HEADER_LABELS: - return "header" - if norm in _FOOTER_LABELS: - return "footer" - if norm in _FOOTNOTE_LABELS: - return "footnote" - if norm in _CAPTION_LABELS: - return "caption" - if norm in _FLOAT_LABELS: - return "float" - return "body" + return _LABEL_ROLES.get(normalize_layout_label(label), "body") def _float_text_label(label: str) -> str: @@ -107,7 +97,7 @@ def detect_text_direction(texts: Iterable[str], language: str | None = None) -> detection language-independent and robust to mixed content. When no strong character is found, the optional ISO 639 language code is used as a fallback hint. - >>> from doctr.utils.reading_order import detect_text_direction + >>> from doctr.models.reading_order import detect_text_direction >>> detect_text_direction(["Hello", "world"]) 'ltr' >>> detect_text_direction(["مرحبا", "بالعالم"]) @@ -120,30 +110,31 @@ def detect_text_direction(texts: Iterable[str], language: str | None = None) -> Returns: either "ltr" or "rtl" """ - rtl_count = 0 - ltr_count = 0 + rtl_count = ltr_count = 0 for text in texts: if not isinstance(text, str): continue for char in text: bidi = unicodedata.bidirectional(char) - if bidi in _RTL_BIDI_CATEGORIES: - rtl_count += 1 - elif bidi == "L": - ltr_count += 1 + rtl_count += bidi in _RTL_BIDI_CATEGORIES + ltr_count += bidi == "L" if rtl_count == 0 and ltr_count == 0: - if isinstance(language, str) and language.split("-")[0].lower() in RTL_LANGUAGES: - return "rtl" - return "ltr" + rtl_hint = isinstance(language, str) and language.split("-")[0].lower() in RTL_LANGUAGES + return "rtl" if rtl_hint else "ltr" return "rtl" if rtl_count > ltr_count else "ltr" +# --- Geometry helpers ------------------------------------------------------------------------------------------------- + + def _geometry_to_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) return float(pts[:, 0].min()), float(pts[:, 1].min()), float(pts[:, 0].max()), float(pts[:, 1].max()) def _to_boxes(geoms: Sequence[Any] | np.ndarray) -> np.ndarray: + """Return the (N, 4) enclosing boxes of geometries in any docTR format""" if isinstance(geoms, np.ndarray) and geoms.ndim == 2 and geoms.shape[1] >= 4: return geoms[:, :4].astype(np.float64) if isinstance(geoms, np.ndarray) and geoms.ndim == 3: @@ -153,6 +144,25 @@ def _to_boxes(geoms: Sequence[Any] | np.ndarray) -> np.ndarray: return np.asarray([_geometry_to_xyxy(geom) for geom in geoms], dtype=np.float64) +def _box_areas(boxes: np.ndarray) -> np.ndarray: + return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + + +def _pairwise_overlaps(boxes: np.ndarray, others: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Width and height of the intersection of each (N, 4) box with each (M, 4) box (0 when they do not overlap)""" + inter_w = np.minimum(boxes[:, None, 2], others[None, :, 2]) - np.maximum(boxes[:, None, 0], others[None, :, 0]) + inter_h = np.minimum(boxes[:, None, 3], others[None, :, 3]) - np.maximum(boxes[:, None, 1], others[None, :, 1]) + return np.clip(inter_w, 0, None), np.clip(inter_h, 0, None) + + +def _is_inside(boxes: np.ndarray, idx: int, target: int, min_coverage: float = 0.5) -> bool: + """Check whether box `idx` lies inside box `target`, by at least `min_coverage` of its area""" + x0, y0, x1, y1 = boxes[idx] + tx0, ty0, tx1, ty1 = boxes[target] + inter = max(min(x1, tx1) - max(x0, tx0), 0.0) * max(min(y1, ty1) - max(y0, ty0), 0.0) + return inter >= min_coverage * max((x1 - x0) * (y1 - y0), 1e-9) + + def _to_canonical_ltr(boxes: np.ndarray, direction: str) -> np.ndarray: """Map boxes to a canonical left-to-right, top-to-bottom space for the given reading direction. @@ -162,13 +172,12 @@ def _to_canonical_ltr(boxes: np.ndarray, direction: str) -> np.ndarray: if direction == "ltr" or boxes.shape[0] == 0: return boxes x0, y0, x1, y1 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] + pivot = float(x0.min() + x1.max()) if direction == "rtl": # mirror horizontally - pivot = float(x0.min() + x1.max()) return np.stack([pivot - x1, y0, pivot - x0, y1], axis=1) if direction == "ttb-ltr": # vertical lines, columns ordered left to right: swap axes return np.stack([y0, x0, y1, x1], axis=1) if direction == "ttb-rtl": # vertical lines, columns ordered right to left: swap axes & mirror - pivot = float(x0.min() + x1.max()) return np.stack([y0, pivot - x1, y1, pivot - x0], axis=1) raise ValueError(f"invalid reading direction '{direction}', should be one of {SUPPORTED_DIRECTIONS[1:]}") @@ -195,6 +204,170 @@ def _strict_rank(primary: np.ndarray, secondary: np.ndarray) -> np.ndarray: return ranks +def deskew_reading_geometries( + geoms: Sequence[Any] | np.ndarray, + region_geoms: Sequence[Any] | np.ndarray | None = None, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, + min_angle: float = 1.0, +) -> tuple[list[Any], list[Any]]: + """De-skew rotated geometries so the reading order can be computed in an upright frame. + + Args: + geoms: geometries of the elements to order, in any docTR format (cf. `sort_reading_order`) + region_geoms: optional geometries of the layout regions, de-skewed together with the elements + page_shape: the page dimensions (height, width). Required for an exact angle on non-square pages with + relative coordinates, since relative coordinates distort angles by the aspect ratio. + angle_geoms: optional reading-oriented 4-point polygons used to estimate the page angle + min_angle: minimum estimated angle (in degrees) to trigger the de-skew. Ordering is affected as soon + as the drift along a line approaches the line height (about 1 degree for a page-wide line), and a + small rigid rotation cannot change the order of an upright page, hence the low default. Beyond 45 + degrees the corner identification is ambiguous (an upstream orientation correction is needed) and + nothing is done. + + Returns: + the (possibly de-skewed) element and region geometries + """ + region_geoms = list(region_geoms) if region_geoms is not None else [] + pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in geoms] + if len(pts) == 0 or any(p.shape[0] != 4 for p in pts): + return list(geoms), region_geoms # straight geometries: nothing to de-skew + height, width = page_shape if page_shape is not None else (1, 1) + scale = np.array([width, height], dtype=np.float64) + angle_pts = [ + np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in (angle_geoms if angle_geoms is not None else []) + ] + if len(angle_pts) > 0 and all(p.shape[0] == 4 for p in angle_pts): + # Detection polygons are already reading-oriented (cf. `estimate_page_angle`): keep their vertex order + angle = estimate_page_angle(np.stack(angle_pts) * scale) + else: + # Normalize the vertex order (TL, TR, BR, BL) so the estimation does not depend on the vertex + angle = estimate_page_angle(np.stack([order_points(p * scale) for p in pts])) + if not np.isfinite(angle) or abs(angle) < min_angle or abs(angle) >= 45: + return list(geoms), region_geoms + # Rigid rotation zeroing the estimated angle, in pixels; the center is irrelevant for ordering purposes + theta = np.deg2rad(angle) + rot = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) + center = (np.concatenate(pts, axis=0) * scale).mean(axis=0) + + def _rotate(points: np.ndarray) -> np.ndarray: + if points.shape[0] != 4: # a straight region: rotate its 4 corners, not only its diagonal + (x0, y0), (x1, y1) = points.min(axis=0), points.max(axis=0) + points = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) + return ((points * scale - center) @ rot.T + center) / scale + + regions = [_rotate(np.asarray(geom, dtype=np.float64).reshape(-1, 2)) for geom in region_geoms] + return [_rotate(p) for p in pts], regions + + +# --- Layout regions --------------------------------------------------------------------------------------------------- + + +def _covering_regions( + boxes: np.ndarray, + regions: np.ndarray, + min_coverage: float, + region_labels: Sequence[str | None] | None = None, +) -> list[int]: + """Return the index of the region covering each (N, 4) box by at least `min_coverage`, -1 if none + + Regions covering a box (nearly) as much as the best one tie: the smallest one wins, unless it is a float tied with + page furniture. + """ + inter_w, inter_h = _pairwise_overlaps(boxes, regions) + coverage = inter_w * inter_h / np.clip(_box_areas(boxes), 1e-9, None)[:, None] + region_areas = _box_areas(regions) + qualifies = coverage >= np.maximum(coverage.max(axis=1, keepdims=True) - 0.1, min_coverage) + # Regions nested in the best region of a box also tie when they cover most of it + reg_w, reg_h = _pairwise_overlaps(regions, regions) + nested = reg_w * reg_h >= 0.9 * np.clip(region_areas, 1e-9, None)[:, None] # nested[i, j]: region i lies inside j + qualifies |= (coverage >= max(min_coverage, 0.75)) & nested[:, coverage.argmax(axis=1)].T + ranked = np.where(qualifies, region_areas[None, :], np.inf) + best = ranked.argmin(axis=1) + result = [int(reg) if np.isfinite(ranked[idx, reg]) else -1 for idx, reg in enumerate(best)] + if region_labels is None or len(region_labels) != regions.shape[0]: + return result + + roles = [layout_label_role(label) for label in region_labels] + is_float = np.array([role == "float" for role in roles], dtype=bool) + is_furniture = np.array([role in _FURNITURE_ROLES for role in roles], dtype=bool) + if not (is_float.any() and is_furniture.any()): + return result + # The text of a float tied with page furniture (e.g. a logo in the page header) is read with the furniture + furniture_ranked = np.where(qualifies & is_furniture[None, :], region_areas[None, :], np.inf) + for idx, reg in enumerate(result): + if reg >= 0 and is_float[reg] and np.isfinite(furniture_ranked[idx].min()): + result[idx] = int(furniture_ranked[idx].argmin()) + return result + + +def _region_groups(regions: np.ndarray, labels: Sequence[str | None], min_overlap: float = 0.1) -> list[int]: + """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 = _box_areas(regions) + 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 + + # Pairwise with scalar min / max, as the comparisons of degenerate (NaN) boxes depend on the argument order + for i, j in ((i, j) for i in range(len(parent)) for j in range(i + 1, len(parent)) if norm[i] == norm[j]): + 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(len(parent))] + + +def assign_layout_labels( + geoms: Sequence[Any] | np.ndarray, + layout_geoms: Sequence[Any] | np.ndarray, + layout_labels: Sequence[str], + min_coverage: float = 0.5, + page_shape: tuple[int, int] | None = None, + angle_geoms: Sequence[Any] | np.ndarray | None = None, +) -> list[str | None]: + """Assign a layout label to each element based on its overlap with the detected layout regions. + + Each element receives the label of the region covering the largest share of its area, provided this share reaches + `min_coverage`; otherwise its label is None (treated as regular body content). Among nested regions, the smallest + one wins (e.g. 'Caption' over 'Picture'), except for a float nested in page furniture (e.g. a logo in the page + header): its text takes the furniture label. + + Args: + geoms: geometries of the elements to label, in any docTR format + layout_geoms: geometries of the layout regions, in any docTR format + layout_labels: labels of the layout regions (e.g. `[region.type for region in page.layout]`) + min_coverage: minimum share of an element's area a region must cover to assign its label + page_shape: the page dimensions (height, width), used to de-skew rotated pages exactly (cf. + `deskew_reading_geometries`) + angle_geoms: optional reading-oriented 4-point polygons (typically the page's word polygons) used to + estimate the page angle on rotated pages (cf. `deskew_reading_geometries`) + + Returns: + the label of each element (None when no region covers it enough) + """ + # De-skew elements and regions together so the coverage is measured on tight boxes in the same frame + geoms, layout_geoms = deskew_reading_geometries(geoms, layout_geoms, page_shape=page_shape, angle_geoms=angle_geoms) + boxes, regions = _to_boxes(geoms), _to_boxes(layout_geoms) + if len(layout_labels) != regions.shape[0]: + raise ValueError(f"Incompatible number of labels ({len(layout_labels)}) and regions ({regions.shape[0]})") + if boxes.shape[0] == 0 or regions.shape[0] == 0: + return [None] * boxes.shape[0] + covering = _covering_regions(boxes, regions, min_coverage, layout_labels) + return [str(layout_labels[reg]) if reg >= 0 else None for reg in covering] + + +# --- Topological sort ------------------------------------------------------------------------------------------------- + + def _key_value_rows( boxes: np.ndarray, voters: np.ndarray, split: float, y_overlap: np.ndarray, y_overlap_threshold: float, span: float ) -> bool: @@ -215,17 +388,13 @@ def _key_value_rows( 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: + if np.count_nonzero(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 + on_left = voters & ((x0 + x1) / 2 <= split) + left, right = np.flatnonzero(on_left), np.flatnonzero(voters & ~on_left) height = float(np.median(y1[voters] - y0[voters])) - if not height > 0: + if right.size == 0 or left.size == 0 or 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: @@ -237,8 +406,43 @@ def _key_value_rows( 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 + return float(np.median((x0[right][paired] - ends[paired]) / height)) >= _KEY_VALUE_MIN_GAP + + +def _column_gutter( + boxes: np.ndarray, voters: np.ndarray, y_overlap: np.ndarray, y_overlap_threshold: float +) -> np.ndarray | None: + """Detect a multi-column page: a vertical line separating the voters into two groups crossed by few of them + + Args: + boxes: the (N, 4) boxes in canonical LTR space + voters: the mask of the boxes used to detect the columns (all of them with fewer than 3 voters) + 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 + + Returns: + the mask of the boxes crossing the gutter for a multi-column page, None for a single-column page + """ + if np.count_nonzero(voters) < 3: + voters = np.ones(boxes.shape[0], dtype=bool) + x0, x1 = boxes[:, 0], boxes[:, 2] + vx0, vx1, num_voters = x0[voters], x1[voters], int(np.count_nonzero(voters)) + if num_voters < 3: + return None + # Measured on the voters only, so that a wide figure does not shift the search window + left_edge = float(vx0.min()) + span = float(vx1.max()) - left_edge or 1.0 + tolerance = max(1, int(0.05 * num_voters)) + centers = (vx0 + vx1) / 2 + lo, hi = left_edge + 0.25 * span, left_edge + 0.75 * span + for split in np.unique(vx1[(vx1 >= lo) & (vx1 <= hi)]): + 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): + return None # rows of keys and values (e.g. a receipt): read row after row + return np.minimum(x1 - split, split - x0) > 0.02 * span + return None def _topological_order( @@ -266,19 +470,17 @@ def _topological_order( x0, y0, x1, y1 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] x_overlap = _overlap_ratios(x0, x1) y_overlap = _overlap_ratios(y0, y1) - # Strict total orders on both axes, so that the induced relations cannot create 2-cycles - x_rank = _strict_rank(x0, x1) - y_rank = _strict_rank(y0, y1) + x_rank, y_rank = _strict_rank(x0, x1), _strict_rank(y0, y1) - # i -> j edges: i must be read before j. The matrices are combined in place: every intermediate is an - # N x N array, so materializing them all at once dominates the memory of this function on dense pages. + # edges[i, j]: i must be read before j. The N x N matrices are combined in place, since materializing them all at + # once dominates the memory of this function on dense pages. x_linked = x_overlap > x_overlap_threshold - edges = y_rank[:, None] < y_rank[None, :] # i is above j - edges &= x_linked + edges = y_rank[:, None] < y_rank[None, :] # i is above j... + edges &= x_linked # ... in the same column 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 + same_row &= x_rank[:, None] < x_rank[None, :] # i is on the left of j on the same row # 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)) @@ -291,88 +493,57 @@ def _topological_order( del same_row np.fill_diagonal(edges, False) - in_degree = edges.sum(axis=0) - emitted = np.zeros(num_boxes, dtype=bool) - order: list[int] = [] - last = -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)) - if num_voters >= 3: - # Measured on the voters only, so that a wide figure does not shift the search window - left_edge = float(vx0.min()) - span = float(vx1.max()) - left_edge or 1.0 - tolerance = max(1, int(0.05 * num_voters)) - centers = (vx0 + vx1) / 2 - lo, hi = left_edge + 0.25 * span, left_edge + 0.75 * span - for split in np.unique(vx1[(vx1 >= lo) & (vx1 <= hi)]): - 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 - - # 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. + # A multi-column page favors the continuation of the current column (cf. `_continuation`) + spanning = (x1 - x0) > 0.5 * (float(x1.max() - x0.min()) or 1.0) + crossing = _column_gutter( + boxes, np.ones(num_boxes, dtype=bool) if voters is None else voters, y_overlap, y_overlap_threshold + ) + multi_column = crossing is not None + # Column components: elements connected through horizontal overlap, except those which must not merge the columns + # (page-wide or crossing the gutter). 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. + bridging = spanning | crossing if multi_column else spanning 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] + def _continuation(last: int, ready: np.ndarray) -> np.ndarray: + """The available elements continuing the reading after `last`: below it in its column (multi-column page), + else on its row, else in its column component""" + if multi_column: + below = ready[(x_overlap[last, ready] > x_overlap_threshold) & (y0[ready] >= y0[last])] + if below.size > 0: + return below + level = ready[y_overlap[last, ready] > y_overlap_threshold] + if level.size > 0 or spanning[last]: + return level if level.size > 0 else ready + same_column = ready[component[ready] == component[last]] + return same_column if same_column.size > 0 else ready + + in_degree = edges.sum(axis=0) + emitted = np.zeros(num_boxes, dtype=bool) + order: list[int] = [] 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) - ready = np.flatnonzero(~emitted) - candidates = ready + candidates = np.flatnonzero(~emitted) else: - # Favor the continuation of the current column: the closest available element located below the - # last emitted one with a horizontal overlap - candidates = ( - ready[(x_overlap[last, ready] > x_overlap_threshold) & (y0[ready] >= y0[last])] - if last >= 0 and multi_column - else np.empty(0, dtype=int) - ) - if candidates.size == 0 and last >= 0: - candidates = ready[y_overlap[last, ready] > y_overlap_threshold] - if candidates.size == 0 and last >= 0 and not spanning[last]: - # Continuation broken (gap, fragment): stay in the same column before switching to another - same_column = ready[component[ready] == component[last]] - candidates = same_column if same_column.size else ready - elif candidates.size == 0: - candidates = ready + candidates = _continuation(order[-1], ready) if order else ready # Topmost first, then leftmost (in canonical space) next_idx = int(candidates[np.lexsort((x0[candidates], y0[candidates]))[0]]) order.append(next_idx) emitted[next_idx] = True in_degree = in_degree - edges[next_idx] - last = next_idx - return order -def _is_inside(boxes: np.ndarray, idx: int, target: int, min_coverage: float = 0.5) -> bool: - """Check whether box `idx` lies inside box `target`, by at least `min_coverage` of its area""" - x0, y0, x1, y1 = boxes[idx] - tx0, ty0, tx1, ty1 = boxes[target] - inter = max(min(x1, tx1) - max(x0, tx0), 0.0) * max(min(y1, ty1) - max(y0, ty0), 0.0) - return inter >= min_coverage * max((x1 - x0) * (y1 - y0), 1e-9) +# --- Floats & captions ------------------------------------------------------------------------------------------------ def _outer_floats(idcs: list[int], boxes: np.ndarray, labels: list[str]) -> list[int]: """Return the floats among `idcs` that are not nested in a larger float""" - areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + areas = _box_areas(boxes) floats = [idx for idx in idcs if labels[idx] in _FLOAT_LABELS] return [ idx @@ -386,7 +557,7 @@ def _float_members(idcs: list[int], boxes: np.ndarray, labels: list[str]) -> dic An element belongs to the smallest outer float covering at least half of it. """ - areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + areas = _box_areas(boxes) outer = _outer_floats(idcs, boxes, labels) members: dict[int, list[int]] = {} for idx in idcs: @@ -411,15 +582,12 @@ def _stands_beside(boxes: np.ndarray, run: list[int], region: list[int]) -> bool """ rx0, ry0 = boxes[run, 0].min(), boxes[run, 1].min() rx1, ry1 = boxes[run, 2].max(), boxes[run, 3].max() - for idx in region: - x0, y0, x1, y1 = boxes[idx] + for x0, y0, x1, y1 in boxes[region]: # Level with the run: at least half of the element height - v_inter = min(y1, ry1) - max(y0, ry0) - if v_inter < 0.5 * max(y1 - y0, 1e-9): + if min(y1, ry1) - max(y0, ry0) < 0.5 * max(y1 - y0, 1e-9): continue # Beside the run: horizontal overlap of at most 10% of the narrower width - h_inter = min(x1, rx1) - max(x0, rx0) - if h_inter <= 0.1 * max(min(x1 - x0, rx1 - rx0), 1e-9): + if min(x1, rx1) - max(x0, rx0) <= 0.1 * max(min(x1 - x0, rx1 - rx0), 1e-9): return True return False @@ -456,25 +624,24 @@ def _is_float(idx: int) -> bool: if not _is_float(out[pos]): pos += 1 continue + # out[pos:end] is a run of floats, read after out[pos - 1]: does it interrupt the region of that element? end = pos while end < len(out) and _is_float(out[end]): end += 1 - # out[pos:end] is a run of floats, read after the element out[pos - 1] group = int(groups[out[pos - 1]]) if pos > 0 else -1 - if group < 0 or end == len(out) or int(groups[out[end]]) != group: + interrupts = group >= 0 and end < len(out) and int(groups[out[end]]) == group + if interrupts and boxes is not None: + region = [idx for idx in out if int(groups[idx]) == group and not _is_float(idx)] + interrupts = _stands_beside(boxes, out[pos:end], region) + if not interrupts: pos = end continue + # Move the run after the rest of the region stop = end while stop < len(out) and not _is_float(out[stop]) and int(groups[out[stop]]) == group: stop += 1 - if boxes is not None: - region = [idx for idx in out if int(groups[idx]) == group and not _is_float(idx)] - if not _stands_beside(boxes, out[pos:end], region): - pos = end - continue out[pos:stop] = out[end:stop] + out[pos:end] - # Check the moved run again: it can interrupt the same region further down - pos = stop - (end - pos) + pos = stop - (end - pos) # check the moved run again: it can interrupt the same region further down return out @@ -532,15 +699,15 @@ def _read_after(target: int, unit: list[int]) -> int: for unit in _caption_units(caption_idcs, region_groups): cx0, cy0 = boxes[unit, 0].min(), boxes[unit, 1].min() cx1, cy1 = boxes[unit, 2].max(), boxes[unit, 3].max() - best_target, best_dist = -1, float("inf") - for target in float_idcs: - dist = _caption_distance((cx0, cy0, cx1, cy1), boxes[target]) + target, best_dist = -1, float("inf") + for flt in float_idcs: + dist = _caption_distance((cx0, cy0, cx1, cy1), boxes[flt]) if dist < best_dist: - best_target, best_dist = target, dist - if best_target >= 0 and best_dist <= max_distance: + target, best_dist = flt, dist + if target >= 0 and best_dist <= max_distance: # Read right before the float when above its center, otherwise after it - above = (cy0 + cy1) / 2 <= (boxes[best_target, 1] + boxes[best_target, 3]) / 2 - pos = order.index(best_target) if above else _read_after(best_target, unit) + above = (cy0 + cy1) / 2 <= (boxes[target, 1] + boxes[target, 3]) / 2 + pos = order.index(target) if above else _read_after(target, unit) else: in_column = [idx for idx in order if min(boxes[idx, 2], cx1) - max(boxes[idx, 0], cx0) > 0] above_idcs = [idx for idx in in_column if (boxes[idx, 1] + boxes[idx, 3]) / 2 < cy0] @@ -553,68 +720,72 @@ def _read_after(target: int, unit: list[int]) -> int: return order -def deskew_reading_geometries( +# --- Public ordering functions ---------------------------------------------------------------------------------------- + + +def _canonical_boxes( geoms: Sequence[Any] | np.ndarray, - region_geoms: Sequence[Any] | np.ndarray | None = None, - page_shape: tuple[int, int] | None = None, - angle_geoms: Sequence[Any] | np.ndarray | None = None, - min_angle: float = 1.0, -) -> tuple[list[Any], list[Any]]: - """De-skew rotated geometries so the reading order can be computed in an upright frame. + direction: str, + labels: Sequence[str | None] | None, + region_groups: Sequence[int] | None, + page_shape: tuple[int, int] | None, + angle_geoms: Sequence[Any] | np.ndarray | None, +) -> np.ndarray: + """Check the inputs of the ordering functions and return the de-skewed boxes in canonical LTR space""" + if direction not in SUPPORTED_DIRECTIONS[1:]: + raise ValueError(f"invalid reading direction '{direction}', should be one of {SUPPORTED_DIRECTIONS[1:]}") + geoms, _ = deskew_reading_geometries(geoms, page_shape=page_shape, angle_geoms=angle_geoms) + boxes = _to_boxes(geoms) + if labels is not None and len(labels) != boxes.shape[0]: + raise ValueError(f"Incompatible number of labels ({len(labels)}) and geometries ({boxes.shape[0]})") + if region_groups is not None and len(region_groups) != boxes.shape[0]: + raise ValueError( + f"Incompatible number of region groups ({len(region_groups)}) and geometries ({boxes.shape[0]})" + ) + return _to_canonical_ltr(boxes, direction) - Args: - geoms: geometries of the elements to order, in any docTR format (cf. `sort_reading_order`) - region_geoms: optional geometries of the layout regions, de-skewed together with the elements - page_shape: the page dimensions (height, width). Required for an exact angle on non-square pages with - relative coordinates, since relative coordinates distort angles by the aspect ratio. - angle_geoms: optional reading-oriented 4-point polygons used to estimate the page angle - min_angle: minimum estimated angle (in degrees) to trigger the de-skew. Ordering is affected as soon - as the drift along a line approaches the line height (about 1 degree for a page-wide line), and a - small rigid rotation cannot change the order of an upright page, hence the low default. Beyond 45 - degrees the corner identification is ambiguous (an upstream orientation correction is needed) and - nothing is done. - Returns: - the (possibly de-skewed) element and region geometries - """ - region_geoms = list(region_geoms) if region_geoms is not None else [] - pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in geoms] - if len(pts) == 0 or any(p.shape[0] != 4 for p in pts): - return list(geoms), region_geoms # straight geometries: nothing to de-skew - height, width = page_shape if page_shape is not None else (1, 1) - scale = np.array([width, height], dtype=np.float64) - angle_source = angle_geoms if angle_geoms is not None else [] - angle_pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in angle_source] - if len(angle_pts) > 0 and all(p.shape[0] == 4 for p in angle_pts): - # Detection polygons are already reading-oriented (cf. `estimate_page_angle`): keep their vertex order - angle = estimate_page_angle(np.stack(angle_pts) * scale) - else: - # Normalize the vertex order (TL, TR, BR, BL) so the estimation does not depend on the vertex - angle = estimate_page_angle(np.stack([order_points(p * scale) for p in pts])) - if not np.isfinite(angle) or abs(angle) < min_angle or abs(angle) >= 45: - return list(geoms), region_geoms - # Rigid rotation zeroing the estimated angle; the center is irrelevant for ordering purposes - theta = np.deg2rad(angle) - rot = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) - all_pts = np.concatenate(pts, axis=0) * scale - center = all_pts.mean(axis=0) +def _sort( + canonical: np.ndarray, + labels: Sequence[str | None] | None, + x_overlap_threshold: float, + y_overlap_threshold: float, + caption_max_distance: float, + column_voters: Sequence[bool] | None, + region_groups: Sequence[int] | None, +) -> list[int]: + """Order boxes in canonical LTR space (cf. `sort_reading_order`)""" + num_boxes = canonical.shape[0] + if num_boxes <= 1: + return list(range(num_boxes)) + voters = np.ones(num_boxes, dtype=bool) if column_voters is None else np.asarray(column_voters, dtype=bool) - def _rotate(points: np.ndarray) -> np.ndarray: - return ((points * scale - center) @ rot.T + center) / scale + def _order(idcs: list[int]) -> list[int]: + sub_order = _topological_order(canonical[idcs], x_overlap_threshold, y_overlap_threshold, voters[idcs]) + return [idcs[i] for i in sub_order] - deskewed = [_rotate(p) for p in pts] + if labels is None: + return _order(list(range(num_boxes))) - def _corners(points: np.ndarray) -> np.ndarray: - # Straight ((xmin, ymin), (xmax, ymax)) regions must be expanded to their 4 corners before rotating, - # otherwise only the diagonal would be rotated and the resulting extent would be underestimated - if points.shape[0] == 4: - return points - (x0, y0), (x1, y1) = points.min(axis=0), points.max(axis=0) - return np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) + norm_labels = [normalize_layout_label(label) for label in labels] + groups: dict[str, list[int]] = {"header": [], "body": [], "caption": [], "footnote": [], "footer": []} + for idx, label in enumerate(norm_labels): + role = layout_label_role(label) + groups["body" if role == "float" else role].append(idx) - region_pts = [np.asarray(geom, dtype=np.float64).reshape(-1, 2) for geom in region_geoms] - regions_out = [_rotate(_corners(p)) for p in region_pts] - return deskewed, regions_out + # Read the elements inside a float (e.g. the text of a figure) right after it + members = _float_members(groups["body"], canonical, norm_labels) + owner = {idx: flt for flt, idcs in members.items() for idx in idcs} + body_order = _order(groups["body"]) + inside = {flt: [idx for idx in body_order if owner.get(idx) == flt] for flt in members} + body_order = [elt for idx in body_order if idx not in owner for elt in (idx, *inside.get(idx, []))] + if region_groups is not None: + # A float standing beside a region must not interrupt it + body_order = _defer_floats(body_order, norm_labels, members, region_groups, canonical) + body_order = _attach_captions( + body_order, _order(groups["caption"]), canonical, norm_labels, caption_max_distance, region_groups + ) + return _order(groups["header"]) + body_order + _order(groups["footnote"]) + _order(groups["footer"]) def sort_reading_order( @@ -667,54 +838,67 @@ def sort_reading_order( Returns: the permutation of the input indices which sorts the elements in reading order """ - if direction not in SUPPORTED_DIRECTIONS[1:]: - raise ValueError(f"invalid reading direction '{direction}', should be one of {SUPPORTED_DIRECTIONS[1:]}") - geoms, _ = deskew_reading_geometries(geoms, page_shape=page_shape, angle_geoms=angle_geoms) - boxes = _to_boxes(geoms) - num_boxes = boxes.shape[0] - if labels is not None and len(labels) != num_boxes: - raise ValueError(f"Incompatible number of labels ({len(labels)}) and geometries ({num_boxes})") - if region_groups is not None and len(region_groups) != num_boxes: - raise ValueError(f"Incompatible number of region groups ({len(region_groups)}) and geometries ({num_boxes})") - if num_boxes <= 1: - return list(range(num_boxes)) - - canonical = _to_canonical_ltr(boxes, direction) - - voters = np.ones(num_boxes, dtype=bool) if column_voters is None else np.asarray(column_voters, dtype=bool) + canonical = _canonical_boxes(geoms, direction, labels, region_groups, page_shape, angle_geoms) + return _sort( + canonical, labels, x_overlap_threshold, y_overlap_threshold, caption_max_distance, column_voters, region_groups + ) - def _order(idcs: list[int]) -> list[int]: - if len(idcs) == 0: - return [] - sub_order = _topological_order(canonical[idcs], x_overlap_threshold, y_overlap_threshold, voters[idcs]) - return [idcs[i] for i in sub_order] - if labels is None: - return _order(list(range(num_boxes))) +class _SegmentJoins: + """Criteria to join two consecutive elements of the reading order in the same segment (cf. + `resolve_reading_segments`), in canonical LTR space""" - norm_labels = [normalize_layout_label(label) for label in labels] - groups: dict[str, list[int]] = {"header": [], "body": [], "caption": [], "footnote": [], "footer": []} - for idx, label in enumerate(norm_labels): - role = layout_label_role(label) - groups["body" if role == "float" else role].append(idx) + def __init__( + self, + boxes: np.ndarray, + norm_labels: list[str], + region_groups: Sequence[int] | None, + x_overlap_threshold: float, + y_overlap_threshold: float, + ) -> None: + self.boxes, self.labels, self.groups = boxes, norm_labels, region_groups + self.roles = [layout_label_role(label) for label in norm_labels] + self.median_height = float(np.median(boxes[:, 3] - boxes[:, 1])) + self.x_overlap_threshold, self.y_overlap_threshold = x_overlap_threshold, y_overlap_threshold + + @staticmethod + def _overlap(start: float, end: float, other_start: float, other_end: float) -> float: + """1D overlap of two intervals, relative to the shorter one""" + inter = min(end, other_end) - max(start, other_start) + return inter / max(min(end - start, other_end - other_start), 1e-9) + + def _same_region(self, prev: int, cur: int) -> bool: + return self.groups is None or self.groups[prev] == self.groups[cur] + + def compatible(self, prev: int, cur: int) -> bool: + """Same label (not a float) and same layout region""" + return self.labels[prev] == self.labels[cur] and self.roles[cur] != "float" and self._same_region(prev, cur) + + def stacked(self, row: np.ndarray, cur: int, paragraph_gap: float) -> bool: + """Right below the last row of the segment, in the same column""" + box = self.boxes[cur] + return ( + self._overlap(row[0], row[2], box[0], box[2]) > self.x_overlap_threshold + and box[1] - row[3] <= paragraph_gap * self.median_height + ) - body_order = _order(groups["body"]) - # Read the elements inside a float (e.g. the text of a figure) right after it - members = _float_members(groups["body"], canonical, norm_labels) - if members: - owner = {idx: flt for flt, idcs in members.items() for idx in idcs} - read_inside: dict[int, list[int]] = {flt: [] for flt in members} - for idx in body_order: - if idx in owner: - read_inside[owner[idx]].append(idx) - body_order = [elt for idx in body_order if idx not in owner for elt in (idx, *read_inside.get(idx, []))] - if region_groups is not None: - # A float standing beside a region must not interrupt it - body_order = _defer_floats(body_order, norm_labels, members, region_groups, canonical) - body_order = _attach_captions( - body_order, _order(groups["caption"]), canonical, norm_labels, caption_max_distance, region_groups - ) - return _order(groups["header"]) + body_order + _order(groups["footnote"]) + _order(groups["footer"]) + def same_row(self, prev: int, cur: int, row_gap: float, aspect: float) -> bool: + """On the row of the previous element, after it: the rest of a line split at a wide word gap (at most + `row_gap` median heights, any gap within a region of body text)""" + box, prev_box = self.boxes[cur], self.boxes[prev] + return ( + self._overlap(prev_box[1], prev_box[3], box[1], box[3]) > self.y_overlap_threshold + and box[0] > prev_box[0] + and ( + (box[0] - prev_box[2]) * aspect <= row_gap * self.median_height + or ( + self.groups is not None + and self.groups[cur] >= 0 + and self._same_region(prev, cur) + and self.roles[cur] == "body" + ) + ) + ) def resolve_reading_segments( @@ -766,10 +950,10 @@ def resolve_reading_segments( a partition of the input indices into reading-ordered segments (each segment being itself in reading order) """ - # De-skew once so the ordering and the vertical-gap measurement share the same upright frame - geoms, _ = deskew_reading_geometries(geoms, page_shape=page_shape, angle_geoms=angle_geoms) + # The elements are ordered on their upright geometries, in which the gaps between them are measured as well + upright, _ = deskew_reading_geometries(geoms, page_shape=page_shape, angle_geoms=angle_geoms) order = sort_reading_order( - geoms, + upright, direction=direction, labels=labels, x_overlap_threshold=x_overlap_threshold, @@ -780,166 +964,30 @@ def resolve_reading_segments( ) if len(order) == 0: return [] - canonical = _to_canonical_ltr(_to_boxes(geoms), direction) - median_height = float(np.median(canonical[:, 3] - canonical[:, 1])) + canonical = _to_canonical_ltr(_to_boxes(upright), direction) 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))] - + joins = _SegmentJoins(canonical, norm_labels, region_groups, x_overlap_threshold, y_overlap_threshold) segments: list[list[int]] = [[order[0]]] 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]) - 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 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] - 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 same_region - and (stacked or same_row) - ) - if same_segment: + stacked, same_row = joins.stacked(row, cur, paragraph_gap), joins.same_row(prev, cur, row_gap, aspect) + joined = joins.compatible(prev, cur) and (stacked or same_row) + if joined: 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() + box = canonical[cur] + row = ( + np.concatenate((np.minimum(row[:2], box[:2]), np.maximum(row[2:], box[2:]))) + if joined and same_row and not stacked + else box.copy() + ) return segments -def assign_layout_labels( - geoms: Sequence[Any] | np.ndarray, - layout_geoms: Sequence[Any] | np.ndarray, - layout_labels: Sequence[str], - min_coverage: float = 0.5, - page_shape: tuple[int, int] | None = None, - angle_geoms: Sequence[Any] | np.ndarray | None = None, -) -> list[str | None]: - """Assign a layout label to each element based on its overlap with the detected layout regions. - - Each element receives the label of the region covering the largest share of its area, provided this share reaches - `min_coverage`; otherwise its label is None (treated as regular body content). Among nested regions, the smallest - one wins (e.g. 'Caption' over 'Picture'), except for a float nested in page furniture (e.g. a logo in the page - header): its text takes the furniture label. - - Args: - geoms: geometries of the elements to label, in any docTR format - layout_geoms: geometries of the layout regions, in any docTR format - layout_labels: labels of the layout regions (e.g. `[region.type for region in page.layout]`) - min_coverage: minimum share of an element's area a region must cover to assign its label - page_shape: the page dimensions (height, width), used to de-skew rotated pages exactly (cf. - `deskew_reading_geometries`) - angle_geoms: optional reading-oriented 4-point polygons (typically the page's word polygons) used to - estimate the page angle on rotated pages (cf. `deskew_reading_geometries`) - - Returns: - the label of each element (None when no region covers it enough) - """ - # De-skew elements and regions together so the coverage is measured on tight boxes in the same frame - geoms, layout_geoms = deskew_reading_geometries(geoms, layout_geoms, page_shape=page_shape, angle_geoms=angle_geoms) - boxes, regions = _to_boxes(geoms), _to_boxes(layout_geoms) - if len(layout_labels) != regions.shape[0]: - raise ValueError(f"Incompatible number of labels ({len(layout_labels)}) and regions ({regions.shape[0]})") - if boxes.shape[0] == 0 or regions.shape[0] == 0: - return [None] * boxes.shape[0] - - return [ - str(layout_labels[reg_idx]) if reg_idx >= 0 else None - for reg_idx in _covering_regions(boxes, regions, min_coverage, layout_labels) - ] - - -def _region_groups(regions: np.ndarray, labels: Sequence[str | None], min_overlap: float = 0.1) -> list[int]: - """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]) - 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, - min_coverage: float, - region_labels: Sequence[str | None] | None = None, -) -> list[int]: - """Return the index of the region covering each (N, 4) box by at least `min_coverage`, -1 if none - - Regions covering a box (nearly) as much as the best one tie: the smallest one wins, unless it is a float tied with - page furniture. - """ - inter_w = np.minimum(boxes[:, None, 2], regions[None, :, 2]) - np.maximum(boxes[:, None, 0], regions[None, :, 0]) - inter_h = np.minimum(boxes[:, None, 3], regions[None, :, 3]) - np.maximum(boxes[:, None, 1], regions[None, :, 1]) - inter = np.clip(inter_w, 0, None) * np.clip(inter_h, 0, None) - areas = np.clip((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]), 1e-9, None) - coverage = inter / areas[:, None] - region_areas = (regions[:, 2] - regions[:, 0]) * (regions[:, 3] - regions[:, 1]) - qualifies = coverage >= np.maximum(coverage.max(axis=1, keepdims=True) - 0.1, min_coverage) - # Regions nested in the best region of a box also tie when they cover most of it - rx0, ry0, rx1, ry1 = regions[:, 0], regions[:, 1], regions[:, 2], regions[:, 3] - reg_inter_w = np.minimum(rx1[:, None], rx1[None, :]) - np.maximum(rx0[:, None], rx0[None, :]) - reg_inter_h = np.minimum(ry1[:, None], ry1[None, :]) - np.maximum(ry0[:, None], ry0[None, :]) - reg_inter = np.clip(reg_inter_w, 0, None) * np.clip(reg_inter_h, 0, None) - nested = reg_inter >= 0.9 * np.clip(region_areas, 1e-9, None)[:, None] # nested[i, j]: region i lies inside j - qualifies |= (coverage >= max(min_coverage, 0.75)) & nested[:, coverage.argmax(axis=1)].T - ranked = np.where(qualifies, region_areas[None, :], np.inf) - best = ranked.argmin(axis=1) - result = [int(reg) if np.isfinite(ranked[idx, reg]) else -1 for idx, reg in enumerate(best)] - if region_labels is None or len(region_labels) != regions.shape[0]: - return result - - roles = [layout_label_role(label) for label in region_labels] - is_float = np.array([role == "float" for role in roles], dtype=bool) - is_furniture = np.array([role in _FURNITURE_ROLES for role in roles], dtype=bool) - if not (is_float.any() and is_furniture.any()): - return result - furniture_ranked = np.where(qualifies & is_furniture[None, :], region_areas[None, :], np.inf) - for idx, reg in enumerate(result): - if reg >= 0 and is_float[reg] and np.isfinite(furniture_ranked[idx].min()): - result[idx] = int(furniture_ranked[idx].argmin()) - return result - - class ReadingOrderPredictor(NestedObject): """Language-independent reading order estimation for document elements. diff --git a/tests/common/test_io_exporters.py b/tests/common/test_io_exporters.py index 945ea32c8e..95a27116d7 100644 --- a/tests/common/test_io_exporters.py +++ b/tests/common/test_io_exporters.py @@ -1,5 +1,6 @@ import json import re +from xml.etree import ElementTree as ET import cv2 import numpy as np @@ -9,20 +10,27 @@ from doctr.io import elements from doctr.io.exporters import ( AsciiDocExporter, + DocumentExportsMixin, HTMLExporter, + KIEPageExportsMixin, MarkdownExporter, + PageExportsMixin, TextExporter, XMLExporter, + _covering_region_indices, _join_caption_lines, + _PageTextExporter, ordered_line_words, page_reading_order, to_json_safe, ) from doctr.io.figures import FigureEncoder +_ORIENTATION = {"value": 0, "confidence": None} + def _word_at(text, x0, y0, x1, y1): - return elements.Word(text, 0.95, ((x0, y0), (x1, y1)), 0.9, {"value": 0, "confidence": None}) + return elements.Word(text, 0.95, ((x0, y0), (x1, y1)), 0.9, _ORIENTATION) def _line_at(text, x0, y0, x1, y1, rtl=False): @@ -35,10 +43,31 @@ def _line_at(text, x0, y0, x1, y1, rtl=False): ]) +def _page(lines, image=None, dims=(1000, 800), **kwargs): + """A page whose lines form a single block, as in the default builder configuration""" + image = np.zeros((10, 10, 3), dtype=np.uint8) if image is None else image + return elements.Page(image, [elements.Block(lines=lines)], 0, dims, **kwargs) + + def _texts(items): return [[" ".join(word.value for word in line.words) for line in item.lines] for item in items] +def _flat_texts(items): + return [text for item in _texts(items) for text in item] + + +def _prediction(value, y0, y1): + return elements.Prediction(value, 0.9, ((0.1, y0), (0.9, y1)), 0.9, _ORIENTATION) + + +def _kie_page(predictions=None): + """A KIE page whose predictions are stored bottom-first""" + if predictions is None: + predictions = {CLASS_NAME: [_prediction("second", 0.5, 0.6), _prediction("first", 0.1, 0.2)]} + return elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) + + 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)] @@ -54,12 +83,53 @@ def _reading_order_page(): elements.LayoutElement("List-item", 0.97, ((0.06, 0.38), (0.48, 0.45))), elements.LayoutElement("Page-footer", 0.97, ((0.35, 0.94), (0.65, 0.98))), ] - return elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), layout=layout - ) + return _page(lines, layout=layout) + + +def _order_cells(): + return [ + 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), + ] + + +def _table_page(num_rows=2, num_cols=1, rotated=False): + """A page with an intro line and one recognized table""" + geo = ((0.1, 0.6), (0.4, 0.66), (0.4, 0.71), (0.1, 0.71)) if rotated else ((0.1, 0.6), (0.4, 0.71)) + cells = [ + elements.TableCell("head", 0.9, ((0.1, 0.6), (0.4, 0.65)), 0, 0, 0, 0), + elements.TableCell("body", 0.9, ((0.1, 0.66), (0.4, 0.71)), 1, 1, 0, 0), + ] + table = elements.Table(cells=cells, num_rows=num_rows, num_cols=num_cols, geometry=geo) + return _page([_line_at("intro line", 0.1, 0.1, 0.9, 0.15)], tables=[table]) + + +def _rotated_line(text, x0, y0, x1, y1, deg, height, width): + """A line rotated by `deg` degrees around the center of a page of the given size (in pixels)""" + angle = np.deg2rad(deg) + rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) + center = np.array([width / 2, height / 2]) + tokens = text.split() + step = (x1 - x0) / len(tokens) + words = [] + for idx, token in enumerate(tokens): + pts = np.array([ + [(x0 + idx * step) * width, y0 * height], + [(x0 + (idx + 0.9) * step) * width, y0 * height], + [(x0 + (idx + 0.9) * step) * width, y1 * height], + [(x0 + idx * step) * width, y1 * height], + ]) + pts = ((pts - center) @ rot.T + center) / [width, height] + words.append(elements.Word(token, 0.9, tuple(tuple(pt) for pt in pts), 0.9, _ORIENTATION)) + return elements.Line(words) -def test_page_items_in_reading_order(): +# --- Reading order --------------------------------------------------------------------------------------------------- + + +def test_page_reading_order(): page = _reading_order_page() items = page.items_in_reading_order() assert all(isinstance(item, elements.Block) for item in items) @@ -69,115 +139,122 @@ def test_page_items_in_reading_order(): assert rendered.index("left line 0 left line 1 left line 2") < rendered.index( "right line 0 right line 1 right line 2" ) + # The JSON export follows the reading order too... + exported = [ + " ".join(word["value"] for line in block["lines"] for word in line["words"]) + for block in page.export()["blocks"] + ] + assert exported[0] == "A Two Column Study" + assert exported[-1] == "Page 3 of 12" + assert exported.index("left line 0 left line 1 left line 2") < exported.index( + "right line 0 right line 1 right line 2" + ) + # ... which can be opted out of + raw = page.export(reading_order=False)["blocks"] + assert len(raw) == 1 and len(raw[0]["lines"]) == 9 + # render() too, with the page furniture that can be dropped, exactly like in the Markdown / HTML exports + assert page.render().splitlines()[0] == "A Two Column Study" + assert "Page 3 of 12" in page.render() + assert "Page 3 of 12" not in page.render(include_furniture=False) + # Multi-block pages are ordered at the block level top = elements.Block([_line_at("first words", 0.1, 0.1, 0.9, 0.15)]) 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"] + # The blocks are stored right-first, but render() linearizes them like the other exporters + left = elements.Block( + lines=[_line_at("left top", 0.08, 0.1, 0.45, 0.13), _line_at("left low", 0.08, 0.2, 0.45, 0.23)] + ) + right = elements.Block( + lines=[_line_at("right top", 0.55, 0.1, 0.92, 0.13), _line_at("right low", 0.55, 0.2, 0.92, 0.23)] + ) + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [right, left], 0, (1000, 800)) + assert page.render() == "left top\n\nleft low\n\nright top\n\nright low" + assert page.render(block_break=" | ") == "left top | left low | right top | right low" + # Single-block and empty pages + single = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [left], 0, (1000, 800)) + assert single.render(block_break=" | ") == "left top | left low" + assert elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [], 0, (1000, 800)).render() == "" # 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", - ] + assert _flat_texts(page_reading_order(_page(lines, dims=dims))[0]) == ["first line here", "second part"] - -def test_page_export_as_markdown(): - page = _reading_order_page() - markdown = page.export_as_markdown() - parts = markdown.split("\n\n") - assert parts[0] == "# A Two Column Study" - assert parts[1] == "left line 0\nleft line 1\nleft line 2" - # The list item belongs to the left column, hence it is read before the right column - assert parts[2] == "- \\- item one" # list item, with the raw OCR dash escaped - assert parts[3] == "right line 0\nright line 1\nright line 2" - assert parts[4] == "Page 3 of 12" - # Page furniture can be dropped - assert "Page 3 of 12" not in page.export_as_markdown(include_furniture=False) - # Markdown structural characters are escaped by default + # Artefacts survive the regrouping + artefact = elements.Artefact("qr_code", 0.9, ((0.1, 0.8), (0.2, 0.9))) page = elements.Page( np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block([_line_at("*bold* #tag [link]", 0.1, 0.1, 0.9, 0.15)])], + [elements.Block(lines=[_line_at("a line", 0.1, 0.1, 0.9, 0.15)], artefacts=[artefact])], 0, (1000, 800), ) - assert page.export_as_markdown() == "\\*bold\\* \\#tag \\[link\\]" - assert page.export_as_markdown(escape=False) == "*bold* #tag [link]" - # Empty pages export to an empty string - assert elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [], 0, (1000, 800)).export_as_markdown() == "" + assert page.export()["blocks"][0]["artefacts"] == [artefact.export()] + # ... also those of a block without any line + artefact = elements.Artefact("qr_code", 0.9, ((0.7, 0.7), (0.9, 0.9))) + orphan = elements.Block(lines=[], artefacts=[artefact], geometry=((0.7, 0.7), (0.9, 0.9)), objectness_score=0.9) + text = elements.Block(lines=[_line_at("some text", 0.1, 0.1, 0.9, 0.15)]) + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [orphan, text], 0, (1000, 800)) + items = page.items_in_reading_order() + assert [artefact.export() for artefact in items[-1].artefacts] == [artefact.export()] + assert page.export()["blocks"][-1]["artefacts"] == [artefact.export()] @pytest.mark.parametrize( - "line, expected", + "height, width, lines, degrees", [ - ("3. not a list", "3\\. not a list"), - ("12) neither", "12\\) neither"), - (" 7.", " 7\\."), - ("3", "3"), - ("2022 was a year", "2022 was a year"), - ("1.5 million", "1.5 million"), - ("- dash", "\\- dash"), - ("# hash", "\\# hash"), + ( + 1000, + 800, + [ + ("big page title", 0.1, 0.05, 0.9, 0.09), + ("left one", 0.1, 0.15, 0.45, 0.19), + ("left two", 0.1, 0.21, 0.45, 0.25), + ("left three", 0.1, 0.27, 0.45, 0.31), + ("right one", 0.55, 0.15, 0.9, 0.19), + ("right two", 0.55, 0.21, 0.9, 0.25), + ("right three", 0.55, 0.27, 0.9, 0.31), + ], + (15, 25), + ), + ( + 800, + 1200, + [ + ("big page title", 0.1, 0.05, 0.9, 0.09), + ("left one", 0.1, 0.15, 0.45, 0.19), + ("left two", 0.1, 0.21, 0.45, 0.25), + ("right one", 0.55, 0.15, 0.9, 0.19), + ("right two", 0.55, 0.21, 0.9, 0.25), + ], + (-35, 35), + ), ], + ids=["portrait", "landscape"], ) -def test_markdown_finalize_line(line, expected): - assert MarkdownExporter().finalize_line(line) == expected - +def test_page_reading_order_rotated(height, width, lines, degrees): + # A rotated page is read like the same page without rotation + def _rotated_page(deg): + return _page([_rotated_line(*line, deg, height, width) for line in lines], dims=(height, width)) -def test_page_export_as_markdown_rtl(): - # Two columns of Arabic text: the right column is read first, and the words of each line are emitted - # from the rightmost to the leftmost one - lines = [ - _line_at("النص في العمود الأيمن", 0.54, 0.1, 0.92, 0.14, rtl=True), - _line_at("النص في العمود الأيسر", 0.08, 0.1, 0.46, 0.14, rtl=True), - ] - page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800)) - markdown = page.export_as_markdown() - assert markdown == "النص في العمود الأيمن\n\nالنص في العمود الأيسر" - # An explicit direction takes precedence over the detection - assert page.export_as_markdown(direction="ltr").startswith("الأيسر") + expected = _rotated_page(0).export_as_markdown() + assert expected.split()[:3] == ["big", "page", "title"] + for deg in degrees: + assert _rotated_page(deg).export_as_markdown().split() == expected.split() -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") +def test_page_reading_order_with_tables(): # 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] - ) + page = _page(lines, tables=[elements.Table(_order_cells(), 2, 2, ((0.05, 0.25), (0.95, 0.28)), 0.95)]) 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]] + assert _flat_texts(blocks) == ["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 = [ @@ -186,20 +263,15 @@ def test_page_export_with_tables(): _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)))], - ) + page = _page(lines, layout=[elements.LayoutElement("Table", 0.9, ((0.08, 0.38), (0.92, 0.54)))]) items, labels, _ = page_reading_order(page) - texts = [text for item in _texts(items) for text in item] + texts = _flat_texts(items) 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) # 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 += [ @@ -213,19 +285,14 @@ def test_page_export_with_tables(): 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] + items, labels, _ = page_reading_order(_page(lines, layout=layout)) + texts = _flat_texts(items) 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_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 = [ @@ -243,232 +310,185 @@ def test_page_reading_order_page_wide_table_region(): ] 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] + items, _, _ = page_reading_order(_page(lines, layout=layout, tables=tables)) + texts = _flat_texts(item for item in items if isinstance(item, elements.Block)) 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() - parts = asciidoc.split("\n\n") - assert parts[0] == "== A Two Column Study" - assert parts[2] == "* {empty}- item one" - assert "Page 3 of 12" not in page.export_as_asciidoc(include_furniture=False) +def test_reading_direction(): + # Two columns of Arabic text: the right column is read first, and the words of each line are emitted + # from the rightmost to the leftmost one + lines = [ + _line_at("النص في العمود الأيمن", 0.54, 0.1, 0.92, 0.14, rtl=True), + _line_at("النص في العمود الأيسر", 0.08, 0.1, 0.46, 0.14, rtl=True), + ] + page = _page(lines) + assert page.export_as_markdown() == "النص في العمود الأيمن\n\nالنص في العمود الأيسر" + # An explicit direction takes precedence over the detection + assert page.export_as_markdown(direction="ltr").startswith("الأيسر") + line = _line_at("one two three", 0.1, 0.1, 0.7, 0.15) + assert [word.render() for word in ordered_line_words(line, "ltr")] == ["one", "two", "three"] + assert [word.render() for word in ordered_line_words(line, "rtl")] == ["three", "two", "one"] + # Vertical pages read their words top to bottom + stacked = elements.Line([_word_at(value, 0.1, 0.1 * idx, 0.2, 0.1 * idx + 0.05) for idx, value in enumerate("abc")]) + assert [word.render() for word in ordered_line_words(stacked, "ttb-rtl")] == ["a", "b", "c"] + assert [word.render() for word in ordered_line_words(stacked, "ttb-ltr")] == ["a", "b", "c"] + # A single-word line short-circuits the per-line detection + single = elements.Line([_word_at("solo", 0.1, 0.1, 0.2, 0.15)]) + assert [word.render() for word in ordered_line_words(single, "auto", auto=True)] == ["solo"] -def test_page_export_as(): + # A vertical direction reaches every exporter page = _reading_order_page() - assert page.export_as("markdown") == page.export_as("md") == page.export_as_markdown() - assert page.export_as("adoc") == page.export_as_asciidoc() - assert page.export_as("text") == page.render() - assert page.export_as("json") == page.export() - assert isinstance(page.export_as("xml")[0], bytes) - assert page.export_as("markdown", include_furniture=False) == page.export_as_markdown(include_furniture=False) - with pytest.raises(ValueError): - page.export_as("yaml") + for export in (page.render, page.export_as_markdown, page.export_as_asciidoc, page.export_as_html): + assert isinstance(export(direction="ttb-ltr"), str) + assert isinstance(page.export_as_xml(direction="ttb-rtl")[0], bytes) -def test_document_export_as_markdown(): - pages = [ - elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block([_line_at(f"page {idx} content", 0.1, 0.1, 0.9, 0.15)])], - idx, - (1000, 800), - ) - for idx in range(2) +def test_document_forwards_reading_order_options(): + left = elements.Block(lines=[_line_at("left col", 0.08, 0.2, 0.45, 0.23)]) + right = elements.Block(lines=[_line_at("right col", 0.55, 0.2, 0.92, 0.23)]) + footer = elements.Block(lines=[_line_at("Page 1 of 2", 0.4, 0.95, 0.6, 0.97)]) + layout = [ + elements.LayoutElement("Text", 0.98, ((0.06, 0.18), (0.48, 0.26))), + elements.LayoutElement("Text", 0.98, ((0.52, 0.18), (0.94, 0.26))), + elements.LayoutElement("Page-footer", 0.97, ((0.35, 0.94), (0.65, 0.98))), ] - doc = elements.Document(pages) - assert doc.export_as_markdown() == "page 0 content\n\n---\n\npage 1 content" - assert doc.export_as_asciidoc() == "page 0 content\n\n<<<\n\npage 1 content" - assert doc.export_as_markdown(page_break="\n\n") == "page 0 content\n\npage 1 content" - assert doc.export_as("markdown") == doc.export_as_markdown() - assert doc.export_as("text") == doc.render() - assert doc.export_as("json") == doc.export() - assert len(doc.export_as("xml")) == 2 - with pytest.raises(ValueError): - doc.export_as("pdf") + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [right, left, footer], 0, (1000, 800), layout=layout) + doc = elements.Document([page, page]) + assert ( + doc.render(page_break=" || ") == "left col\n\nright col\n\nPage 1 of 2 || left col\n\nright col\n\nPage 1 of 2" + ) + assert "Page 1 of 2" not in doc.render(include_furniture=False) + assert doc.render(block_break=" ", page_break=" || ").startswith("left col right col") + assert doc.export_as("text", include_furniture=False) == doc.render(include_furniture=False) + ordered = doc.export()["pages"][0]["blocks"] + assert [block["lines"][0]["words"][0]["value"] for block in ordered] == ["left", "right", "Page"] + raw = doc.export(reading_order=False)["pages"][0]["blocks"] + assert [block["lines"][0]["words"][0]["value"] for block in raw] == ["right", "left", "Page"] + assert doc.export_as("json", reading_order=False) == doc.export(reading_order=False) + json.dumps(doc.export(reading_order=False)) + + # KIE documents + doc = elements.KIEDocument([_kie_page()]) + assert doc.render() == f"{CLASS_NAME}: first\n\n{CLASS_NAME}: second" + # The reading direction reaches KIEPage.render through the document + assert doc.render(direction="ttb-ltr") == f"{CLASS_NAME}: first\n\n{CLASS_NAME}: second" + assert [p["value"] for p in doc.export()["pages"][0]["predictions"][CLASS_NAME]] == ["first", "second"] + raw = doc.export(reading_order=False)["pages"][0]["predictions"][CLASS_NAME] + assert [p["value"] for p in raw] == ["second", "first"] + json.dumps(doc.export(reading_order=False)) -def test_kie_page_export_as_markdown(): - predictions = { - CLASS_NAME: [ - elements.Prediction("second", 0.9, ((0.1, 0.5), (0.9, 0.6)), 0.9, {"value": 0, "confidence": None}), - elements.Prediction("first", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None}), - ] - } - page = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) +def test_kie_page_exports(): + page = _kie_page() + # Predictions are exported in reading order + assert page.render() == f"{CLASS_NAME}: first\n\n{CLASS_NAME}: second" assert page.export_as_markdown() == f"**{CLASS_NAME}**\n\n- first\n- second" assert page.export_as_asciidoc() == f"*{CLASS_NAME}*\n\n* first\n* second" assert page.export_as("md") == page.export_as_markdown() with pytest.raises(ValueError): page.export_as("yaml") - doc = elements.KIEDocument([page]) - assert doc.export_as_markdown() == page.export_as_markdown() + assert elements.KIEDocument([page]).export_as_markdown() == page.export_as_markdown() + assert [p["value"] for p in page.export()["predictions"][CLASS_NAME]] == ["first", "second"] + assert [p["value"] for p in page.export(reading_order=False)["predictions"][CLASS_NAME]] == ["second", "first"] + xml = page.export_as_xml()[0].decode() + assert xml.index(">first<") < xml.index(">second<") + # Classes without any prediction are skipped by every text exporter + kie = _kie_page({CLASS_NAME: [_prediction("second", 0.5, 0.6), _prediction("first", 0.1, 0.2)], "empty": []}) + assert TextExporter().export_kie_page(kie) == f"{CLASS_NAME}:\n\nfirst\nsecond" + assert "empty" not in MarkdownExporter().export_kie_page(kie) + assert "empty" not in AsciiDocExporter().export_kie_page(kie) + assert "empty" not in HTMLExporter().export_kie_page(kie) + # export_document dispatches on the page type and uses the format-specific page break + mixed = elements.Document([_page([_line_at("a page", 0.1, 0.1, 0.9, 0.15)]), kie]) + assert TextExporter().export_document(mixed) == f"a page\n\n\n\n{CLASS_NAME}:\n\nfirst\nsecond" + assert TextExporter().export_document(mixed, page_break=" // ").startswith("a page // ") -def test_page_export_as_markdown_list_items(): - # Three separate single-line list items, each covered by its own List-item region -> three bullets - lines = [_line_at(f"item number {idx}", 0.1, 0.1 + 0.1 * idx, 0.5, 0.13 + 0.1 * idx) for idx in range(3)] - layout = [ - elements.LayoutElement("List-item", 0.9, ((0.08, 0.09 + 0.1 * idx), (0.52, 0.14 + 0.1 * idx))) - for idx in range(3) - ] - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), layout=layout - ) - assert page.export_as_markdown() == "- item number 0\n- item number 1\n- item number 2" - assert page.export_as_asciidoc() == "* item number 0\n* item number 1\n* item number 2" - # each list item is its own block in reading order - items = page.items_in_reading_order() - assert len(items) == 3 - assert all(len(item.lines) == 1 for item in items) +# --- Text formats ---------------------------------------------------------------------------------------------------- -def test_page_export_as_markdown_wrapped_list_item(): - # A single list item wrapped over three visual lines (one region) must render as ONE bullet, while a - # second item (another region) is a second bullet. - lines = [ - _line_at("first item wrapping over", 0.1, 0.10, 0.9, 0.13), - _line_at("several visual lines here", 0.1, 0.14, 0.9, 0.17), - _line_at("until it finally ends", 0.1, 0.18, 0.6, 0.21), - _line_at("second short item", 0.1, 0.26, 0.5, 0.29), - ] - layout = [ - elements.LayoutElement("List-item", 0.9, ((0.08, 0.09), (0.92, 0.22))), - elements.LayoutElement("List-item", 0.9, ((0.08, 0.25), (0.52, 0.30))), - ] - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), layout=layout - ) - assert page.export_as_markdown() == ( - "- first item wrapping over several visual lines here until it finally ends\n- second short item" - ) - items = page.items_in_reading_order() - assert len(items) == 2 - assert len(items[0].lines) == 3 and len(items[1].lines) == 1 +def test_page_export_as_markdown(): + page = _reading_order_page() + parts = page.export_as_markdown().split("\n\n") + assert parts[0] == "# A Two Column Study" + assert parts[1] == "left line 0\nleft line 1\nleft line 2" + # The list item belongs to the left column, hence it is read before the right column + assert parts[2] == "- \\- item one" # list item, with the raw OCR dash escaped + assert parts[3] == "right line 0\nright line 1\nright line 2" + assert parts[4] == "Page 3 of 12" + # Page furniture can be dropped + assert "Page 3 of 12" not in page.export_as_markdown(include_furniture=False) + # AsciiDoc follows the same structure + parts = page.export_as_asciidoc().split("\n\n") + assert parts[0] == "== A Two Column Study" + assert parts[2] == "* {empty}- item one" + assert "Page 3 of 12" not in page.export_as_asciidoc(include_furniture=False) + # Markdown structural characters are escaped by default + page = _page([_line_at("*bold* #tag [link]", 0.1, 0.1, 0.9, 0.15)]) + assert page.export_as_markdown() == "\\*bold\\* \\#tag \\[link\\]" + assert page.export_as_markdown(escape=False) == "*bold* #tag [link]" + # Empty pages export to an empty string + assert elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [], 0, (1000, 800)).export_as_markdown() == "" -def test_page_export_as_markdown_rotated_page(): - height, width = 1000, 800 - - def _rot_line(text, x0, y0, x1, y1, deg): - angle = np.deg2rad(deg) - rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) - center = np.array([width / 2, height / 2]) - tokens = text.split() - step = (x1 - x0) / len(tokens) - words = [] - for idx, token in enumerate(tokens): - pts = np.array([ - [(x0 + idx * step) * width, y0 * height], - [(x0 + (idx + 0.9) * step) * width, y0 * height], - [(x0 + (idx + 0.9) * step) * width, y1 * height], - [(x0 + idx * step) * width, y1 * height], - ]) - pts = ((pts - center) @ rot.T + center) / [width, height] - words.append( - elements.Word(token, 0.9, tuple(tuple(pt) for pt in pts), 0.9, {"value": 0, "confidence": None}) - ) - return elements.Line(words) +@pytest.mark.parametrize( + "line, expected", + [ + ("3. not a list", "3\\. not a list"), + ("12) neither", "12\\) neither"), + (" 7.", " 7\\."), + ("3", "3"), + ("2022 was a year", "2022 was a year"), + ("1.5 million", "1.5 million"), + ("- dash", "\\- dash"), + ("# hash", "\\# hash"), + ], +) +def test_markdown_finalize_line(line, expected): + assert MarkdownExporter().finalize_line(line) == expected + + +def test_page_export_list_items(): + # Three separate single-line list items, each covered by its own List-item region -> three bullets + lines = [_line_at(f"item number {idx}", 0.1, 0.1 + 0.1 * idx, 0.5, 0.13 + 0.1 * idx) for idx in range(3)] layout = [ - ("big page title", 0.1, 0.05, 0.9, 0.09), - ("left one", 0.1, 0.15, 0.45, 0.19), - ("left two", 0.1, 0.21, 0.45, 0.25), - ("left three", 0.1, 0.27, 0.45, 0.31), - ("right one", 0.55, 0.15, 0.9, 0.19), - ("right two", 0.55, 0.21, 0.9, 0.25), - ("right three", 0.55, 0.27, 0.9, 0.31), + elements.LayoutElement("List-item", 0.9, ((0.08, 0.09 + 0.1 * idx), (0.52, 0.14 + 0.1 * idx))) + for idx in range(3) ] - expected = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_rot_line(*args, 0) for args in layout])], - 0, - (height, width), - ).export_as_markdown() - assert expected.split()[:3] == ["big", "page", "title"] - for deg in (15, 25): - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_rot_line(*args, deg) for args in layout])], - 0, - (height, width), - ) - assert page.export_as_markdown().split() == expected.split() - - -def test_page_export_as_markdown_rotated_landscape_page(): - height, width = 800, 1200 - - def _rot_line(text, x0, y0, x1, y1, deg): - angle = np.deg2rad(deg) - rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) - center = np.array([width / 2, height / 2]) - tokens = text.split() - step = (x1 - x0) / len(tokens) - words = [] - for idx, token in enumerate(tokens): - pts = np.array([ - [(x0 + idx * step) * width, y0 * height], - [(x0 + (idx + 0.9) * step) * width, y0 * height], - [(x0 + (idx + 0.9) * step) * width, y1 * height], - [(x0 + idx * step) * width, y1 * height], - ]) - pts = ((pts - center) @ rot.T + center) / [width, height] - words.append( - elements.Word(token, 0.9, tuple(tuple(pt) for pt in pts), 0.9, {"value": 0, "confidence": None}) - ) - return elements.Line(words) + page = _page(lines, layout=layout) + assert page.export_as_markdown() == "- item number 0\n- item number 1\n- item number 2" + assert page.export_as_asciidoc() == "* item number 0\n* item number 1\n* item number 2" + # each list item is its own block in reading order + items = page.items_in_reading_order() + assert len(items) == 3 + assert all(len(item.lines) == 1 for item in items) + # list items render as one
  • per item in HTML + lines = [_line_at(f"item {idx} ", 0.1, 0.1 + 0.1 * idx, 0.5, 0.13 + 0.1 * idx) for idx in range(2)] + page = _page(lines, layout=layout[:2]) + assert page.export_as_html() == "" + # A single list item wrapped over three visual lines (one region) must render as ONE bullet, while a + # second item (another region) is a second bullet. + lines = [ + _line_at("first item wrapping over", 0.1, 0.10, 0.9, 0.13), + _line_at("several visual lines here", 0.1, 0.14, 0.9, 0.17), + _line_at("until it finally ends", 0.1, 0.18, 0.6, 0.21), + _line_at("second short item", 0.1, 0.26, 0.5, 0.29), + ] layout = [ - ("big page title", 0.1, 0.05, 0.9, 0.09), - ("left one", 0.1, 0.15, 0.45, 0.19), - ("left two", 0.1, 0.21, 0.45, 0.25), - ("right one", 0.55, 0.15, 0.9, 0.19), - ("right two", 0.55, 0.21, 0.9, 0.25), + elements.LayoutElement("List-item", 0.9, ((0.08, 0.09), (0.92, 0.22))), + elements.LayoutElement("List-item", 0.9, ((0.08, 0.25), (0.52, 0.30))), ] - expected = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_rot_line(*args, 0) for args in layout])], - 0, - (height, width), - ).export_as_markdown() - for deg in (-35, 35): - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_rot_line(*args, deg) for args in layout])], - 0, - (height, width), - ) - assert page.export_as_markdown().split() == expected.split() - - -def test_exporter_classes_direct_use(): - # The exporter classes are usable directly, and export_document dispatches per page type - page = _reading_order_page() - md = MarkdownExporter() - adoc = AsciiDocExporter() - assert md.export_page(page) == page.export_as_markdown() - assert adoc.export_page(page) == page.export_as_asciidoc() - - class _Doc: - pages = [page, page] - - assert md.export_document(_Doc()) == "\n\n---\n\n".join([page.export_as_markdown()] * 2) - assert adoc.export_document(_Doc(), page_break="\n\n") == "\n\n".join([page.export_as_asciidoc()] * 2) - # page_reading_order returns (items, labels, direction) - items, labels, direction = page_reading_order(page) - assert len(items) == len(labels) and direction == "ltr" + page = _page(lines, layout=layout) + assert page.export_as_markdown() == ( + "- first item wrapping over several visual lines here until it finally ends\n- second short item" + ) + items = page.items_in_reading_order() + assert len(items) == 2 + assert len(items[0].lines) == 3 and len(items[1].lines) == 1 def test_page_export_as_html(): @@ -478,22 +498,29 @@ def test_page_export_as_html(): assert html.startswith("

    ") assert page.export_as("html") == html assert HTMLExporter().export_page(page) == html - # list items render as one
  • per item - lines = [_line_at(f"item {idx} ", 0.1, 0.1 + 0.1 * idx, 0.5, 0.13 + 0.1 * idx) for idx in range(2)] - layout = [ - elements.LayoutElement("List-item", 0.9, ((0.08, 0.09 + 0.1 * idx), (0.52, 0.14 + 0.1 * idx))) - for idx in range(2) - ] - lp = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), layout=layout - ) - assert lp.export_as_html() == "
      \n
    • item 0 <x>
      • \n
    • item 1 <x>
      • \n
    " + # Page furniture can be dropped from the HTML export too + assert "Page 3 of 12" in html + assert "Page 3 of 12" not in page.export_as_html(include_furniture=False) + # Words recognized as empty strings must not produce an empty paragraph / list item + blank = elements.Block(lines=[elements.Line([_word_at("", 0.1, 0.1, 0.4, 0.15)])]) + text = elements.Block(lines=[_line_at("real text", 0.1, 0.5, 0.9, 0.55)]) + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [blank, text], 0, (1000, 800)) + assert page.render() == "real text" + assert page.export_as_markdown() == "real text" + assert page.export_as_html() == "

    real text

    " -def test_export_mixins_carry_full_api(): +def test_export_as_dispatch(): + page = _reading_order_page() + assert page.export_as("markdown") == page.export_as("md") == page.export_as_markdown() + assert page.export_as("adoc") == page.export_as_asciidoc() + assert page.export_as("text") == page.render() + assert page.export_as("json") == page.export() + assert isinstance(page.export_as("xml")[0], bytes) + assert page.export_as("markdown", include_furniture=False) == page.export_as_markdown(include_furniture=False) + with pytest.raises(ValueError): + page.export_as("yaml") # The element export surface comes from the mixins in doctr.io.exporters, with the API unchanged - from doctr.io.exporters import DocumentExportsMixin, KIEPageExportsMixin, PageExportsMixin - for method in ( "render", "export_as_xml", @@ -506,308 +533,75 @@ def test_export_mixins_carry_full_api(): assert getattr(elements.KIEPage, method) is getattr(KIEPageExportsMixin, method) assert getattr(elements.Document, method) is getattr(DocumentExportsMixin, method) assert elements.Page.items_in_reading_order is PageExportsMixin.items_in_reading_order - page = _reading_order_page() - # dispatcher covers every format - for fmt in ("markdown", "md", "asciidoc", "adoc", "html", "text", "txt", "json", "dict", "xml", "hocr"): - page.export_as(fmt) - with pytest.raises(ValueError): - page.export_as("pptx") - - -def test_page_render(): - left = elements.Block( - lines=[_line_at("left top", 0.08, 0.1, 0.45, 0.13), _line_at("left low", 0.08, 0.2, 0.45, 0.23)] - ) - right = elements.Block( - lines=[_line_at("right top", 0.55, 0.1, 0.92, 0.13), _line_at("right low", 0.55, 0.2, 0.92, 0.23)] - ) - # The blocks are stored right-first, but render() linearizes them like the other exporters - page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [right, left], 0, (1000, 800)) - assert page.render() == "left top\n\nleft low\n\nright top\n\nright low" - assert page.render(block_break=" | ") == "left top | left low | right top | right low" - # Single-block and empty pages - single = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [left], 0, (1000, 800)) - assert single.render(block_break=" | ") == "left top | left low" - assert elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [], 0, (1000, 800)).render() == "" - - -def test_page_render_includes_tables_and_furniture_flag(): - page = _reading_order_page() - assert page.render().splitlines()[0] == "A Two Column Study" - assert "Page 3 of 12" in page.render() - # Page furniture can be dropped, exactly like in the Markdown / HTML exports - assert "Page 3 of 12" not in page.render(include_furniture=False) - # Recognized tables are part of the plain text render - cells = [ - elements.TableCell("head", 0.9, ((0.1, 0.6), (0.4, 0.65)), 0, 0, 0, 0), - elements.TableCell("body", 0.9, ((0.1, 0.66), (0.4, 0.71)), 1, 1, 0, 0), - ] - table = elements.Table(cells=cells, num_rows=2, num_cols=1, geometry=((0.1, 0.6), (0.4, 0.71))) - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("intro line", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), - tables=[table], - ) - assert page.render() == "intro line\n\nhead\nbody" - - -def test_export_is_json_serializable(): - import json - - # Straight boxes built from a detection array carry np.float32 coordinates - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [ - elements.Block( - lines=[ - elements.Line([ - elements.Word( - "np", - np.float32(0.9), - ((np.float32(0.1), np.float32(0.1)), (np.float32(0.4), np.float32(0.2))), - np.float32(0.8), - {"value": np.int64(0), "confidence": None}, - ) - ]) - ] - ) - ], - np.int64(0), - (np.int64(1000), np.int64(800)), - {"value": np.float32(0.0), "confidence": np.float32(1.0)}, - ) - exported = page.export() - assert json.loads(json.dumps(exported)) is not None - word = exported["blocks"][0]["lines"][0]["words"][0] - assert [coord for point in word["geometry"] for coord in point] == pytest.approx([0.1, 0.1, 0.4, 0.2]) - assert all(type(coord) is float for point in word["geometry"] for coord in point) - assert type(word["confidence"]) is float and type(word["objectness_score"]) is float - assert type(word["crop_orientation"]["value"]) is int - assert exported["page_idx"] == 0 and type(exported["page_idx"]) is int - assert exported["dimensions"] == (1000, 800) and all(type(dim) is int for dim in exported["dimensions"]) - # Rotated lines / blocks expose their geometry as a numpy array, exported as nested tuples - poly = np.asarray([[0.1, 0.1], [0.4, 0.1], [0.4, 0.2], [0.1, 0.2]], dtype=np.float32) - rot_word = elements.Word("rot", 0.9, poly, 0.8, {"value": 0, "confidence": None}) - rot_page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[elements.Line([rot_word])])], - 0, - (1000, 800), - ) - assert isinstance(rot_page.blocks[0].geometry, np.ndarray) # untouched in memory - rot_export = rot_page.export() - assert isinstance(rot_export["blocks"][0]["geometry"], tuple) - json.dumps(rot_export) - # ... and a whole document round-trips through JSON - doc = elements.Document([page, rot_page]) - assert elements.Document.from_dict(json.loads(json.dumps(doc.export()))) is not None - - -def test_export(): - page = _reading_order_page() - rendered = [ - " ".join(word["value"] for line in block["lines"] for word in line["words"]) - for block in page.export()["blocks"] - ] - assert rendered[0] == "A Two Column Study" - assert rendered[-1] == "Page 3 of 12" - assert rendered.index("left line 0 left line 1 left line 2") < rendered.index( - "right line 0 right line 1 right line 2" - ) - # The linearization can be opted out of - raw = page.export(reading_order=False)["blocks"] - assert len(raw) == 1 and len(raw[0]["lines"]) == 9 - # Artefacts survive the regrouping - artefact = elements.Artefact("qr_code", 0.9, ((0.1, 0.8), (0.2, 0.9))) - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("a line", 0.1, 0.1, 0.9, 0.15)], artefacts=[artefact])], - 0, - (1000, 800), - ) - assert page.export()["blocks"][0]["artefacts"] == [artefact.export()] - - -def test_xml_export(): - import re - - page = _reading_order_page() - xml = page.export_as_xml()[0].decode() - words = re.findall(r'class="ocrx_word"[^>]*>([^<]*)<', xml) - assert words[:4] == ["A", "Two", "Column", "Study"] - assert words[-4:] == ["Page", "3", "of", "12"] - assert words.index("left") < words.index("right") - # Recognized tables are serialized as an hOCR text area (they used to be dropped) - cells = [ - elements.TableCell("head", 0.9, ((0.1, 0.6), (0.4, 0.65)), 0, 0, 0, 0), - elements.TableCell("body", 0.9, ((0.1, 0.66), (0.4, 0.71)), 1, 1, 0, 0), - ] - table = elements.Table(cells=cells, num_rows=2, num_cols=1, geometry=((0.1, 0.6), (0.4, 0.71))) - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("intro line", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), - tables=[table], - ) - xml = page.export_as_xml()[0].decode() - assert 'id="table_1"' in xml and ">head<" in xml and ">body<" in xml - - -def test_kie_page_export(): - predictions = { - CLASS_NAME: [ - elements.Prediction("second", 0.9, ((0.1, 0.5), (0.9, 0.6)), 0.9, {"value": 0, "confidence": None}), - elements.Prediction("first", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None}), - ] - } - page = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - assert [p["value"] for p in page.export()["predictions"][CLASS_NAME]] == ["first", "second"] - assert [p["value"] for p in page.export(reading_order=False)["predictions"][CLASS_NAME]] == ["second", "first"] - xml = page.export_as_xml()[0].decode() - assert xml.index(">first<") < xml.index(">second<") - - -def test_kie_page_render_keeps_reading_order(): - predictions = { - CLASS_NAME: [ - elements.Prediction("second", 0.9, ((0.1, 0.5), (0.9, 0.6)), 0.9, {"value": 0, "confidence": None}), - elements.Prediction("first", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None}), - ] - } - page = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - assert page.render() == f"{CLASS_NAME}: first\n\n{CLASS_NAME}: second" + # The dispatcher covers every format, for every element + page = _page([_line_at("page text", 0.1, 0.1, 0.9, 0.15)]) + kie_page = _kie_page({CLASS_NAME: [_prediction("value", 0.1, 0.2)]}) + formats = ("markdown", "md", "asciidoc", "adoc", "html", "text", "txt", "json", "dict", "xml", "hocr") + for element in (page, kie_page, elements.Document([page]), elements.KIEDocument([kie_page])): + for fmt in formats: + assert element.export_as(fmt) is not None + with pytest.raises(ValueError, match="unsupported export format"): + element.export_as("pptx") + # The document-level named methods carry the page break of their format + doc = elements.Document([page, page]) + assert doc.export_as_html() == "

    page text

    page text

    " + assert doc.export_as_html(page_break="\n") == "

    page text

    \n

    page text

    " + assert doc.export_as_markdown() == "page text\n\n---\n\npage text" + assert doc.export_as_asciidoc() == "page text\n\n<<<\n\npage text" + assert doc.export_as_markdown(page_break="\n\n") == "page text\n\npage text" + assert doc.export_as("markdown") == doc.export_as_markdown() + assert doc.export_as("text") == doc.render() + assert doc.export_as("json") == doc.export() + assert len(doc.export_as("xml")) == 2 and len(doc.export_as_xml()) == 2 + with pytest.raises(ValueError): + doc.export_as("pdf") + kie_doc = elements.KIEDocument([kie_page, kie_page]) + kie_html = f"

    {CLASS_NAME}

    \n
      \n
    • value
      • \n
    " + assert kie_doc.export_as_html() == f"{kie_html}
    {kie_html}" + assert kie_page.export_as_html() == kie_html + assert kie_page.export_as_html(direction="rtl") == kie_page.export_as_html() -def test_xml_exporter_class(): - from xml.etree import ElementTree as ET +def test_exporter_classes(): + # The exporter classes are usable directly, and export_document dispatches per page type page = _reading_order_page() - xml_bytes, tree = XMLExporter().export_page(page) - assert isinstance(xml_bytes, bytes) and isinstance(tree, ET.ElementTree) - # The mixin method delegates to the exporter class, so the output is identical - assert page.export_as_xml()[0] == xml_bytes - assert page.export_as("xml")[0] == xml_bytes + md, adoc = MarkdownExporter(), AsciiDocExporter() + assert md.export_page(page) == page.export_as_markdown() + assert adoc.export_page(page) == page.export_as_asciidoc() - predictions = { - CLASS_NAME: [elements.Prediction("hi", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None})] - } - kie = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - assert kie.export_as_xml()[0] == XMLExporter().export_kie_page(kie)[0] + class _Doc: + pages = [page, page] - doc = elements.Document([page, page]) - doc_xml = XMLExporter().export_document(doc) - assert len(doc_xml) == 2 and doc.export_as_xml()[0][0] == doc_xml[0][0] + assert md.export_document(_Doc()) == "\n\n---\n\n".join([page.export_as_markdown()] * 2) + assert adoc.export_document(_Doc(), page_break="\n\n") == "\n\n".join([page.export_as_asciidoc()] * 2) + # page_reading_order returns (items, labels, direction) + items, labels, direction = page_reading_order(page) + assert len(items) == len(labels) and direction == "ltr" + # The base text exporter is abstract + table = elements.Table(cells=[], num_rows=0, num_cols=0, geometry=((0.0, 0.0), (0.0, 0.0))) + with pytest.raises(NotImplementedError): + _PageTextExporter().render_table(table) + with pytest.raises(NotImplementedError): + _PageTextExporter().class_header("cls") -def test_xml_export_text_size(): - # x_size (text height in pixels) and x_fsize (font size in points) are estimated from the - # element box height; the dpi kwarg scales the pixel-to-point conversion (default 72) - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("one two", 0.1, 0.1, 0.9, 0.2)])], - 0, - (1000, 800), - ) - xml = page.export_as_xml()[0].decode() - # the line spans y=0.1..0.2 of a 1000px-high page -> 100px text height, 100pt at 72 dpi - assert "x_size 100; x_fsize 100;" in xml - assert "x_size 0;" not in xml - xml = page.export_as_xml(dpi=144)[0].decode() - assert "x_size 100; x_fsize 50;" in xml - - # table rows and KIE predictions carry the estimated sizes as well - cells = [ - elements.TableCell("head", 0.9, ((0.1, 0.6), (0.4, 0.65)), 0, 0, 0, 0), - elements.TableCell("body", 0.9, ((0.1, 0.66), (0.4, 0.71)), 1, 1, 0, 0), - ] - table = elements.Table(cells=cells, num_rows=2, num_cols=1, geometry=((0.1, 0.6), (0.4, 0.71))) - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("intro line", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), - tables=[table], +def test_table_exports(): + page = _page( + [_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)], + tables=[elements.Table(_order_cells(), 2, 2, ((0.1, 0.55), (0.7, 0.65)), 0.95)], ) - xml = page.export_as_xml()[0].decode() - # row 0 spans y=0.6..0.65 -> 50px, row 1 spans y=0.66..0.71 -> 50px - assert "x_size 50; x_fsize 50;" in xml - - predictions = { - CLASS_NAME: [elements.Prediction("hi", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None})] - } - kie = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - xml = kie.export_as_xml()[0].decode() - assert "x_size 100; x_fsize 100;" in xml - - -def _table_page(num_rows=2, num_cols=1, rotated=False): - """A page with an intro line and one recognized table""" - geo = ((0.1, 0.6), (0.4, 0.66), (0.4, 0.71), (0.1, 0.71)) if rotated else ((0.1, 0.6), (0.4, 0.71)) - cells = [ - elements.TableCell("head", 0.9, ((0.1, 0.6), (0.4, 0.65)), 0, 0, 0, 0), - elements.TableCell("body", 0.9, ((0.1, 0.66), (0.4, 0.71)), 1, 1, 0, 0), + assert page.export_as_markdown().split("\n\n") == [ + "before the table", + "| Name | Qty |\n| --- | --- |\n| Bolt | 12\\|3 |", + "after the table", ] - table = elements.Table(cells=cells, num_rows=num_rows, num_cols=num_cols, geometry=geo) - return elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("intro line", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), - tables=[table], - ) - - -def test_to_json_safe_covers_every_container(): - # 0-d arrays and numpy scalars collapse to their Python equivalent - assert to_json_safe(np.array(3.5)) == 3.5 and type(to_json_safe(np.array(3.5))) is float - assert to_json_safe(np.int64(7)) == 7 and type(to_json_safe(np.int64(7))) is int - assert to_json_safe(np.bool_(True)) is True - # arrays become nested tuples, so `create_obj_patch` still recognizes an exported geometry - assert to_json_safe(np.asarray([[0.0, 1.0], [2.0, 3.0]])) == ((0.0, 1.0), (2.0, 3.0)) - # lists stay lists, sets are normalized to lists, dict keys to str - assert to_json_safe([np.float32(1.0), (np.int8(2),)]) == [1.0, (2,)] - assert sorted(to_json_safe({np.int64(1), np.int64(2)})) == [1, 2] - assert to_json_safe({np.int64(1): np.float32(0.5)}) == {"1": 0.5} - # anything already built-in is returned untouched - assert to_json_safe("text") == "text" and to_json_safe(None) is None - json.dumps(to_json_safe({"a": [np.float32(0.1)], "b": np.arange(3)})) - - -def test_page_artefacts(): - artefact = elements.Artefact("qr_code", 0.9, ((0.7, 0.7), (0.9, 0.9))) - orphan = elements.Block(lines=[], artefacts=[artefact], geometry=((0.7, 0.7), (0.9, 0.9)), objectness_score=0.9) - text = elements.Block(lines=[_line_at("some text", 0.1, 0.1, 0.9, 0.15)]) - page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [orphan, text], 0, (1000, 800)) - items = page.items_in_reading_order() - assert [artefact.export() for artefact in items[-1].artefacts] == [artefact.export()] - assert page.export()["blocks"][-1]["artefacts"] == [artefact.export()] - - -def test_ordered_line_words_directions(): - line = _line_at("one two three", 0.1, 0.1, 0.7, 0.15) - assert [word.render() for word in ordered_line_words(line, "ltr")] == ["one", "two", "three"] - assert [word.render() for word in ordered_line_words(line, "rtl")] == ["three", "two", "one"] - # Vertical pages read their words top to bottom - stacked = elements.Line([_word_at(value, 0.1, 0.1 * idx, 0.2, 0.1 * idx + 0.05) for idx, value in enumerate("abc")]) - assert [word.render() for word in ordered_line_words(stacked, "ttb-rtl")] == ["a", "b", "c"] - assert [word.render() for word in ordered_line_words(stacked, "ttb-ltr")] == ["a", "b", "c"] - # A single-word line short-circuits the per-line detection - single = elements.Line([_word_at("solo", 0.1, 0.1, 0.2, 0.15)]) - assert [word.render() for word in ordered_line_words(single, "auto", auto=True)] == ["solo"] - - -def test_vertical_direction_reaches_every_exporter(): - page = _reading_order_page() - for export in (page.render, page.export_as_markdown, page.export_as_asciidoc, page.export_as_html): - assert isinstance(export(direction="ttb-ltr"), str) - assert isinstance(page.export_as_xml(direction="ttb-rtl")[0], bytes) - - -def test_html_export_with_tables_and_furniture(): + 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") + # Recognized tables are part of the plain text render and of the HTML export page = _table_page() + assert page.render() == "intro line\n\nhead\nbody" html = page.export_as_html() assert "" in html and "" in html and "" in html assert html.index("

    intro line

    ") < html.index("
    headbody
    ") @@ -819,183 +613,126 @@ def test_html_export_with_tables_and_furniture(): geometry=((0.1, 0.6), (0.4, 0.65)), ) assert HTMLExporter().render_table(single) == "
    \n\n
    only
    " - # Page furniture can be dropped from the HTML export too - assert "Page 3 of 12" in _reading_order_page().export_as_html() - assert "Page 3 of 12" not in _reading_order_page().export_as_html(include_furniture=False) - - -def test_render_table_with_empty_grid(): + # An empty table contributes nothing to the page export empty = elements.Table(cells=[], num_rows=0, num_cols=0, geometry=((0.0, 0.0), (0.0, 0.0))) for exporter in (MarkdownExporter(), AsciiDocExporter(), HTMLExporter(), TextExporter()): assert exporter.render_table(empty) == "" - # An empty table contributes nothing to the page export - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("only text", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), - tables=[empty], - ) + page = _page([_line_at("only text", 0.1, 0.1, 0.9, 0.15)], tables=[empty]) assert page.export_as_markdown() == "only text" assert page.export_as_html() == "

    only text

    " -def test_text_exporter_kie_and_document(): - predictions = { - CLASS_NAME: [ - elements.Prediction("second", 0.9, ((0.1, 0.5), (0.9, 0.6)), 0.9, {"value": 0, "confidence": None}), - elements.Prediction("first", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None}), - ], - "empty": [], - } - kie = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - # Classes without any prediction are skipped by every text exporter - assert TextExporter().export_kie_page(kie) == f"{CLASS_NAME}:\n\nfirst\nsecond" - assert "empty" not in MarkdownExporter().export_kie_page(kie) - assert "empty" not in AsciiDocExporter().export_kie_page(kie) - assert "empty" not in HTMLExporter().export_kie_page(kie) - # export_document dispatches on the page type and uses the format-specific page break - page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("a page", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), - ) - mixed = elements.Document([page, kie]) - assert TextExporter().export_document(mixed) == f"a page\n\n\n\n{CLASS_NAME}:\n\nfirst\nsecond" - assert TextExporter().export_document(mixed, page_break=" // ").startswith("a page // ") +# --- XML / JSON ------------------------------------------------------------------------------------------------------ -def test_base_text_exporter_is_abstract(): - from doctr.io.exporters import _PageTextExporter - - table = elements.Table(cells=[], num_rows=0, num_cols=0, geometry=((0.0, 0.0), (0.0, 0.0))) - with pytest.raises(NotImplementedError): - _PageTextExporter().render_table(table) - with pytest.raises(NotImplementedError): - _PageTextExporter().class_header("cls") - - -def test_blank_lines_are_skipped(): - # Words recognized as empty strings must not produce an empty paragraph / list item - blank = elements.Block(lines=[elements.Line([_word_at("", 0.1, 0.1, 0.4, 0.15)])]) - text = elements.Block(lines=[_line_at("real text", 0.1, 0.5, 0.9, 0.55)]) - page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [blank, text], 0, (1000, 800)) - assert page.render() == "real text" - assert page.export_as_markdown() == "real text" - assert page.export_as_html() == "

    real text

    " +def test_xml_export(): + page = _reading_order_page() + xml = page.export_as_xml()[0].decode() + words = re.findall(r'class="ocrx_word"[^>]*>([^<]*)<', xml) + assert words[:4] == ["A", "Two", "Column", "Study"] + assert words[-4:] == ["Page", "3", "of", "12"] + assert words.index("left") < words.index("right") + # The mixin method delegates to the exporter class, so the output is identical + xml_bytes, tree = XMLExporter().export_page(page) + assert isinstance(xml_bytes, bytes) and isinstance(tree, ET.ElementTree) + assert page.export_as_xml()[0] == xml_bytes + assert page.export_as("xml")[0] == xml_bytes + kie = _kie_page({CLASS_NAME: [_prediction("hi", 0.1, 0.2)]}) + assert kie.export_as_xml()[0] == XMLExporter().export_kie_page(kie)[0] + doc = elements.Document([page, page]) + doc_xml = XMLExporter().export_document(doc) + assert len(doc_xml) == 2 and doc.export_as_xml()[0][0] == doc_xml[0][0] + # Recognized tables are serialized as an hOCR text area (they used to be dropped) + xml = _table_page().export_as_xml()[0].decode() + assert 'id="table_1"' in xml and ">head<" in xml and ">body<" in xml + # The linearization can be opted out of: blocks first, then tables, in their stored order + xml = _table_page().export_as_xml(reading_order=False)[0].decode() + assert xml.index("intro") < xml.index(">head<") -def test_xml_export_edge_cases(): - from doctr.io.exporters import _covering_region_indices + # x_size (text height in pixels) and x_fsize (font size in points) are estimated from the + # element box height; the dpi kwarg scales the pixel-to-point conversion (default 72) + page = _page([_line_at("one two", 0.1, 0.1, 0.9, 0.2)]) + xml = page.export_as_xml()[0].decode() + # the line spans y=0.1..0.2 of a 1000px-high page -> 100px text height, 100pt at 72 dpi + assert "x_size 100; x_fsize 100;" in xml + assert "x_size 0;" not in xml + assert "x_size 100; x_fsize 50;" in page.export_as_xml(dpi=144)[0].decode() + # table rows and KIE predictions carry the estimated sizes as well + # row 0 spans y=0.6..0.65 -> 50px, row 1 spans y=0.66..0.71 -> 50px + assert "x_size 50; x_fsize 50;" in _table_page().export_as_xml()[0].decode() + assert "x_size 100; x_fsize 100;" in kie.export_as_xml()[0].decode() assert _covering_region_indices([], [((0.0, 0.0), (1.0, 1.0))]) == [] # Rotated geometries are rejected with an explicit error, for blocks, tables and KIE predictions alike poly = np.asarray([[0.1, 0.1], [0.4, 0.1], [0.4, 0.2], [0.1, 0.2]], dtype=np.float32) - rot_page = elements.Page( - np.zeros((10, 10, 3), dtype=np.uint8), - [ - elements.Block( - lines=[elements.Line([elements.Word("rot", 0.9, poly, 0.9, {"value": 0, "confidence": None})])] - ) - ], - 0, - (1000, 800), - ) with pytest.raises(TypeError): - rot_page.export_as_xml() + _page([elements.Line([elements.Word("rot", 0.9, poly, 0.9, _ORIENTATION)])]).export_as_xml() with pytest.raises(TypeError): XMLExporter().export_page(_table_page(rotated=True)) - rot_kie = elements.KIEPage( - np.zeros((10, 10, 3), dtype=np.uint8), - {CLASS_NAME: [elements.Prediction("rot", 0.9, poly, 0.9, {"value": 0, "confidence": None})]}, - 0, - (1000, 800), - ) with pytest.raises(TypeError): - rot_kie.export_as_xml() - # The linearization can be opted out of: blocks first, then tables, in their stored order - xml = _table_page().export_as_xml(reading_order=False)[0].decode() - assert xml.index("intro") < xml.index(">head<") - - -def test_document_forwards_reading_order_options(): - left = elements.Block(lines=[_line_at("left col", 0.08, 0.2, 0.45, 0.23)]) - right = elements.Block(lines=[_line_at("right col", 0.55, 0.2, 0.92, 0.23)]) - footer = elements.Block(lines=[_line_at("Page 1 of 2", 0.4, 0.95, 0.6, 0.97)]) - layout = [ - elements.LayoutElement("Text", 0.98, ((0.06, 0.18), (0.48, 0.26))), - elements.LayoutElement("Text", 0.98, ((0.52, 0.18), (0.94, 0.26))), - elements.LayoutElement("Page-footer", 0.97, ((0.35, 0.94), (0.65, 0.98))), - ] - page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [right, left, footer], 0, (1000, 800), layout=layout) - doc = elements.Document([page, page]) - - assert ( - doc.render(page_break=" || ") == "left col\n\nright col\n\nPage 1 of 2 || left col\n\nright col\n\nPage 1 of 2" - ) - assert "Page 1 of 2" not in doc.render(include_furniture=False) - assert doc.render(block_break=" ", page_break=" || ").startswith("left col right col") - assert doc.export_as("text", include_furniture=False) == doc.render(include_furniture=False) - - ordered = doc.export()["pages"][0]["blocks"] - assert [block["lines"][0]["words"][0]["value"] for block in ordered] == ["left", "right", "Page"] - raw = doc.export(reading_order=False)["pages"][0]["blocks"] - assert [block["lines"][0]["words"][0]["value"] for block in raw] == ["right", "left", "Page"] - assert doc.export_as("json", reading_order=False) == doc.export(reading_order=False) - json.dumps(doc.export(reading_order=False)) - - -def test_kie_document_forwards_reading_order_options(): - predictions = { - CLASS_NAME: [ - elements.Prediction("second", 0.9, ((0.1, 0.5), (0.9, 0.6)), 0.9, {"value": 0, "confidence": None}), - elements.Prediction("first", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None}), - ] - } - page = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - doc = elements.KIEDocument([page]) - assert doc.render() == f"{CLASS_NAME}: first\n\n{CLASS_NAME}: second" - # The reading direction reaches KIEPage.render through the document - assert doc.render(direction="ttb-ltr") == f"{CLASS_NAME}: first\n\n{CLASS_NAME}: second" - assert [p["value"] for p in doc.export()["pages"][0]["predictions"][CLASS_NAME]] == ["first", "second"] - raw = doc.export(reading_order=False)["pages"][0]["predictions"][CLASS_NAME] - assert [p["value"] for p in raw] == ["second", "first"] - json.dumps(doc.export(reading_order=False)) + _kie_page({CLASS_NAME: [elements.Prediction("rot", 0.9, poly, 0.9, _ORIENTATION)]}).export_as_xml() -def test_every_format_is_reachable_from_a_document(): +def test_json_export(): + # Straight boxes built from a detection array carry np.float32 coordinates page = elements.Page( np.zeros((10, 10, 3), dtype=np.uint8), - [elements.Block(lines=[_line_at("page text", 0.1, 0.1, 0.9, 0.15)])], - 0, - (1000, 800), + [ + elements.Block( + lines=[ + elements.Line([ + elements.Word( + "np", + np.float32(0.9), + ((np.float32(0.1), np.float32(0.1)), (np.float32(0.4), np.float32(0.2))), + np.float32(0.8), + {"value": np.int64(0), "confidence": None}, + ) + ]) + ] + ) + ], + np.int64(0), + (np.int64(1000), np.int64(800)), + {"value": np.float32(0.0), "confidence": np.float32(1.0)}, ) - predictions = { - CLASS_NAME: [elements.Prediction("value", 0.9, ((0.1, 0.1), (0.9, 0.2)), 0.9, {"value": 0, "confidence": None})] - } - kie_page = elements.KIEPage(np.zeros((10, 10, 3), dtype=np.uint8), predictions, 0, (1000, 800)) - formats = ("markdown", "md", "asciidoc", "adoc", "html", "text", "txt", "json", "dict", "xml", "hocr") + exported = page.export() + assert json.loads(json.dumps(exported)) is not None + word = exported["blocks"][0]["lines"][0]["words"][0] + assert [coord for point in word["geometry"] for coord in point] == pytest.approx([0.1, 0.1, 0.4, 0.2]) + assert all(type(coord) is float for point in word["geometry"] for coord in point) + assert type(word["confidence"]) is float and type(word["objectness_score"]) is float + assert type(word["crop_orientation"]["value"]) is int + assert exported["page_idx"] == 0 and type(exported["page_idx"]) is int + assert exported["dimensions"] == (1000, 800) and all(type(dim) is int for dim in exported["dimensions"]) + # Rotated lines / blocks expose their geometry as a numpy array, exported as nested tuples + poly = np.asarray([[0.1, 0.1], [0.4, 0.1], [0.4, 0.2], [0.1, 0.2]], dtype=np.float32) + rot_page = _page([elements.Line([elements.Word("rot", 0.9, poly, 0.8, _ORIENTATION)])]) + assert isinstance(rot_page.blocks[0].geometry, np.ndarray) # untouched in memory + rot_export = rot_page.export() + assert isinstance(rot_export["blocks"][0]["geometry"], tuple) + json.dumps(rot_export) + # ... and a whole document round-trips through JSON + doc = elements.Document([page, rot_page]) + assert elements.Document.from_dict(json.loads(json.dumps(doc.export()))) is not None - for element in (page, kie_page, elements.Document([page]), elements.KIEDocument([kie_page])): - for fmt in formats: - assert element.export_as(fmt) is not None - with pytest.raises(ValueError, match="unsupported export format"): - element.export_as("pptx") + # 0-d arrays and numpy scalars collapse to their Python equivalent + assert to_json_safe(np.array(3.5)) == 3.5 and type(to_json_safe(np.array(3.5))) is float + assert to_json_safe(np.int64(7)) == 7 and type(to_json_safe(np.int64(7))) is int + assert to_json_safe(np.bool_(True)) is True + # arrays become nested tuples, so `create_obj_patch` still recognizes an exported geometry + assert to_json_safe(np.asarray([[0.0, 1.0], [2.0, 3.0]])) == ((0.0, 1.0), (2.0, 3.0)) + # lists stay lists, sets are normalized to lists, dict keys to str + assert to_json_safe([np.float32(1.0), (np.int8(2),)]) == [1.0, (2,)] + assert sorted(to_json_safe({np.int64(1), np.int64(2)})) == [1, 2] + assert to_json_safe({np.int64(1): np.float32(0.5)}) == {"1": 0.5} + # anything already built-in is returned untouched + assert to_json_safe("text") == "text" and to_json_safe(None) is None + json.dumps(to_json_safe({"a": [np.float32(0.1)], "b": np.arange(3)})) - # The document-level named methods carry the page break of their format - doc = elements.Document([page, page]) - assert doc.export_as_html() == "

    page text

    page text

    " - assert doc.export_as_html(page_break="\n") == "

    page text

    \n

    page text

    " - assert doc.export_as_markdown().count("\n\n---\n\n") == 1 - assert doc.export_as_asciidoc().count("\n\n<<<\n\n") == 1 - assert len(doc.export_as_xml()) == 2 - kie_doc = elements.KIEDocument([kie_page, kie_page]) - kie_html = f"

    {CLASS_NAME}

    \n
      \n
    • value
      • \n
    " - assert kie_doc.export_as_html() == f"{kie_html}
    {kie_html}" - assert kie_page.export_as_html() == f"

    {CLASS_NAME}

    \n
      \n
    • value
      • \n
    " - assert kie_page.export_as_html(direction="rtl") == kie_page.export_as_html() +# --- Figures --------------------------------------------------------------------------------------------------------- def _figure_page(caption="Figure 1 quarterly revenue", picture=((0.12, 0.30), (0.88, 0.60))): @@ -1017,7 +754,7 @@ def _figure_page(caption="Figure 1 quarterly revenue", picture=((0.12, 0.30), (0 if caption is not None: lines.append(_line_at(caption, 0.2, 0.63, 0.8, 0.66)) layout.append(elements.LayoutElement("Caption", 0.96, ((0.18, 0.62), (0.82, 0.67)))) - return elements.Page(image, [elements.Block(lines=lines)], 0, (1000, 800), layout=layout) + return _page(lines, image=image, layout=layout) def test_page_reading_order_with_figures(): @@ -1043,7 +780,7 @@ def test_page_reading_order_with_figures(): assert all(isinstance(item, elements.Block) for item in plain.items_in_reading_order(include_figures=True)) -def test_page_export_markdown_images(): +def test_page_export_images(tmp_path): page = _figure_page() # Placeholder by default, the text is kept markdown = page.export_as_markdown() @@ -1051,34 +788,69 @@ def test_page_export_markdown_images(): assert markdown.index("Revenue grew steadily") < markdown.index("") assert markdown.index("") < markdown.index("Figure 1 quarterly revenue") assert "axis label" in markdown - - assert "" not in page.export_as_markdown(images="none") - assert "Figure 1 quarterly revenue" in page.export_as_markdown(images="none") + assert "// image" in page.export_as_asciidoc() + assert "" in page.export_as_html() + # No image at all, but the text is kept + for fmt in ("markdown", "asciidoc", "html"): + exported = getattr(page, f"export_as_{fmt}")(images="none") + assert "image" not in exported and "base64" not in exported + assert "axis label" in exported + assert "Figure 1 quarterly revenue" in exported # Embedded: the caption becomes the alt text and the inner text is dropped embedded = page.export_as_markdown(images="embedded") assert "![Figure 1 quarterly revenue](data:image/png;base64," in embedded assert "axis label" not in embedded # The caption is also a visible italic line under the image - assert "](data:image/png;base64," in embedded assert embedded.count("Figure 1 quarterly revenue") == 2 assert re.search( r"!\[Figure 1 quarterly revenue\]\(data:image/png;base64,[^)]+\)\n\n\*Figure 1 quarterly revenue\*", embedded ) 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 + asciidoc = page.export_as_asciidoc(images="embedded") + assert ".Figure 1 quarterly revenue\nimage::data:image/png;base64," in asciidoc + html = page.export_as_html(images="embedded") + assert '
    Figure 1 quarterly revenue
    " in html + assert 'alt="Figure 1 quarterly revenue"' in html + + # Referenced: the crops are written to disk + with pytest.raises(ValueError, match="image_dir"): + page.export_as_markdown(images="referenced") + encoder = FigureEncoder("referenced", image_dir=tmp_path / "page", path_prefix="assets/") + markdown = page.export_as_markdown(images=encoder) + assert re.search(r"!\[Figure 1 quarterly revenue\]\(assets/page1_figure1-[0-9a-f]{8}\.png\)", markdown) + assert len(encoder.written) == 1 and encoder.written[0].name.startswith("page1_figure1-") + crop = cv2.imread(str(encoder.written[0])) + assert crop.shape[0] < page.dimensions[0] and crop.shape[1] < page.dimensions[1] + assert crop[..., 2].mean() > 200 # OpenCV reads BGR + + # Documents pass the image options through to their pages + second = _figure_page() + second.page_idx = 1 + doc = elements.Document([page, second]) + assert doc.export_as_markdown().count("") == 2 + assert doc.export_as_markdown(images="none").count("") == 0 + encoder = FigureEncoder("referenced", image_dir=tmp_path / "doc") + doc.export_as_html(images=encoder) + assert len(encoder.written) == 2 + + # Without a caption, the alt text is empty + page.layout = [region for region in page.layout if region.type != "Caption"] + page._reading_order_cache = None + assert 'alt=""' in page.export_as_html(images="embedded") + + # The inner text is also dropped 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), + page = _page( + lines, + image=np.full((100, 80, 3), 200, dtype=np.uint8), layout=[elements.LayoutElement("Picture", 0.9, ((0.20, 0.20), (0.60, 0.60)))], ) markdown = page.export_as_markdown(images="embedded") @@ -1086,60 +858,112 @@ def test_page_export_markdown_images(): assert "AXIS" not in markdown and "LABEL" not in markdown -def test_page_export_referenced_images(tmp_path): - page = _figure_page() - encoder = FigureEncoder("referenced", image_dir=tmp_path, path_prefix="assets/") - markdown = page.export_as_markdown(images=encoder) - assert re.search(r"!\[Figure 1 quarterly revenue\]\(assets/page1_figure1-[0-9a-f]{8}\.png\)", markdown) - assert len(encoder.written) == 1 and encoder.written[0].name.startswith("page1_figure1-") - crop = cv2.imread(str(encoder.written[0])) - assert crop.shape[0] < page.dimensions[0] and crop.shape[1] < page.dimensions[1] - assert crop[..., 2].mean() > 200 # OpenCV reads BGR +def test_page_export_images_fallbacks(): + # A page restored from JSON has no pixels: placeholders, and the inner text is kept + restored = elements.Page.from_dict(_figure_page().export()) + markdown = restored.export_as_markdown(images="embedded") + assert "" in markdown + assert "axis label" in markdown + assert "Figure 1 quarterly revenue" in markdown + # A figure without pixels falls back to a placeholder and keeps its text and caption + class SkippingEncoder(FigureEncoder): + def source(self, page, region, index): + return None -def test_page_export_asciidoc_and_html_images(): page = _figure_page() - assert "// image" in page.export_as_asciidoc() - assert "" in page.export_as_html() + markdown = page.export_as_markdown(images=SkippingEncoder("embedded")) + assert "" in markdown + assert "axis label" in markdown + assert "Figure 1 quarterly revenue" in markdown + # With its pixels, the inner text is dropped + assert "axis label" not in page.export_as_markdown(images="embedded") - asciidoc = page.export_as_asciidoc(images="embedded") - assert ".Figure 1 quarterly revenue\nimage::data:image/png;base64," in asciidoc + # The plain text and JSON exports ignore the figures + assert "image" not in page.render() + assert page.render().startswith("Annual Report") + assert TextExporter().export_page(page, images="embedded") == page.render() + exported = page.export() + assert len(exported["blocks"]) == 5 + assert len(exported["layout"]) == 5 + json.dumps(exported) - html = page.export_as_html(images="embedded") - assert '
    Figure 1 quarterly revenue
    " in html - assert 'alt="Figure 1 quarterly revenue"' in html - # Without a caption, the alt text is empty - page.layout = [region for region in page.layout if region.type != "Caption"] - page._reading_order_cache = None - assert 'alt=""' in page.export_as_html(images="embedded") +@pytest.mark.parametrize( + "exporter, escape, caption, expected", + [ + (MarkdownExporter, True, 'Fig [1], "a" ', ['![Fig \\[1\\], "a" \\](data:image/png;base64,']), + (MarkdownExporter, False, 'Fig [1], "a" ', ['![Fig \\[1\\], "a" ](data:image/png;base64,']), + ( + AsciiDocExporter, + True, + 'Fig [1], "a" ', + ['.Fig [1], "a" \\\nimage::data:image/png;base64,', '["Fig [1], \\"a\\" "]'], + ), + (AsciiDocExporter, True, "... continued", [".{empty}... continued\nimage::data:image/png;base64,"]), + (AsciiDocExporter, True, None, ["image::data:image/png;base64,", "[]\n\nThe trend continued"]), + ( + HTMLExporter, + True, + 'Fig [1], "a" ', + ['alt="Fig [1], "a" <b>"', '
    Fig [1], "a" <b>
    '], + ), + ], +) +def test_page_export_figure_captions(exporter, escape, caption, expected): + exported = exporter().export_page(_figure_page(caption=caption), escape=escape, images="embedded") + assert all(part in exported for part in expected) + + +def test_page_export_caption_formatting(): + # The caption is a separate paragraph under the image, escaped like the other text + for escape, caption, expected in ( + (True, "Figure 1: *bold* claim", "\n\n*Figure 1: \\*bold\\* claim*"), + (False, "Figure 1: plain", "\n\n*Figure 1: plain*"), + ): + page = _figure_page(caption=caption) + markdown = page.export_as_markdown(images="embedded", escape=escape) + image_end = markdown.index(")", markdown.index("](data:image/png;base64,")) + assert markdown[image_end + 1 :].startswith(expected) + # No caption line without pixels + assert "\n\n*" not in page.export_as_markdown(images="placeholder", escape=escape) + # No caption line without a caption + assert "\n\n*" not in _figure_page(caption=None).export_as_markdown(images="embedded") + + # AsciiDoc: the caption is the block title (without the automatic "Figure N." prefix) and the alt text + page = _figure_page(caption="Figure 1: {author} revenue") + asciidoc = page.export_as_asciidoc(images="embedded") + assert '[caption=""]\n.Figure 1: \\{author\\} revenue\nimage::data:image/png;base64,' in asciidoc + assert '["Figure 1: \\{author\\} revenue"]' in asciidoc + raw = page.export_as_asciidoc(images="embedded", escape=False) + assert '[caption=""]\n.Figure 1: {author} revenue\n' in raw and '["Figure 1: {author} revenue"]' in raw + assert "[caption=" not in _figure_page(caption=None).export_as_asciidoc(images="embedded") + + # A caption hyphenated over two lines is joined back + page = _figure_page(caption=None) + caption_lines = [_line_at("Figure 1: quarterly rev-", 0.2, 0.63, 0.8, 0.66), _line_at("enue", 0.2, 0.67, 0.4, 0.7)] + page.blocks = [*page.blocks, elements.Block(lines=caption_lines)] + page.layout = [*page.layout, elements.LayoutElement("Caption", 0.96, ((0.18, 0.62), (0.82, 0.705)))] + assert "![Figure 1: quarterly revenue](data:image/png;base64," in page.export_as_markdown(images="embedded") + assert "
    Figure 1: quarterly revenue
    " in page.export_as_html(images="embedded") @pytest.mark.parametrize( - "exporter, escape, caption, expected", + "lines, expected", [ - (MarkdownExporter, True, 'Fig [1], "a" ', ['![Fig \\[1\\], "a" \\](data:image/png;base64,']), - (MarkdownExporter, False, 'Fig [1], "a" ', ['![Fig \\[1\\], "a" ](data:image/png;base64,']), - ( - AsciiDocExporter, - True, - 'Fig [1], "a" ', - ['.Fig [1], "a" \\\nimage::data:image/png;base64,', '["Fig [1], \\"a\\" "]'], - ), - (AsciiDocExporter, True, "... continued", [".{empty}... continued\nimage::data:image/png;base64,"]), - (AsciiDocExporter, True, None, ["image::data:image/png;base64,", "[]\n\nThe trend continued"]), - ( - HTMLExporter, - True, - 'Fig [1], "a" ', - ['alt="Fig [1], "a" <b>"', '
    Fig [1], "a" <b>
    '], - ), + (["Figure 1: examples across dif-", "ferent categories"], "Figure 1: examples across different categories"), + (["the Encoder pro-", "duces features"], "the Encoder produces features"), + (["models trained on DocLayNet-", "trained data"], "models trained on DocLayNet-trained data"), + (["a state-of-the-", "art model"], "a state-of-the-art model"), + (["ends with a dash -", "then more"], "ends with a dash -then more"), + (["Section 2-", "Results"], "Section 2- Results"), + (["one line"], "one line"), + (["", " padded ", ""], "padded"), + ([], ""), ], ) -def test_page_export_figure_captions(exporter, escape, caption, expected): - exported = exporter().export_page(_figure_page(caption=caption), escape=escape, images="embedded") - assert all(part in exported for part in expected) +def test_join_caption_lines(lines, expected): + assert _join_caption_lines(lines) == expected @pytest.mark.parametrize( @@ -1288,193 +1112,17 @@ def _geom(geometry): elements.TableCell("B", 0.9, _geom((((x0 + x1) / 2, y0), (x1, y1))), 0, 0, 1, 1), ] tables.append(elements.Table(cells, 1, 2, _geom(((x0, y0), (x1, y1))), 0.9)) - page = elements.Page( - np.full((1000, 800, 3), 255, dtype=np.uint8), - [elements.Block(lines=lines)], - 0, - (1000, 800), + page = _page( + lines, + image=np.full((1000, 800, 3), 255, dtype=np.uint8), layout=[elements.LayoutElement(label, 0.9, _geom(geometry)) for label, geometry in layout], tables=tables, ) assert re.sub(r"\(data:[^)]+\)", "(...)", page.export_as_markdown(images="embedded")) == expected -def test_page_export_images_without_page_image(): - # A page restored from JSON has no pixels: placeholders, and the inner text is kept - restored = elements.Page.from_dict(_figure_page().export()) - markdown = restored.export_as_markdown(images="embedded") - assert "" in markdown - assert "axis label" in markdown - assert "Figure 1 quarterly revenue" in markdown - - -def test_text_and_json_exports_ignore_figures(): - page = _figure_page() - assert "image" not in page.render() - assert page.render().startswith("Annual Report") - assert TextExporter().export_page(page, images="embedded") == page.render() - # Figures do not leak into the JSON blocks - exported = page.export() - assert len(exported["blocks"]) == 5 - assert len(exported["layout"]) == 5 - json.dumps(exported) - - -@pytest.mark.parametrize( - "picture", - [ - ((0.12, 0.30), (0.88, 0.60)), - ((0.12, 0.30), (0.88, 0.30), (0.88, 0.60), (0.12, 0.60)), - ], -) -def test_page_export_as_xml_figures(picture): - page = _figure_page(picture=picture) - xml_bytes, _ = page.export_as_xml() - xml = xml_bytes.decode() - assert 'class="ocr_photo" id="figure_1"' in xml - assert 'title="bbox 96 300 704 600"' in xml - assert "ocr_photo" in xml.split('name="ocr-capabilities" content="')[1].split('"')[0] - # The raw order also carries the figures - assert 'class="ocr_photo"' in page.export_as_xml(reading_order=False)[0].decode() - - -def test_document_export_passes_images_through(tmp_path): - first, second = _figure_page(), _figure_page() - second.page_idx = 1 - doc = elements.Document([first, second]) - assert doc.export_as_markdown().count("") == 2 - assert doc.export_as_markdown(images="none").count("") == 0 - encoder = FigureEncoder("referenced", image_dir=tmp_path) - doc.export_as_html(images=encoder) - assert len(encoder.written) == 2 - - -def test_asciidoc_figure_title_and_attributes(): - page = _figure_page(caption="Figure 1: {author} revenue") - asciidoc = page.export_as_asciidoc(images="embedded") - assert '[caption=""]\n.Figure 1: \\{author\\} revenue\nimage::data:image/png;base64,' in asciidoc - assert '["Figure 1: \\{author\\} revenue"]' in asciidoc - raw = page.export_as_asciidoc(images="embedded", escape=False) - assert '[caption=""]\n.Figure 1: {author} revenue\n' in raw and '["Figure 1: {author} revenue"]' in raw - assert "[caption=" not in _figure_page(caption=None).export_as_asciidoc(images="embedded") - - -@pytest.mark.parametrize( - "lines, expected", - [ - (["Figure 1: examples across dif-", "ferent categories"], "Figure 1: examples across different categories"), - (["the Encoder pro-", "duces features"], "the Encoder produces features"), - (["models trained on DocLayNet-", "trained data"], "models trained on DocLayNet-trained data"), - (["a state-of-the-", "art model"], "a state-of-the-art model"), - (["ends with a dash -", "then more"], "ends with a dash -then more"), - (["Section 2-", "Results"], "Section 2- Results"), - (["one line"], "one line"), - (["", " padded ", ""], "padded"), - ([], ""), - ], -) -def test_join_caption_lines(lines, expected): - assert _join_caption_lines(lines) == expected - - -def test_figure_caption_hyphenation(): - page = _figure_page(caption=None) - caption_lines = [_line_at("Figure 1: quarterly rev-", 0.2, 0.63, 0.8, 0.66), _line_at("enue", 0.2, 0.67, 0.4, 0.7)] - page.blocks = [*page.blocks, elements.Block(lines=caption_lines)] - page.layout = [*page.layout, elements.LayoutElement("Caption", 0.96, ((0.18, 0.62), (0.82, 0.705)))] - assert "![Figure 1: quarterly revenue](data:image/png;base64," in page.export_as_markdown(images="embedded") - assert "
    Figure 1: quarterly revenue
    " in page.export_as_html(images="embedded") - - -@pytest.mark.parametrize("fmt", ["markdown", "asciidoc", "html"]) -def test_page_export_images_none(fmt): - page = _figure_page() - exported = getattr(page, f"export_as_{fmt}")(images="none") - assert "image" not in exported and "base64" not in exported - assert "axis label" in exported - assert "Figure 1 quarterly revenue" in exported - - -def test_page_export_unresolved_figure_keeps_its_text(): - # A figure without pixels falls back to a placeholder and keeps its text and caption - class SkippingEncoder(FigureEncoder): - def source(self, page, region, index): - return None - - page = _figure_page() - markdown = page.export_as_markdown(images=SkippingEncoder("embedded")) - assert "" in markdown - assert "axis label" in markdown - assert "Figure 1 quarterly revenue" in markdown - # With its pixels, the inner text is dropped - markdown = page.export_as_markdown(images="embedded") - assert "axis label" not in markdown - - -def test_page_export_figures_in_page_furniture(): - page = _figure_page() - page.page[10:40, 20:120] = (0, 0, 255) - page.blocks[0].lines.append(_line_at("ACME", 0.04, 0.015, 0.12, 0.035)) - page.layout = [ - *page.layout, - elements.LayoutElement("Page-header", 0.9, ((0.0, 0.0), (0.2, 0.05))), - elements.LayoutElement("Picture", 0.9, ((0.02, 0.01), (0.15, 0.04))), - ] - # With the furniture, the logo is a figure like any other - embedded = page.export_as_markdown(images="embedded") - assert embedded.count("![") == 2 and "ACME" not in embedded - placeholder = page.export_as_markdown(images="placeholder") - assert placeholder.count("") == 2 and "ACME" in placeholder - # Without the furniture, neither the logo nor its text - for images in ("none", "placeholder", "embedded"): - exported = page.export_as_markdown(images=images, include_furniture=False) - assert "ACME" not in exported - assert "Annual Report" in exported and "Figure 1 quarterly revenue" in exported - assert page.export_as_markdown(images="embedded", include_furniture=False).count("![") == 1 - assert page.export_as_markdown(images="placeholder", include_furniture=False).count("") == 1 - - -def test_page_export_caption_split_across_columns(): - # The last caption line is closer to the table of the other column than to its own figure - image = np.full((1000, 800, 3), 255, dtype=np.uint8) - image[115:330, 420:720] = (0, 0, 255) - lines = [ - _line_at("Figure 5: Prediction performance of a", 0.52, 0.352, 0.91, 0.366), - _line_at("network trained on increasing fractions", 0.52, 0.37, 0.91, 0.384), - _line_at("significantly better predictions.", 0.52, 0.419, 0.73, 0.434), - _line_at("Body text of the left column", 0.08, 0.47, 0.48, 0.49), - ] - cells = [ - elements.TableCell("A", 0.9, ((0.102, 0.248), (0.29, 0.443)), 0, 0, 0, 0), - elements.TableCell("B", 0.9, ((0.29, 0.248), (0.477, 0.443)), 0, 0, 1, 1), - ] - page = elements.Page( - image, - [elements.Block(lines=lines)], - 0, - (1000, 800), - layout=[ - elements.LayoutElement("Picture", 0.9, ((0.526, 0.115), (0.903, 0.33))), - elements.LayoutElement("Caption", 0.9, ((0.515, 0.348), (0.918, 0.438))), - elements.LayoutElement("Table", 0.9, ((0.1, 0.245), (0.48, 0.445))), - elements.LayoutElement("Text", 0.9, ((0.07, 0.465), (0.49, 0.495))), - ], - tables=[elements.Table(cells, 1, 2, ((0.102, 0.248), (0.477, 0.443)), 0.9)], - ) - caption = ( - "Figure 5: Prediction performance of a network trained on increasing fractions " - "significantly better predictions." - ) - html = page.export_as_html(images="embedded") - assert f"
    {caption}
    " in html - # The whole caption is read right after its figure - markdown = page.export_as_markdown(images="placeholder") - figure, table = markdown.index(""), markdown.index("| A | B |") - assert figure < markdown.index("Figure 5") < markdown.index("significantly better") < table - - -def test_page_export_figure_does_not_split_a_paragraph(): - """A figure beside a paragraph is exported after its last line""" +def test_page_export_figure_placement(): + # A figure beside a paragraph is exported after its last line lines = [_line_at("Wide title spanning the page", 0.1, 0.05, 0.9, 0.09)] # A 3-column author grid across the gutter: not detected as multi-column for row in range(4): @@ -1492,8 +1140,7 @@ def test_page_export_figure_does_not_split_a_paragraph(): ] image = np.zeros((1000, 800, 3), dtype=np.uint8) image[310:667, 416:728] = 255 - page = elements.Page(image, [elements.Block(lines=lines)], 0, (1000, 800), layout=layout) - markdown = page.export_as_markdown(images="embedded") + markdown = _page(lines, image=image, layout=layout).export_as_markdown(images="embedded") assert ( markdown.index("last line of the paragraph") < markdown.index("![](data:image/png") @@ -1514,48 +1161,76 @@ def test_page_export_figure_does_not_split_a_paragraph(): 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) + items, _, _ = page_reading_order(_page(lines, dims=(1810, 1280), layout=layout)) assert len(items) == 1 and len(items[0].lines) == 6 - -def test_page_export_figure_stacked_in_a_loose_region(): - """A loose text region enclosing a figure does not push the figure after its last line""" + # A loose text region enclosing a figure does not push the figure after its last line lines = [_line_at(f"above {idx}", 0.1, 0.1 + idx * 0.04, 0.9, 0.13 + idx * 0.04) for idx in range(3)] lines += [_line_at(f"below {idx}", 0.1, 0.7 + idx * 0.04, 0.9, 0.73 + idx * 0.04) for idx in range(3)] layout = [ elements.LayoutElement("Text", 0.9, ((0.08, 0.09), (0.92, 0.83))), elements.LayoutElement("Picture", 0.9, ((0.2, 0.3), (0.8, 0.6))), ] - page = elements.Page( - np.zeros((1000, 800, 3), dtype=np.uint8), [elements.Block(lines=lines)], 0, (1000, 800), layout=layout - ) - markdown = page.export_as_markdown() + markdown = _page(lines, image=np.zeros((1000, 800, 3), dtype=np.uint8), layout=layout).export_as_markdown() assert markdown.index("above 2") < markdown.index("") < markdown.index("below 0") + # A caption whose last line is closer to the table of the other column than to its own figure + image = np.full((1000, 800, 3), 255, dtype=np.uint8) + image[115:330, 420:720] = (0, 0, 255) + lines = [ + _line_at("Figure 5: Prediction performance of a", 0.52, 0.352, 0.91, 0.366), + _line_at("network trained on increasing fractions", 0.52, 0.37, 0.91, 0.384), + _line_at("significantly better predictions.", 0.52, 0.419, 0.73, 0.434), + _line_at("Body text of the left column", 0.08, 0.47, 0.48, 0.49), + ] + cells = [ + elements.TableCell("A", 0.9, ((0.102, 0.248), (0.29, 0.443)), 0, 0, 0, 0), + elements.TableCell("B", 0.9, ((0.29, 0.248), (0.477, 0.443)), 0, 0, 1, 1), + ] + layout = [ + elements.LayoutElement("Picture", 0.9, ((0.526, 0.115), (0.903, 0.33))), + elements.LayoutElement("Caption", 0.9, ((0.515, 0.348), (0.918, 0.438))), + elements.LayoutElement("Table", 0.9, ((0.1, 0.245), (0.48, 0.445))), + elements.LayoutElement("Text", 0.9, ((0.07, 0.465), (0.49, 0.495))), + ] + page = _page( + lines, image=image, layout=layout, tables=[elements.Table(cells, 1, 2, ((0.102, 0.248), (0.477, 0.443)), 0.9)] + ) + caption = ( + "Figure 5: Prediction performance of a network trained on increasing fractions " + "significantly better predictions." + ) + assert f"
    {caption}
    " in page.export_as_html(images="embedded") + # The whole caption is read right after its figure + markdown = page.export_as_markdown(images="placeholder") + figure, table = markdown.index(""), markdown.index("| A | B |") + assert figure < markdown.index("Figure 5") < markdown.index("significantly better") < table -@pytest.mark.parametrize( - "escape, caption, expected", - [ - (True, "Figure 1: *bold* claim", "\n\n*Figure 1: \\*bold\\* claim*"), - (False, "Figure 1: plain", "\n\n*Figure 1: plain*"), - ], -) -def test_page_export_markdown_caption_paragraph(escape, caption, expected): - page = _figure_page(caption=caption) - markdown = page.export_as_markdown(images="embedded", escape=escape) - image_end = markdown.index(")", markdown.index("](data:image/png;base64,")) - assert markdown[image_end + 1 :].startswith(expected) - # No caption line without a caption or without pixels - assert "\n\n*" not in _figure_page(caption=None).export_as_markdown(images="embedded") - placeholder = page.export_as_markdown(images="placeholder", escape=escape) - assert "\n\n*" not in placeholder +def test_page_export_figures_in_page_furniture(): + # A logo in the page header is read with it, and its text is labeled as page furniture + page = _figure_page() + page.page[10:40, 20:120] = (0, 0, 255) + page.blocks[0].lines.append(_line_at("ACME", 0.04, 0.015, 0.12, 0.035)) + page.layout = [ + *page.layout, + elements.LayoutElement("Page-header", 0.9, ((0.0, 0.0), (0.2, 0.05))), + elements.LayoutElement("Picture", 0.9, ((0.02, 0.01), (0.15, 0.04))), + ] + # With the furniture, the logo is a figure like any other + embedded = page.export_as_markdown(images="embedded") + assert embedded.count("![") == 2 and "ACME" not in embedded + placeholder = page.export_as_markdown(images="placeholder") + assert placeholder.count("") == 2 and "ACME" in placeholder + # Without the furniture, neither the logo nor its text + for images in ("none", "placeholder", "embedded"): + exported = page.export_as_markdown(images=images, include_furniture=False) + assert "ACME" not in exported + assert "Annual Report" in exported and "Figure 1 quarterly revenue" in exported + assert page.export_as_markdown(images="embedded", include_furniture=False).count("![") == 1 + assert page.export_as_markdown(images="placeholder", include_furniture=False).count("") == 1 -def test_page_reading_order_logo_in_page_header(): - """A logo in the page header is read with it, and its text is labeled as page furniture""" + # Same with a page header spanning the page, with more text beside the logo page = _figure_page() page.page[10:40, 20:120] = (0, 0, 255) page.blocks[0].lines.append(_line_at("ACME", 0.04, 0.015, 0.12, 0.035)) @@ -1614,7 +1289,18 @@ def test_page_export_figure_partly_in_page_furniture(bottom, in_furniture): assert ("tick" in exported) is (images == "placeholder" and not in_furniture) -def test_page_export_referenced_mode_needs_an_encoder(): - page = _figure_page() - with pytest.raises(ValueError, match="image_dir"): - page.export_as_markdown(images="referenced") +@pytest.mark.parametrize( + "picture", + [ + ((0.12, 0.30), (0.88, 0.60)), + ((0.12, 0.30), (0.88, 0.30), (0.88, 0.60), (0.12, 0.60)), + ], +) +def test_page_export_as_xml_figures(picture): + page = _figure_page(picture=picture) + xml = page.export_as_xml()[0].decode() + assert 'class="ocr_photo" id="figure_1"' in xml + assert 'title="bbox 96 300 704 600"' in xml + assert "ocr_photo" in xml.split('name="ocr-capabilities" content="')[1].split('"')[0] + # The raw order also carries the figures + assert 'class="ocr_photo"' in page.export_as_xml(reading_order=False)[0].decode() diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index 0c3b801929..252960d3a7 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -10,11 +10,35 @@ words_per_page = 10 -boxes_1 = {CLASS_NAME: np.random.rand(words_per_page, 6)} # dict format -boxes_1[CLASS_NAME][:2] *= boxes_1[CLASS_NAME][2:4] -boxes_2 = np.random.rand(words_per_page, 6) # array format -boxes_2[:2] *= boxes_2[2:4] +def _orientations(num): + return [{"value": 0, "confidence": None}] * num + + +def _word_args(words, boxes, shape=(1000, 1000), confidences=None): + """Pack builder arguments for a single page from word values and boxes""" + confidences = [0.95] * len(words) if confidences is None else confidences + return ( + [np.zeros((*shape, 3), dtype=np.uint8)], + [np.asarray(boxes, dtype=np.float32)], + [np.full(len(words), 0.9, dtype=np.float32)], + [list(zip(words, confidences))], + [shape], + [_orientations(len(words))], + ) + + +def _words(page): + return [word.value for block in page.blocks for line in block.lines for word in line.words] + + +def _rot_poly(x0, y0, x1, y1, deg, cx=0.5, cy=0.5): + a = np.deg2rad(deg) + ca, sa = np.cos(a), np.sin(a) + return [ + [cx + (x - cx) * ca - (y - cy) * sa, cy + (x - cx) * sa + (y - cy) * ca] + for x, y in [(x0, y0), (x1, y0), (x1, y1), (x0, y1)] + ] def test_documentbuilder(): @@ -34,46 +58,31 @@ def test_documentbuilder(): [objectness_scores, objectness_scores], [("hello", 1.0)] * 3, [(100, 200)] * num_pages, - [{"value": 0, "confidence": None}] * 3, + _orientations(3), ) - out = doc_builder( + args = ( pages, [boxes, boxes], [objectness_scores, objectness_scores], [[("hello", 1.0)] * words_per_page] * num_pages, [(100, 200)] * num_pages, - [[{"value": 0, "confidence": None}] * words_per_page] * num_pages, + [_orientations(words_per_page)] * num_pages, ) + out = doc_builder(*args) assert isinstance(out, Document) assert len(out.pages) == num_pages - assert all(isinstance(page.page, np.ndarray) for page in out.pages) and all( - page.page.shape == (100, 200, 3) for page in out.pages - ) + assert all(isinstance(page.page, np.ndarray) and page.page.shape == (100, 200, 3) for page in out.pages) # 1 Block & 1 line per page assert len(out.pages[0].blocks) == 1 and len(out.pages[0].blocks[0].lines) == 1 assert len(out.pages[0].blocks[0].lines[0].words) == words_per_page # Resolve lines doc_builder = builder.DocumentBuilder(resolve_lines=True, resolve_blocks=True) - out = doc_builder( - pages, - [boxes, boxes], - [objectness_scores, objectness_scores], - [[("hello", 1.0)] * words_per_page] * num_pages, - [(100, 200)] * num_pages, - [[{"value": 0, "confidence": None}] * words_per_page] * num_pages, - ) + assert isinstance(doc_builder(*args), Document) # No detection - boxes = np.zeros((0, 4)) - objectness_scores = np.zeros([0]) out = doc_builder( - pages, - [boxes, boxes], - [objectness_scores, objectness_scores], - [[], []], - [(100, 200)] * num_pages, - [[]] * num_pages, + pages, [np.zeros((0, 4))] * 2, [np.zeros([0])] * 2, [[], []], [(100, 200)] * num_pages, [[]] * num_pages ) assert len(out.pages[0].blocks) == 0 @@ -82,15 +91,14 @@ def test_documentbuilder(): [[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]], ]) - objectness_scores = np.array([0.99, 0.99]) doc_builder_2 = builder.DocumentBuilder(resolve_blocks=False, resolve_lines=False, export_as_straight_boxes=True) out = doc_builder_2( [np.zeros((100, 100, 3))], [boxes], - [objectness_scores], + [np.array([0.99, 0.99])], [[("hello", 0.99), ("word", 0.99)]], [(100, 100)], - [[{"value": 0, "confidence": None}] * 2], + [_orientations(2)], ) 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 @@ -105,6 +113,23 @@ def test_documentbuilder(): "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False, keep_reading_order=False)" ) + # A single mismatched argument must raise, whichever one it is + boxes = np.random.rand(words_per_page, 6) + boxes[:2] *= boxes[2:4] + kwargs = dict( + pages=pages, + boxes=[boxes[:, :4]] * num_pages, + objectness_scores=[np.array([0.9] * words_per_page)] * num_pages, + text_preds=[[("hello", 0.99)] * words_per_page] * num_pages, + page_shapes=[(100, 200)] * num_pages, + crop_orientations=[_orientations(words_per_page)] * num_pages, + ) + # sanity: consistent arguments build fine + assert isinstance(builder.DocumentBuilder()(**kwargs), Document) + for key in ("pages", "objectness_scores", "text_preds", "page_shapes", "crop_orientations"): + with pytest.raises(ValueError): + builder.DocumentBuilder()(**{**kwargs, key: kwargs[key] + [kwargs[key][0]]}) + def test_kiedocumentbuilder(): num_pages = 2 @@ -123,44 +148,33 @@ def test_kiedocumentbuilder(): [objectness_scores, objectness_scores], [{CLASS_NAME: ("hello", 1.0)}] * 3, [(100, 200), (100, 200)], - [{CLASS_NAME: [{"value": 0, "confidence": None}] * 3}], + [{CLASS_NAME: _orientations(3)}], ) - out = doc_builder( + args = ( pages, [predictions, predictions], [objectness_scores, objectness_scores], [{CLASS_NAME: [("hello", 1.0)] * words_per_page}] * num_pages, [(100, 200), (100, 200)], - [{CLASS_NAME: [{"value": 0, "confidence": None}] * words_per_page}] * num_pages, + [{CLASS_NAME: _orientations(words_per_page)}] * num_pages, ) + out = doc_builder(*args) assert isinstance(out, KIEDocument) assert len(out.pages) == num_pages - assert all(isinstance(page.page, np.ndarray) for page in out.pages) and all( - page.page.shape == (100, 200, 3) for page in out.pages - ) - # 1 Block & 1 line per page + assert all(isinstance(page.page, np.ndarray) and page.page.shape == (100, 200, 3) for page in out.pages) assert len(out.pages[0].predictions) == 1 assert len(out.pages[0].predictions[CLASS_NAME]) == words_per_page # Resolve lines doc_builder = builder.KIEDocumentBuilder(resolve_lines=True, resolve_blocks=True) - out = doc_builder( - pages, - [predictions, predictions], - [objectness_scores, objectness_scores], - [{CLASS_NAME: [("hello", 1.0)] * words_per_page}] * num_pages, - [(100, 200), (100, 200)], - [{CLASS_NAME: [{"value": 0, "confidence": None}] * words_per_page}] * num_pages, - ) + assert isinstance(doc_builder(*args), KIEDocument) # No detection predictions = {CLASS_NAME: np.zeros((0, 4))} - objectness_scores = {CLASS_NAME: np.zeros((1))} - out = doc_builder( pages, [predictions, predictions], - [objectness_scores, objectness_scores], + [{CLASS_NAME: np.zeros(1)}] * 2, [{CLASS_NAME: []}, {CLASS_NAME: []}], [(100, 200), (100, 200)], [{CLASS_NAME: []}, {CLASS_NAME: []}], @@ -174,15 +188,14 @@ def test_kiedocumentbuilder(): [[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]], ]) } - objectness_scores = {CLASS_NAME: np.array([0.99, 0.99])} doc_builder_2 = builder.KIEDocumentBuilder(resolve_blocks=False, resolve_lines=False, export_as_straight_boxes=True) out = doc_builder_2( [np.zeros((100, 100, 3))], [predictions], - [objectness_scores], + [{CLASS_NAME: np.array([0.99, 0.99])}], [{CLASS_NAME: [("hello", 0.99), ("word", 0.99)]}], [(100, 100)], - [{CLASS_NAME: [{"value": 0, "confidence": None}] * 2}], + [{CLASS_NAME: _orientations(2)}], ) assert out.pages[0].predictions[CLASS_NAME][0].geometry == ((0.05, 0.1), (0.2, 0.25)) assert out.pages[0].predictions[CLASS_NAME][1].geometry == ((0.45, 0.5), (0.6, 0.65)) @@ -194,67 +207,53 @@ def test_kiedocumentbuilder(): "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False, keep_reading_order=False)" ) - -def test_documentbuilder_layout(): - - doc_builder = builder.DocumentBuilder() - boxes = np.array([[0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4]]) - objectness_scores = np.array([0.9, 0.9]) - regions = [ - { - "boxes": np.array([[0.05, 0.02, 0.95, 0.08], [0.05, 0.2, 0.95, 0.5]], dtype=np.float32), - "class_names": ["Title", "Text"], - "scores": [0.95, 0.88], - } - ] - out = doc_builder( + # Layout regions are attached to the page and exported + args = ( [np.zeros((100, 100, 3))], - [boxes], - [objectness_scores], - [[("hello", 0.99), ("world", 0.99)]], + [{CLASS_NAME: np.array([[0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4]])}], + [{CLASS_NAME: np.array([0.9, 0.9])}], + [{CLASS_NAME: [("hello", 0.99), ("world", 0.99)]}], [(100, 100)], - [[{"value": 0, "confidence": None}] * 2], - regions=regions, + [{CLASS_NAME: _orientations(2)}], ) - page = out.pages[0] + regions = { + "boxes": np.array([[0.05, 0.02, 0.95, 0.08], [0.05, 0.2, 0.95, 0.5]], dtype=np.float32), + "class_names": ["Title", "Text"], + "scores": [0.95, 0.88], + } + page = builder.KIEDocumentBuilder()(*args, regions=[regions]).pages[0] + assert all(isinstance(region, LayoutElement) for region in page.layout) + assert [region.type for region in page.layout] == ["Title", "Text"] + assert page.export()["layout"] == [region.export() for region in page.layout] + # no regions -> empty layout + assert builder.KIEDocumentBuilder()(*args).pages[0].layout == [] + + +def test_documentbuilder_layout(): + doc_builder = builder.DocumentBuilder() + args = _word_args(["hello", "world"], [[0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4]], shape=(100, 100)) + regions = { + "boxes": np.array([[0.05, 0.02, 0.95, 0.08], [0.05, 0.2, 0.95, 0.5]], dtype=np.float32), + "class_names": ["Title", "Text"], + "scores": [0.95, 0.88], + } + page = doc_builder(*args, regions=[regions]).pages[0] # Layout regions are attached as LayoutElement and exported - assert len(page.layout) == 2 assert all(isinstance(region, LayoutElement) for region in page.layout) assert [region.type for region in page.layout] == ["Title", "Text"] assert page.layout[0].confidence == pytest.approx(0.95) assert np.allclose(np.array(page.layout[0].geometry), [[0.05, 0.02], [0.95, 0.08]], atol=1e-6) assert page.export()["layout"] == [region.export() for region in page.layout] - # no regions -> empty layout - out_no_layout = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [objectness_scores], - [[("hello", 0.99), ("world", 0.99)]], - [(100, 100)], - [[{"value": 0, "confidence": None}] * 2], - ) - assert out_no_layout.pages[0].layout == [] - assert out_no_layout.pages[0].export()["layout"] == [] - + page = doc_builder(*args).pages[0] + assert page.layout == [] and page.export()["layout"] == [] # Rotated layout polygons (4, 2) are converted to a 4-point geometry - rotated_regions = [ - { - "boxes": np.array([[[0.1, 0.1], [0.4, 0.12], [0.39, 0.3], [0.09, 0.28]]], dtype=np.float32), - "class_names": ["Table"], - "scores": [0.7], - } - ] - out_rot = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [objectness_scores], - [[("hello", 0.99), ("world", 0.99)]], - [(100, 100)], - [[{"value": 0, "confidence": None}] * 2], - regions=rotated_regions, - ) - region = out_rot.pages[0].layout[0] + rotated = { + "boxes": np.array([[[0.1, 0.1], [0.4, 0.12], [0.39, 0.3], [0.09, 0.28]]], dtype=np.float32), + "class_names": ["Table"], + "scores": [0.7], + } + region = doc_builder(*args, regions=[rotated]).pages[0].layout[0] assert region.type == "Table" assert isinstance(region.geometry, tuple) and len(region.geometry) == 4 @@ -270,28 +269,20 @@ def test_documentbuilder_layout(): # 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) + lines = doc_builder._resolve_lines(boxes, word_regions=doc_builder._word_regions(boxes, columns, ["Text"] * 3)) 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 - ) + word_regions = doc_builder._word_regions(boxes, overlapping, ["Text"] * 2) + assert len(doc_builder._resolve_lines(boxes, word_regions=word_regions)) == 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], + page = doc_builder( + *_word_args(["word"] * 18, boxes, shape=(100, 100)), 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 + ).pages[0] + assert [len(line.words) for block in page.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)) @@ -306,154 +297,6 @@ def test_documentbuilder_layout(): 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 { - "geometry": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], - "score": score, - "row_start": rs, - "row_end": re, - "col_start": cs, - "col_end": ce, - } - - -def test_documentbuilder_tables(): - doc_builder = builder.DocumentBuilder(resolve_lines=True) - - # 4 words inside a top table, 2 inside a bottom table, 1 caption outside both - def wbox(cx, cy, w=0.04, h=0.02): - return [cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2] - - words = [ - ("Name", 0.17, 0.15), - ("Age", 0.36, 0.15), - ("Alice", 0.17, 0.27), - ("30", 0.36, 0.27), - ("City", 0.17, 0.63), - ("Pop", 0.36, 0.63), - ("caption", 0.30, 0.92), - ] - boxes = np.array([wbox(cx, cy) for _, cx, cy in words], dtype=np.float32) - text_preds = [[(w, 0.95) for w, _, _ in words]] - objectness_scores = np.full(len(words), 0.9, dtype=np.float32) - orientations = [[{"value": 0, "confidence": None}] * len(words)] - - # The OCR pipeline passes a list of grids (one per cropped table region), in page-relative coordinates. - # The bottom table uses offset (1-based) logical coordinates to exercise local re-indexing. - table_top = { - "cells": [ - _table_cell(0.10, 0.10, 0.25, 0.20, 0, 0, 0, 0), - _table_cell(0.28, 0.10, 0.45, 0.20, 0, 0, 1, 1), - _table_cell(0.10, 0.22, 0.25, 0.32, 1, 1, 0, 0), - _table_cell(0.28, 0.22, 0.45, 0.32, 1, 1, 1, 1), - ], - "num_rows": 2, - "num_cols": 2, - } - table_bottom = { - "cells": [ - _table_cell(0.10, 0.58, 0.25, 0.68, 1, 1, 1, 1), - _table_cell(0.28, 0.58, 0.45, 0.68, 1, 1, 2, 2), - ], - "num_rows": 99, # deliberately wrong dims -> recomputed from local coordinates - "num_cols": 99, - } - - out = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [objectness_scores], - text_preds, - [(100, 100)], - orientations, - tables=[[table_top, table_bottom]], - ) - page = out.pages[0] - - # One Table per provided grid - assert len(page.tables) == 2 - assert all(isinstance(t, Table) for t in page.tables) - assert page.tables[0].to_grid() == [["Name", "Age"], ["Alice", "30"]] - # bottom table re-indexed from offset coordinates to a local 0-based 1 x 2 grid - assert (page.tables[1].num_rows, page.tables[1].num_cols) == (1, 2) - assert page.tables[1].to_grid() == [["City", "Pop"]] - - # Words assigned to a table are removed from the blocks; the caption remains - remaining = [w.value for b in page.blocks for line in b.lines for w in line.words] - assert remaining == ["caption"] - - # Tables are part of the page export - exported = page.export() - assert len(exported["tables"]) == 2 - assert page.tables[0].to_grid() == [["Name", "Age"], ["Alice", "30"]] - - # A single grid (dict) is also accepted -> one table - out_single = doc_builder( - [np.zeros((100, 100, 3))], - [boxes[:4]], - [objectness_scores[:4]], - [text_preds[0][:4]], - [(100, 100)], - [orientations[0][:4]], - tables=[table_top], - ) - assert len(out_single.pages[0].tables) == 1 - - # No tables -> empty page.tables and every word is kept in the blocks - out_none = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [objectness_scores], - text_preds, - [(100, 100)], - orientations, - ) - assert out_none.pages[0].tables == [] - assert out_none.pages[0].export()["tables"] == [] - kept = sorted(w.value for b in out_none.pages[0].blocks for line in b.lines for w in line.words) - assert kept == sorted(w for w, _, _ in words) - - -def test_kiedocumentbuilder_layout(): - from doctr.io.elements import LayoutElement - - doc_builder = builder.KIEDocumentBuilder() - predictions = {CLASS_NAME: np.array([[0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4]])} - objectness_scores = {CLASS_NAME: np.array([0.9, 0.9])} - regions = [ - { - "boxes": np.array([[0.05, 0.02, 0.95, 0.08], [0.05, 0.2, 0.95, 0.5]], dtype=np.float32), - "class_names": ["Title", "Text"], - "scores": [0.95, 0.88], - } - ] - out = doc_builder( - [np.zeros((100, 100, 3))], - [predictions], - [objectness_scores], - [{CLASS_NAME: [("hello", 0.99), ("world", 0.99)]}], - [(100, 100)], - [{CLASS_NAME: [{"value": 0, "confidence": None}] * 2}], - regions=regions, - ) - page = out.pages[0] - assert len(page.layout) == 2 - assert all(isinstance(region, LayoutElement) for region in page.layout) - assert [region.type for region in page.layout] == ["Title", "Text"] - assert page.export()["layout"] == [region.export() for region in page.layout] - - # no regions -> empty layout - out_no_layout = doc_builder( - [np.zeros((100, 100, 3))], - [predictions], - [objectness_scores], - [{CLASS_NAME: [("hello", 0.99), ("world", 0.99)]}], - [(100, 100)], - [{CLASS_NAME: [{"value": 0, "confidence": None}] * 2}], - ) - assert out_no_layout.pages[0].layout == [] - - @pytest.mark.parametrize( "input_boxes, sorted_idxs", [ @@ -470,6 +313,8 @@ def test_kiedocumentbuilder_layout(): ], [0, 1], ], # rot + # Boxes with zero height do not produce a NaN ordering (division by the median height) + [[[0.5, 0.2, 0.6, 0.2], [0.1, 0.2, 0.2, 0.2]], [1, 0]], ], ) def test_sort_boxes(input_boxes, sorted_idxs): @@ -611,6 +456,8 @@ def _justified_lines(scale=1.0): None, [[0, 10, 1, 11, 2, 12, 3, 4], [5, 13, 6, 14, 7, 15, 8, 9]], ], + # Boxes with zero height cannot be grouped into lines (no median height) + [[[0.5, 0.2, 0.6, 0.2], [0.1, 0.2, 0.2, 0.2]], None, [[1], [0]]], ], ) def test_resolve_lines(input_boxes, shape, lines): @@ -618,7 +465,7 @@ def test_resolve_lines(input_boxes, shape, lines): assert doc_builder._resolve_lines(np.asarray(input_boxes), shape) == lines -def test_points_in_polygons(): +def test_table_geometry_helpers(): polys = np.array( [ [[0.1, 0.1], [0.4, 0.1], [0.4, 0.3], [0.1, 0.3]], # axis-aligned quad @@ -628,17 +475,10 @@ def test_points_in_polygons(): ) points = np.array([[0.2, 0.2], [0.6, 0.7], [0.95, 0.95], [0.05, 0.05]], dtype=np.float32) inside = builder.DocumentBuilder._points_in_polygons(points, polys) - assert inside.shape == (4, 2) - assert inside[0].tolist() == [True, False] - assert inside[1].tolist() == [False, True] - assert not inside[2].any() - assert not inside[3].any() + assert inside.tolist() == [[True, False], [False, True], [False, False], [False, False]] # empty inputs yield empty masks instead of raising assert builder.DocumentBuilder._points_in_polygons(np.zeros((0, 2)), polys).shape == (0, 2) assert builder.DocumentBuilder._points_in_polygons(points, np.zeros((0, 4, 2))).shape == (4, 0) - - -def test_as_cell_polygon(): # flat straight box (xmin, ymin, xmax, ymax) -> (4, 2) polygon poly = builder.DocumentBuilder._as_cell_polygon([0.1, 0.2, 0.5, 0.6]) assert poly.shape == (4, 2) @@ -648,221 +488,243 @@ def test_as_cell_polygon(): assert np.allclose(builder.DocumentBuilder._as_cell_polygon(quad), quad) -def test_documentbuilder_tables_straight_geometry(): - # Straight-mode table predictions store cells as flat (xmin, ymin, xmax, ymax) boxes. - # Regression test: this format used to crash the word-to-cell assignment with an IndexError. - doc_builder = builder.DocumentBuilder() - - def wbox(cx, cy, w=0.04, h=0.02): - return [cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2] - - words = [("Name", 0.17, 0.15), ("Age", 0.36, 0.15), ("caption", 0.30, 0.92)] - boxes = np.array([wbox(cx, cy) for _, cx, cy in words], dtype=np.float32) - text_preds = [[(w, 0.95) for w, _, _ in words]] - objectness_scores = np.full(len(words), 0.9, dtype=np.float32) - orientations = [[{"value": 0, "confidence": None}] * len(words)] - - def straight_cell(x0, y0, x1, y1, cs): - return { - "geometry": [x0, y0, x1, y1], # flat straight-box format, as emitted by the table predictor - "score": 0.9, - "row_start": 0, - "row_end": 0, - "col_start": cs, - "col_end": cs, - } - - table = { - "cells": [straight_cell(0.10, 0.10, 0.25, 0.20, 0), straight_cell(0.28, 0.10, 0.45, 0.20, 1)], - "score": 0.9, +def _table_cell(x0, y0, x1, y1, rs, re, cs, ce, score=0.9): + return { + "geometry": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], + "score": score, + "row_start": rs, + "row_end": re, + "col_start": cs, + "col_end": ce, } - out = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [objectness_scores], - text_preds, - [(100, 100)], - orientations, - tables=[[table]], - ) - page = out.pages[0] - assert len(page.tables) == 1 - assert page.tables[0].to_grid() == [["Name", "Age"]] - # words assigned to the table are removed from the blocks; the caption remains - remaining = [w.value for b in page.blocks for line in b.lines for w in line.words] - assert remaining == ["caption"] +# Word centers: a 2 x 2 table at the top, a 1 x 2 table below it and a caption outside both +_TABLE_WORDS = { + "Name": (0.17, 0.15), + "Age": (0.36, 0.15), + "Alice": (0.17, 0.27), + "30": (0.36, 0.27), + "City": (0.17, 0.63), + "Pop": (0.36, 0.63), + "caption": (0.30, 0.92), +} -def test_documentbuilder_argument_length_validation(): - # A single mismatched argument must raise, whichever one it is - num_pages = 2 - doc_builder = builder.DocumentBuilder() - pages = [np.zeros((100, 200, 3))] * num_pages - boxes = np.random.rand(words_per_page, 6) - boxes[:2] *= boxes[2:4] - page_boxes = [boxes[:, :4]] * num_pages - objectness = [np.array([0.9] * words_per_page)] * num_pages - text_preds = [[("hello", 0.99)] * words_per_page] * num_pages - page_shapes = [(100, 200)] * num_pages - crop_orientations = [[{"value": 0, "confidence": None}] * words_per_page] * num_pages +def _table_words(*names, confidences=None): + boxes = [[cx - 0.02, cy - 0.01, cx + 0.02, cy + 0.01] for cx, cy in (_TABLE_WORDS[name] for name in names)] + return _word_args(names, boxes, shape=(100, 100), confidences=confidences) - args = dict( - pages=pages, - boxes=page_boxes, - objectness_scores=objectness, - text_preds=text_preds, - page_shapes=page_shapes, - crop_orientations=crop_orientations, - ) - # sanity: consistent arguments build fine - assert isinstance(doc_builder(**args), Document) - # each argument mismatched on its own must raise - for key in ("pages", "objectness_scores", "text_preds", "page_shapes", "crop_orientations"): - bad_args = dict(args) - bad_args[key] = args[key] + [args[key][0]] - with pytest.raises(ValueError): - doc_builder(**bad_args) - - -def test_sort_boxes_degenerate_heights(): - # Boxes with zero height must not produce a NaN ordering (division by median height) +def test_documentbuilder_tables(): doc_builder = builder.DocumentBuilder() - boxes = np.array([[0.5, 0.2, 0.6, 0.2], [0.1, 0.2, 0.2, 0.2]], dtype=np.float32) - idxs, _ = doc_builder._sort_boxes(boxes) - assert sorted(np.asarray(idxs).tolist()) == [0, 1] + # The OCR pipeline passes a list of grids (one per cropped table region), in page-relative coordinates. + # The bottom table uses offset (1-based) logical coordinates to exercise local re-indexing. + table_top = { + "cells": [ + _table_cell(0.10, 0.10, 0.25, 0.20, 0, 0, 0, 0), + _table_cell(0.28, 0.10, 0.45, 0.20, 0, 0, 1, 1), + _table_cell(0.10, 0.22, 0.25, 0.32, 1, 1, 0, 0), + _table_cell(0.28, 0.22, 0.45, 0.32, 1, 1, 1, 1), + ], + "num_rows": 2, + "num_cols": 2, + } + table_bottom = { + "cells": [_table_cell(0.10, 0.58, 0.25, 0.68, 1, 1, 1, 1), _table_cell(0.28, 0.58, 0.45, 0.68, 1, 1, 2, 2)], + "num_rows": 99, # deliberately wrong dims -> recomputed from local coordinates + "num_cols": 99, + } + args = _table_words(*_TABLE_WORDS) + page = doc_builder(*args, tables=[[table_top, table_bottom]]).pages[0] + # One Table per provided grid + assert len(page.tables) == 2 and all(isinstance(table, Table) for table in page.tables) + assert page.tables[0].to_grid() == [["Name", "Age"], ["Alice", "30"]] + # bottom table re-indexed from offset coordinates to a local 0-based 1 x 2 grid + assert (page.tables[1].num_rows, page.tables[1].num_cols) == (1, 2) + assert page.tables[1].to_grid() == [["City", "Pop"]] + # Words assigned to a table are removed from the blocks; the caption remains + assert _words(page) == ["caption"] + # Tables are part of the page export + assert len(page.export()["tables"]) == 2 + # A single grid (dict) is also accepted -> one table + assert len(doc_builder(*_table_words("Name", "Age", "Alice", "30"), tables=[table_top]).pages[0].tables) == 1 + # No tables -> empty page.tables and every word is kept in the blocks + page = doc_builder(*args).pages[0] + assert page.tables == [] and page.export()["tables"] == [] + assert sorted(_words(page)) == sorted(_TABLE_WORDS) + # Straight-mode table predictions store cells as flat (xmin, ymin, xmax, ymax) boxes + # Regression test: this format used to crash the word-to-cell assignment with an IndexError + straight = { + "cells": [ + {**_table_cell(0, 0, 0, 0, 0, 0, col, col), "geometry": box} + for col, box in enumerate([[0.10, 0.10, 0.25, 0.20], [0.28, 0.10, 0.45, 0.20]]) + ], + "score": 0.9, + } + page = doc_builder(*_table_words("Name", "Age", "caption"), tables=[[straight]]).pages[0] + assert len(page.tables) == 1 and page.tables[0].to_grid() == [["Name", "Age"]] + assert _words(page) == ["caption"] -def test_documentbuilder_table_cell_scores(): # Each cell keeps the score of the table structure model, next to the recognition confidence of its words - doc_builder = builder.DocumentBuilder() - boxes = np.array([[0.15, 0.14, 0.19, 0.16], [0.34, 0.14, 0.38, 0.16]], dtype=np.float32) table = { "cells": [ _table_cell(0.10, 0.10, 0.25, 0.20, 0, 0, 0, 0, score=0.8), _table_cell(0.28, 0.10, 0.45, 0.20, 0, 0, 1, 1, score=0.6), - # no word inside this cell - _table_cell(0.10, 0.22, 0.25, 0.32, 1, 1, 0, 0, score=0.4), + _table_cell(0.10, 0.22, 0.25, 0.32, 1, 1, 0, 0, score=0.4), # no word inside this cell ], "num_rows": 2, "num_cols": 2, } - out = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [np.array([0.9, 0.9])], - [[("Name", 0.95), ("Age", 0.85)]], - [(100, 100)], - [[{"value": 0, "confidence": None}] * 2], - tables=[table], - ) - cells = out.pages[0].tables[0].cells + page = doc_builder(*_table_words("Name", "Age", confidences=[0.95, 0.85]), tables=[table]).pages[0] + cells = page.tables[0].cells assert [cell.value for cell in cells] == ["Name", "Age", ""] assert [cell.objectness_score for cell in cells] == [0.8, 0.6, 0.4] assert [cell.confidence for cell in cells] == pytest.approx([0.95, 0.85, 0.4]) - assert out.pages[0].tables[0].confidence == pytest.approx(0.6) - exported = out.pages[0].export()["tables"][0]["cells"] - assert [cell["objectness_score"] for cell in exported] == [0.8, 0.6, 0.4] + assert page.tables[0].confidence == pytest.approx(0.6) + assert [cell["objectness_score"] for cell in page.export()["tables"][0]["cells"]] == [0.8, 0.6, 0.4] + # Multi-word cells are read line by line, left to right + words = ["world", "hello", "second", "line", "solo"] + boxes = [ + [0.20, 0.13, 0.28, 0.16], # "world" - first line, right + [0.11, 0.13, 0.19, 0.16], # "hello" - first line, left + [0.11, 0.17, 0.20, 0.20], # "second" - second line, left + [0.21, 0.17, 0.27, 0.20], # "line" - second line, right + [0.36, 0.13, 0.44, 0.16], # "solo" - other cell + ] + table = { + "cells": [_table_cell(0.10, 0.10, 0.30, 0.22, 0, 0, 0, 0), _table_cell(0.32, 0.10, 0.50, 0.22, 0, 0, 1, 1)], + "num_rows": 1, + "num_cols": 2, + } + page = doc_builder(*_word_args(words, boxes), tables=[[table]]).pages[0] + assert len(page.tables) == 1 and page.tables[0].to_grid() == [["hello world second line", "solo"]] + # Every word ended up in the table, so nothing is left in the regular blocks + assert page.blocks == [] + # A degenerate cell whose words all share the same center still produces a single row + table = {"cells": [_table_cell(0.10, 0.10, 0.30, 0.22, 0, 0, 0, 0)], "num_rows": 1, "num_cols": 1} + boxes = [[0.15, 0.15, 0.15, 0.15], [0.16, 0.15, 0.16, 0.15]] + assert doc_builder(*_word_args(["a", "b"], boxes), tables=[[table]]).pages[0].tables[0].to_grid() == [["a b"]] -def test_documentbuilder_tables_empty_cells(): - # A table prediction with no cells (e.g. a false-positive "Table" region where the table model - # finds nothing) must not crash the document build - doc_builder = builder.DocumentBuilder() - boxes = np.array([[0.1, 0.1, 0.2, 0.2]], dtype=np.float32) - out = doc_builder( - [np.zeros((100, 100, 3))], - [boxes], - [np.array([0.9])], - [[("hello", 0.99)]], - [(100, 100)], - [[{"value": 0, "confidence": None}]], - tables=[[{"cells": [], "num_rows": 0, "num_cols": 0}]], - ) - assert out.pages[0].tables == [] - # the word stays in the regular blocks since it was never consumed by a table - assert [w.value for b in out.pages[0].blocks for line in b.lines for w in line.words] == ["hello"] + # A word in the gutter goes to its nearest cell, a word far away from the table stays in the regular blocks + words = ["left", "right", "outside", "faraway"] + boxes = [ + [0.11, 0.12, 0.19, 0.16], # inside cell 0 + [0.33, 0.12, 0.41, 0.16], # inside cell 1 + [0.305, 0.12, 0.315, 0.16], # in the gutter, closest to cell 1 + [0.11, 0.80, 0.25, 0.85], # far outside the table + ] + table = { + "cells": [_table_cell(0.10, 0.10, 0.30, 0.18, 0, 0, 0, 0), _table_cell(0.32, 0.10, 0.50, 0.18, 0, 0, 1, 1)], + "num_rows": 1, + "num_cols": 2, + } + page = doc_builder(*_word_args(words, boxes), tables=[[table]]).pages[0] + grid = page.tables[0].to_grid() + assert grid[0][0] == "left" and "outside" in grid[0][1] + assert _words(page) == ["faraway"] + + # Rotated pages keep 4-point polygons for the table and its cells + cell_boxes = [ + (0.10, 0.10, 0.28, 0.20), + (0.30, 0.10, 0.48, 0.20), + (0.10, 0.22, 0.28, 0.32), + (0.30, 0.22, 0.48, 0.32), + ] + word_boxes = [ + (0.13, 0.13, 0.25, 0.17), + (0.33, 0.13, 0.45, 0.17), + (0.13, 0.25, 0.25, 0.29), + (0.33, 0.25, 0.45, 0.29), + ] + table = { + "cells": [ + {**_table_cell(0, 0, 0, 0, rs, re, cs, ce), "geometry": _rot_poly(*box, 20)} + for box, (rs, re, cs, ce) in zip(cell_boxes, [(0, 0, 0, 0), (0, 0, 1, 1), (1, 1, 0, 0), (1, 1, 1, 1)]) + ], + "num_rows": 2, + "num_cols": 2, + } + polys = [_rot_poly(*box, 20) for box in word_boxes] + page = doc_builder(*_word_args(["Name", "Age", "Alice", "30"], polys), tables=[[table]]).pages[0] + assert len(page.tables) == 1 and len(page.tables[0].geometry) == 4 + assert all(len(cell.geometry) == 4 for cell in page.tables[0].cells) + assert page.tables[0].to_grid() == [["Name", "Age"], ["Alice", "30"]] + # The export is still JSON-serializable + json.dumps(page.export()) + + +@pytest.mark.parametrize( + "words, boxes, grids, tables", + [ + # A table prediction with no cells (e.g. a false-positive "Table" region where the table model finds nothing) + # is skipped: the words stay in the regular blocks since they were never consumed by a table + [["hello"], [[0.1, 0.1, 0.2, 0.2]], [], [{"cells": [], "num_rows": 0, "num_cols": 0}]], + [["body", "text"], [[0.10, 0.10, 0.25, 0.15], [0.30, 0.10, 0.45, 0.15]], [], [{"cells": []}]], + # A detected table grid on a page without any recognized word + [[], np.zeros((0, 4)), [[[""]]], [{"cells": [_table_cell(0.10, 0.10, 0.30, 0.20, 0, 0, 0, 0)]}]], + ], +) +def test_documentbuilder_tables_without_content(words, boxes, grids, tables): + page = builder.DocumentBuilder()(*_word_args(words, boxes, shape=(100, 100)), tables=[tables]).pages[0] + assert [table.to_grid() for table in page.tables] == grids + assert _words(page) == words + if not words: + assert page.blocks == [] and page.render() == "" def test_documentbuilder_keep_reading_order(): # Two columns of 3 lines each: a naive top-down sort interleaves the columns - left = [[0.1, 0.1 + 0.2 * idx, 0.3, 0.2 + 0.2 * idx] for idx in range(3)] - right = [[0.6, 0.1 + 0.2 * idx, 0.8, 0.2 + 0.2 * idx] for idx in range(3)] - boxes = np.asarray(left + right) - words = [(f"L{idx}", 0.9) for idx in range(3)] + [(f"R{idx}", 0.9) for idx in range(3)] - crop_orientations = [{"value": 0, "confidence": None}] * len(words) - args = ( - [np.zeros((100, 100, 3), dtype=np.uint8)], - [boxes], - [np.ones(len(words))], - [words], - [(100, 100)], - [crop_orientations], - ) - - doc = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True)(*args) - page = doc.pages[0] - assert page.render(block_break=" ").split() == ["L0", "L1", "L2", "R0", "R1", "R2"] + words = ["L0", "L1", "L2", "R0", "R1", "R2"] + boxes = [[x, 0.1 + 0.2 * idx, x + 0.2, 0.2 + 0.2 * idx] for x in (0.1, 0.6) for idx in range(3)] + args = _word_args(words, boxes, shape=(100, 100)) + page = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True)(*args).pages[0] + assert page.render(block_break=" ").split() == words # The flag reorders the stored blocks themselves - assert [word.value for block in page.blocks for line in block.lines for word in line.words] == [ - "L0", - "L1", - "L2", - "R0", - "R1", - "R2", - ] - # Without the flag, the stored blocks keep their original (interleaved) order + assert _words(page) == words + # Without the flag, the stored blocks keep their original (interleaved) order, but the exports are ordered page = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=False)(*args).pages[0] - assert [word.value for block in page.blocks for line in block.lines for word in line.words] != [ - "L0", - "L1", - "L2", - "R0", - "R1", - "R2", - ] - assert page.render(block_break=" ").split() == ["L0", "L1", "L2", "R0", "R1", "R2"] - exported = [word["value"] for b in page.export()["blocks"] for ln in b["lines"] for word in ln["words"]] - assert exported == ["L0", "L1", "L2", "R0", "R1", "R2"] + assert _words(page) != words + assert page.render(block_break=" ").split() == words + assert [word["value"] for b in page.export()["blocks"] for ln in b["lines"] for word in ln["words"]] == words # Layout regions are used to place page furniture: the top line is labeled as a page footer -> emitted last - boxes = np.asarray([[0.1, 0.05, 0.9, 0.1], [0.1, 0.3, 0.9, 0.4], [0.1, 0.5, 0.9, 0.6]]) - words = [("footer", 0.9), ("first", 0.9), ("second", 0.9)] + boxes = [[0.1, 0.05, 0.9, 0.1], [0.1, 0.3, 0.9, 0.4], [0.1, 0.5, 0.9, 0.6]] regions = {"boxes": np.asarray([[0.05, 0.02, 0.95, 0.15]]), "class_names": ["Page-footer"], "scores": [0.9]} - doc = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True)( - [np.zeros((100, 100, 3), dtype=np.uint8)], - [boxes], - [np.ones(3)], - [words], - [(100, 100)], - [[{"value": 0, "confidence": None}] * 3], - regions=[regions], - ) - assert doc.pages[0].render(block_break=" ").split() == ["first", "second", "footer"] + doc_builder = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True) + page = doc_builder(*_word_args(["footer", "first", "second"], boxes, shape=(100, 100)), regions=[regions]).pages[0] + assert page.render(block_break=" ").split() == ["first", "second", "footer"] + # The stored order, render() and Markdown agree: no conflicting re-sort across the exports + boxes = [[x, 0.15 + 0.03 * r, x + 0.37, 0.17 + 0.03 * r] for x in (0.08, 0.55) for r in range(3)] + boxes.append([0.3, 0.05, 0.7, 0.08]) # a footer, geometrically at the top + regions = { + "boxes": np.asarray([[0.06, 0.13, 0.47, 0.26], [0.53, 0.13, 0.94, 0.26], [0.25, 0.03, 0.75, 0.10]]), + "class_names": ["Text", "Text", "Page-footer"], + "scores": [0.9, 0.9, 0.9], + } + expected = [*words, "foot"] + page = doc_builder(*_word_args(expected, boxes, shape=(100, 100)), regions=[regions]).pages[0] + assert _words(page) == expected + # render() follows the stored order and is idempotent + assert page.render(block_break=" ").split() == expected == page.render(block_break=" ").split() + # markdown re-derives reading order from the lines and lands on the same order + assert [token for token in page.export_as_markdown().replace("#", " ").split() if token in expected] == expected -def test_documentbuilder_keep_reading_order_keeps_figures(): # The builder stores its linearization back in the page cache: the figures must be kept in it - image = np.zeros((100, 100, 3), dtype=np.uint8) - image[30:60, 10:90] = 255 - boxes = np.asarray([[0.1, 0.05, 0.9, 0.1], [0.3, 0.4, 0.5, 0.45], [0.2, 0.65, 0.8, 0.7], [0.1, 0.8, 0.9, 0.85]]) - words = [("intro", 0.9), ("inner", 0.9), ("caption", 0.9), ("outro", 0.9)] + args = _word_args( + ["intro", "inner", "caption", "outro"], + [[0.1, 0.05, 0.9, 0.1], [0.3, 0.4, 0.5, 0.45], [0.2, 0.65, 0.8, 0.7], [0.1, 0.8, 0.9, 0.85]], + shape=(100, 100), + ) + args[0][0][30:60, 10:90] = 255 regions = { "boxes": np.asarray([[0.08, 0.28, 0.92, 0.62], [0.15, 0.63, 0.85, 0.72]]), "class_names": ["Picture", "Caption"], "scores": [0.9, 0.9], } - args = ( - [image], - [boxes], - [np.ones(4)], - [words], - [(100, 100)], - [[{"value": 0, "confidence": None}] * 4], - ) exports = {} for keep in (False, True): page = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=keep)(*args, regions=[regions]).pages[0] @@ -873,67 +735,11 @@ def test_documentbuilder_keep_reading_order_keeps_figures(): assert "inner" not in exports[True] -def _rot_poly(x0, y0, x1, y1, deg, cx=0.5, cy=0.5): - a = np.deg2rad(deg) - ca, sa = np.cos(a), np.sin(a) - return [ - [cx + (x - cx) * ca - (y - cy) * sa, cy + (x - cx) * sa + (y - cy) * ca] - for x, y in [(x0, y0), (x1, y0), (x1, y1), (x0, y1)] - ] - - -def test_documentbuilder_keep_reading_order_rotated(): - deg = 25 - height, width = 1000, 800 - angle = np.deg2rad(deg) - rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) - center = np.array([width / 2, height / 2]) - - def _rot_box(x0, y0, x1, y1): - pts = np.array([ - [x0 * width, y0 * height], - [x1 * width, y0 * height], - [x1 * width, y1 * height], - [x0 * width, y1 * height], - ]) - return ((pts - center) @ rot.T + center) / [width, height] - - left = [_rot_box(0.1, 0.1 + 0.2 * idx, 0.3, 0.2 + 0.2 * idx) for idx in range(3)] - right = [_rot_box(0.6, 0.1 + 0.2 * idx, 0.8, 0.2 + 0.2 * idx) for idx in range(3)] - polys = np.asarray(left + right, dtype=np.float32) - words = [(f"L{idx}", 0.9) for idx in range(3)] + [(f"R{idx}", 0.9) for idx in range(3)] - doc = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True)( - [np.zeros((height, width, 3), dtype=np.uint8)], - [polys], - [np.ones(len(words))], - [words], - [(height, width)], - [[{"value": 0, "confidence": None}] * len(words)], +@pytest.mark.parametrize("resolve_lines, resolve_blocks", [(True, True), (True, False), (False, True), (False, False)]) +def test_documentbuilder_keep_reading_order_flags(resolve_lines, resolve_blocks): + doc_builder = builder.DocumentBuilder( + resolve_lines=resolve_lines, resolve_blocks=resolve_blocks, keep_reading_order=True ) - assert doc.pages[0].render(block_break=" ").split() == ["L0", "L1", "L2", "R0", "R1", "R2"] - - -_FLAG_COMBINATIONS = [ - (True, True), - (True, False), - (False, True), - (False, False), -] - - -def _run(doc_builder, boxes, words, regions=None): - return doc_builder( - [np.zeros((100, 100, 3), dtype=np.uint8)], - [np.asarray(boxes, dtype=np.float32)], - [np.ones(len(words))], - [words], - [(100, 100)], - [[{"value": 0, "confidence": None}] * len(words)], - regions=[regions] if regions is not None else None, - ).pages[0] - - -def _two_columns_of_words(): # Two columns, three rows each, two words per row; shuffled so the input order is not the reading order cells = [] for column, tag in ((0.08, "L"), (0.55, "R")): @@ -941,265 +747,67 @@ def _two_columns_of_words(): y = 0.12 + 0.04 * row cells.append(([column, y, column + 0.15, y + 0.02], f"{tag}{row}a")) cells.append(([column + 0.17, y, column + 0.34, y + 0.02], f"{tag}{row}b")) - order = [7, 0, 11, 3, 5, 9, 1, 6, 10, 2, 8, 4] - cells = [cells[idx] for idx in order] - boxes = [box for box, _ in cells] - words = [(text, 0.9) for _, text in cells] - expected = "L0a L0b L1a L1b L2a L2b R0a R0b R1a R1b R2a R2b".split() - return boxes, words, expected - - -def test_documentbuilder_keep_reading_order_all_flag_combinations(): - boxes, words, expected = _two_columns_of_words() - # two column text regions - regions = { + cells = [cells[idx] for idx in [7, 0, 11, 3, 5, 9, 1, 6, 10, 2, 8, 4]] + args = _word_args([text for _, text in cells], [box for box, _ in cells], shape=(100, 100)) + columns = { "boxes": np.asarray([[0.06, 0.10, 0.45, 0.24], [0.53, 0.10, 0.92, 0.24]]), "class_names": ["Text", "Text"], "scores": [0.9, 0.9], } - for resolve_lines, resolve_blocks in _FLAG_COMBINATIONS: - for page_regions in (None, regions): - page = _run( - builder.DocumentBuilder( - resolve_lines=resolve_lines, resolve_blocks=resolve_blocks, keep_reading_order=True - ), - boxes, - words, - page_regions, - ) - assert page.render(block_break=" ").split() == expected - - -def test_documentbuilder_keep_reading_order_without_resolve_lines(): - # resolve_lines=False still reads row by row (lines are resolved internally only to order the content) - boxes, words, expected = _two_columns_of_words() - page = _run( - builder.DocumentBuilder(resolve_lines=False, resolve_blocks=False, keep_reading_order=True), boxes, words - ) - assert page.render(block_break=" ").split() == expected - # each column collapses onto a single line since lines are not exposed - assert all(len(block.lines) == 1 for block in page.blocks) - - -def test_documentbuilder_keep_reading_order_prefers_layout_furniture(): + for regions in (None, columns): + page = doc_builder(*args, regions=None if regions is None else [regions]).pages[0] + assert page.render(block_break=" ").split() == "L0a L0b L1a L1b L2a L2b R0a R0b R1a R1b R2a R2b".split() + if not resolve_lines and not resolve_blocks and regions is None: + # lines are resolved internally only to order the content: each column collapses onto a single line + assert all(len(block.lines) == 1 for block in page.blocks) + + # Layout regions place the page furniture, whatever its position on the page boxes = [ [0.1, 0.05, 0.9, 0.09], # footer text (near the top of the page) [0.1, 0.30, 0.9, 0.38], # first body line [0.1, 0.45, 0.9, 0.53], # second body line [0.1, 0.15, 0.9, 0.20], # header text ] - words = [("foot", 0.9), ("first", 0.9), ("second", 0.9), ("head", 0.9)] regions = { "boxes": np.asarray([[0.05, 0.03, 0.95, 0.11], [0.05, 0.13, 0.95, 0.22]]), "class_names": ["Page-footer", "Page-header"], "scores": [0.9, 0.9], } - for resolve_lines, resolve_blocks in _FLAG_COMBINATIONS: - page = _run( - builder.DocumentBuilder( - resolve_lines=resolve_lines, resolve_blocks=resolve_blocks, keep_reading_order=True - ), - boxes, - words, - regions, - ) - assert page.render(block_break=" ").split() == ["head", "first", "second", "foot"] + page = doc_builder(*_word_args(["foot", "first", "second", "head"], boxes, shape=(100, 100)), regions=[regions]) + assert page.pages[0].render(block_break=" ").split() == ["head", "first", "second", "foot"] + +def test_documentbuilder_rotated_pages(): + # Two columns of words rotated by 25 degrees around the center of a non-square page are read column by column + height, width = 1000, 800 + angle = np.deg2rad(25) + rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) + center = np.array([width / 2, height / 2]) -def test_documentbuilder_keep_reading_order_without_resolve_lines_rotated(): - angle = np.deg2rad(20.0) - ca, sa = np.cos(angle), np.sin(angle) + def _rot_box(x0, y0, x1, y1): + pts = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) * [width, height] + return ((pts - center) @ rot.T + center) / [width, height] - def _rot(x0, y0, x1, y1, cx=0.5, cy=0.5): - return [ - [cx + (x - cx) * ca - (y - cy) * sa, cy + (x - cx) * sa + (y - cy) * ca] - for x, y in [(x0, y0), (x1, y0), (x1, y1), (x0, y1)] - ] + words = ["L0", "L1", "L2", "R0", "R1", "R2"] + polys = [_rot_box(x, 0.1 + 0.2 * idx, x + 0.2, 0.2 + 0.2 * idx) for x in (0.1, 0.6) for idx in range(3)] + doc_builder = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True) + page = doc_builder(*_word_args(words, polys, shape=(height, width))).pages[0] + assert page.render(block_break=" ").split() == words + # Same without resolve_lines, for two words per row cells = [] for column, tag in ((0.08, "L"), (0.55, "R")): for row in range(3): y = 0.15 + 0.04 * row - cells.append((_rot(column, y, column + 0.15, y + 0.02), f"{tag}{row}a")) - cells.append((_rot(column + 0.17, y, column + 0.34, y + 0.02), f"{tag}{row}b")) - boxes = np.asarray([box for box, _ in cells], dtype=np.float32) - words = [(text, 0.9) for _, text in cells] - - page = _run( - builder.DocumentBuilder(resolve_lines=False, resolve_blocks=False, keep_reading_order=True), boxes, words - ) + cells.append((_rot_poly(column, y, column + 0.15, y + 0.02, 20), f"{tag}{row}a")) + cells.append((_rot_poly(column + 0.17, y, column + 0.34, y + 0.02, 20), f"{tag}{row}b")) + doc_builder = builder.DocumentBuilder(resolve_lines=False, resolve_blocks=False, keep_reading_order=True) + page = doc_builder(*_word_args([text for _, text in cells], [box for box, _ in cells], shape=(100, 100))).pages[0] assert page.render(block_break=" ").split() == "L0a L0b L1a L1b L2a L2b R0a R0b R1a R1b R2a R2b".split() - -def test_documentbuilder_keep_reading_order_no_double_sort_across_exports(): - boxes = [[x, 0.15 + 0.03 * r, x + 0.37, 0.17 + 0.03 * r] for x in (0.08, 0.55) for r in range(3)] - boxes.append([0.3, 0.05, 0.7, 0.08]) # a footer, geometrically at the top - words = [("L0", 0.9), ("L1", 0.9), ("L2", 0.9), ("R0", 0.9), ("R1", 0.9), ("R2", 0.9), ("foot", 0.9)] - regions = { - "boxes": np.asarray([[0.06, 0.13, 0.47, 0.26], [0.53, 0.13, 0.94, 0.26], [0.25, 0.03, 0.75, 0.10]]), - "class_names": ["Text", "Text", "Page-footer"], - "scores": [0.9, 0.9, 0.9], - } - page = _run(builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=True), boxes, words, regions) - - expected = ["L0", "L1", "L2", "R0", "R1", "R2", "foot"] - block_order = [word.value for block in page.blocks for line in block.lines for word in line.words] - assert block_order == expected - # render() follows the stored order and is idempotent (no conflicting re-sort) - assert page.render(block_break=" ").split() == expected - assert page.render(block_break=" ").split() == page.render(block_break=" ").split() - # markdown re-derives reading order from the lines and lands on the same order - markdown_words = [token for token in page.export_as_markdown().replace("#", " ").split() if token in expected] - assert markdown_words == expected - - -def _word_args(words, boxes, shape=(1000, 1000)): - """Pack builder arguments for a single page from (value, box) pairs""" - return ( - [np.zeros((*shape, 3), dtype=np.uint8)], - [np.asarray(boxes, dtype=np.float32)], - [np.full(len(words), 0.9, dtype=np.float32)], - [[(value, 0.95) for value in words]], - [shape], - [[{"value": 0, "confidence": None}] * len(words)], - ) - - -def test_documentbuilder_tables_multi_word_cells(): - words = ["world", "hello", "second", "line", "solo"] - boxes = [ - [0.20, 0.13, 0.28, 0.16], # "world" - first line, right - [0.11, 0.13, 0.19, 0.16], # "hello" - first line, left - [0.11, 0.17, 0.20, 0.20], # "second" - second line, left - [0.21, 0.17, 0.27, 0.20], # "line" - second line, right - [0.36, 0.13, 0.44, 0.16], # "solo" - other cell - ] - table = { - "cells": [ - _table_cell(0.10, 0.10, 0.30, 0.22, 0, 0, 0, 0), - _table_cell(0.32, 0.10, 0.50, 0.22, 0, 0, 1, 1), - ], - "num_rows": 1, - "num_cols": 2, - } - out = builder.DocumentBuilder(resolve_lines=True)(*_word_args(words, boxes), tables=[[table]]) - page = out.pages[0] - assert len(page.tables) == 1 - assert page.tables[0].to_grid() == [["hello world second line", "solo"]] - # Every word ended up in the table, so nothing is left in the regular blocks - assert page.blocks == [] - # A degenerate cell whose words all share the same center still produces a single row - words = ["a", "b"] - boxes = [[0.15, 0.15, 0.15, 0.15], [0.16, 0.15, 0.16, 0.15]] - out = builder.DocumentBuilder(resolve_lines=True)( - *_word_args(words, boxes), - tables=[[{"cells": [_table_cell(0.10, 0.10, 0.30, 0.22, 0, 0, 0, 0)], "num_rows": 1, "num_cols": 1}]], - ) - assert out.pages[0].tables[0].to_grid() == [["a b"]] - - -def test_documentbuilder_tables_nearest_cell_fallback(): - words = ["left", "right", "outside", "faraway"] - boxes = [ - [0.11, 0.12, 0.19, 0.16], # inside cell 0 - [0.33, 0.12, 0.41, 0.16], # inside cell 1 - [0.305, 0.12, 0.315, 0.16], # in the gutter, closest to cell 1 - [0.11, 0.80, 0.25, 0.85], # far outside the table - ] - table = { - "cells": [ - _table_cell(0.10, 0.10, 0.30, 0.18, 0, 0, 0, 0), - _table_cell(0.32, 0.10, 0.50, 0.18, 0, 0, 1, 1), - ], - "num_rows": 1, - "num_cols": 2, - } - page = builder.DocumentBuilder(resolve_lines=True)(*_word_args(words, boxes), tables=[[table]]).pages[0] - grid = page.tables[0].to_grid() - assert grid[0][0] == "left" - assert "outside" in grid[0][1] - # The word far away from the table is untouched and stays in the regular blocks - body = [word.value for block in page.blocks for line in block.lines for word in line.words] - assert body == ["faraway"] - - -def test_documentbuilder_tables_rotated_geometry(): - deg = 20 - words = ["Name", "Age", "Alice", "30"] - cell_boxes = [ - (0.10, 0.10, 0.28, 0.20), - (0.30, 0.10, 0.48, 0.20), - (0.10, 0.22, 0.28, 0.32), - (0.30, 0.22, 0.48, 0.32), - ] - word_boxes = [ - (0.13, 0.13, 0.25, 0.17), - (0.33, 0.13, 0.45, 0.17), - (0.13, 0.25, 0.25, 0.29), - (0.33, 0.25, 0.45, 0.29), - ] - polys = np.asarray([_rot_poly(*box, deg) for box in word_boxes], dtype=np.float32) # (N, 4, 2) - table = { - "cells": [ - {**_table_cell(0, 0, 0, 0, rs, re, cs, ce), "geometry": _rot_poly(*box, deg)} - for box, (rs, re, cs, ce) in zip(cell_boxes, [(0, 0, 0, 0), (0, 0, 1, 1), (1, 1, 0, 0), (1, 1, 1, 1)]) - ], - "num_rows": 2, - "num_cols": 2, - } - args = _word_args(words, polys) - out = builder.DocumentBuilder(resolve_lines=True, export_as_straight_boxes=False)(*args, tables=[[table]]) - page = out.pages[0] - assert len(page.tables) == 1 - table_out = page.tables[0] - # Rotated pages keep 4-point polygons for the table and its cells - assert len(table_out.geometry) == 4 - assert all(len(cell.geometry) == 4 for cell in table_out.cells) - assert table_out.to_grid() == [["Name", "Age"], ["Alice", "30"]] - # The export is still JSON-serializable - json.dumps(page.export()) - - -def test_documentbuilder_tables_without_words(): - # A detected table grid on a page without any recognized word must not raise - table = {"cells": [_table_cell(0.10, 0.10, 0.30, 0.20, 0, 0, 0, 0)], "num_rows": 1, "num_cols": 1} - out = builder.DocumentBuilder(resolve_lines=True)( - [np.zeros((100, 100, 3), dtype=np.uint8)], - [np.zeros((0, 4), dtype=np.float32)], - [np.zeros(0, dtype=np.float32)], - [[]], - [(100, 100)], - [[]], - tables=[[table]], - ) - page = out.pages[0] - assert len(page.tables) == 1 and page.tables[0].to_grid() == [[""]] - assert page.blocks == [] and page.render() == "" - - -def test_documentbuilder_sort_boxes_mixed_rotation(): - deg = 30 - rotated = [_rot_poly(0.10, 0.10 + 0.12 * idx, 0.60, 0.16 + 0.12 * idx, deg) for idx in range(3)] - axis_aligned = [ - [[0.10, 0.70], [0.60, 0.70], [0.60, 0.76], [0.10, 0.76]], # not rotated at all - ] - polys = np.asarray(rotated + axis_aligned, dtype=np.float32) - words = ["r0", "r1", "r2", "flat"] - out = builder.DocumentBuilder(resolve_lines=True, export_as_straight_boxes=False)( - *_word_args(words, polys, shape=(1000, 1000)) - ) - page = out.pages[0] - rendered = [word.value for block in page.blocks for line in block.lines for word in line.words] - assert sorted(rendered) == ["flat", "r0", "r1", "r2"] + # Words rotated by 30 degrees mixed with an axis-aligned one + polys = [_rot_poly(0.10, 0.10 + 0.12 * idx, 0.60, 0.16 + 0.12 * idx, 30) for idx in range(3)] + polys.append([[0.10, 0.70], [0.60, 0.70], [0.60, 0.76], [0.10, 0.76]]) # not rotated at all + page = builder.DocumentBuilder()(*_word_args(["r0", "r1", "r2", "flat"], polys)).pages[0] + assert sorted(_words(page)) == ["flat", "r0", "r1", "r2"] json.dumps(page.export()) - - -def test_documentbuilder_tables_skips_empty_grids(): - # A detected table region whose structure prediction has no cell at all is skipped entirely - words = ["body", "text"] - boxes = [[0.10, 0.10, 0.25, 0.15], [0.30, 0.10, 0.45, 0.15]] - grids = [{"cells": [], "num_rows": 0, "num_cols": 0}] - page = builder.DocumentBuilder(resolve_lines=True)(*_word_args(words, boxes), tables=[grids]).pages[0] - assert page.tables == [] - rendered = [word.value for block in page.blocks for line in block.lines for word in line.words] - assert rendered == words diff --git a/tests/common/test_models_reading_order.py b/tests/common/test_models_reading_order.py index f0ad29258d..bc458397a8 100644 --- a/tests/common/test_models_reading_order.py +++ b/tests/common/test_models_reading_order.py @@ -11,7 +11,138 @@ resolve_reading_segments, sort_reading_order, ) -from doctr.models.reading_order.base import _region_groups +from doctr.models.reading_order.base import _defer_floats, _region_groups + + +def _box(x0, y0, x1, y1): + return ((x0, y0), (x1, y1)) + + +def _column(x0, x1, y0, step, count, height=0.03): + return [_box(x0, y0 + step * i, x1, y0 + step * i + height) for i in range(count)] + + +def _rotated_box(box, deg, width=800, height=1000): + angle = np.deg2rad(deg) + rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) + center = np.array([width / 2, height / 2]) + (x0, y0), (x1, y1) = box + pts = np.array([[x0 * width, y0 * height], [x1 * width, y0 * height], [x1 * width, y1 * height]]) + pts = np.concatenate([pts, [[x0 * width, y1 * height]]]) + return ((pts - center) @ rot.T + center) / [width, height] + + +# Two columns of 4 lines +TWO_COLUMNS = _column(0.1, 0.45, 0.1, 0.2, 4, 0.15) + _column(0.55, 0.9, 0.1, 0.2, 4, 0.15) + +# Page furniture, a figure with its caption, body text and a footnote +LABELED_PAGE = [ + _box(0.1, 0.92, 0.9, 0.97), # 0: page footer + _box(0.1, 0.02, 0.9, 0.06), # 1: page header + _box(0.1, 0.1, 0.9, 0.4), # 2: body text + _box(0.1, 0.45, 0.5, 0.7), # 3: figure + _box(0.1, 0.71, 0.5, 0.75), # 4: caption below the figure + _box(0.55, 0.45, 0.9, 0.88), # 5: body on the right of the figure + _box(0.1, 0.8, 0.5, 0.84), # 6: footnote +] +PAGE_LABELS = ["Page-footer", "Page-header", "Text", "Picture", "Caption", "Text", "Footnote"] + +# A page-wide figure with text inside, its caption, a figure in the left column with its caption, two columns of text +FIGURES_PAGE = [ + _box(0.10, 0.05, 0.90, 0.30), # 0: figure (page wide) + _box(0.15, 0.10, 0.30, 0.12), # 1: text inside the figure, left side + _box(0.70, 0.10, 0.85, 0.12), # 2: text inside the figure, right side + _box(0.10, 0.32, 0.90, 0.35), # 3: its caption + _box(0.10, 0.40, 0.45, 0.70), # 4: figure in the left column + _box(0.10, 0.72, 0.45, 0.80), # 5: its caption + *_column(0.55, 0.90, 0.40, 0.05, 8), # 6-13: right column + *_column(0.10, 0.45, 0.82, 0.04, 3), # 14-16: left column +] +FIGURES_LABELS = ["Picture", "Picture", "Picture", "Caption", "Picture", "Caption"] + ["Text"] * 11 + +# Lines of a real receipt (CORD): the right-aligned amounts are level with their item, far from it +RECEIPT = [ + _box(0.096, 0.409, 0.478, 0.439), # 0: Enting2 Dua Bandeng Ori + _box(0.650, 0.407, 0.753, 0.431), # 1: 31.000 + _box(0.094, 0.444, 0.178, 0.467), # 2: 01069 + _box(0.243, 0.440, 0.411, 0.469), # 3: 1 31.000 + _box(0.094, 0.471, 0.561, 0.500), # 4: TAHU BAKSO Grg Ayam VC [biji + _box(0.651, 0.470, 0.754, 0.492), # 5: 68.000 + _box(0.091, 0.506, 0.174, 0.529), # 6: 07165 + _box(0.243, 0.504, 0.411, 0.529), # 7: 10 6.800 + _box(0.087, 0.564, 0.430, 0.593), # 8: QTY 11 ITEMS 2 + _box(0.088, 0.599, 0.170, 0.622), # 9: TOTAL + _box(0.635, 0.597, 0.741, 0.619), # 10: 99.000 + _box(0.086, 0.629, 0.172, 0.655), # 11: KARTU + _box(0.637, 0.629, 0.742, 0.652), # 12: 99.000 +] + +# 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 +CRUMPLED_TOTALS = [_box(60, 1200 + 60 * i, 600, 1240 + 60 * i) for i in range(3)] + [ + _box(114, 1376, 282, 1420), # 3: label + _box(420, 1388, 520, 1431), # 4: its amount + _box(112, 1408, 226, 1451), # 5: label + _box(420, 1422, 520, 1465), # 6: its amount + _box(110, 1445, 162, 1483), # 7: label + _box(410, 1456, 520, 1502), # 8: its amount + _box(110, 1481, 162, 1518), # 9: last label + *[_box(60, 1560 + 60 * i, 600, 1600 + 60 * i) for i in range(2)], # 10-11: full-width lines +] + +# A 3-line caption whose last line is closer to the table of the other column +SPLIT_CAPTION = [ + _box(0.526, 0.115, 0.903, 0.33), # 0: figure + _box(0.102, 0.248, 0.477, 0.443), # 1: table + _box(0.52, 0.352, 0.91, 0.366), # 2: caption line 1 + _box(0.52, 0.37, 0.91, 0.384), # 3: caption line 2 + _box(0.52, 0.419, 0.73, 0.434), # 4: caption line 3 +] +SPLIT_CAPTION_LABELS = ["Picture", "Table", "Caption", "Caption", "Caption"] + + +def _paragraph_beside_a_figure(): + """A single-column page (3-column author grid) with a paragraph ending one line below the figure beside it""" + boxes, labels, groups = [[0.1, 0.05, 0.9, 0.09]], ["Title"], [0] + for row in range(4): + for x in (0.15, 0.41, 0.68): + boxes.append([x, 0.14 + row * 0.03, x + 0.18, 0.154 + row * 0.03]) + labels.append(None) + groups.append(-1) + for y in np.linspace(0.30, 0.64, 12): + boxes.append([0.08, y, 0.48, y + 0.014]) + labels.append("Text") + groups.append(1) + boxes += [[0.08, 0.662, 0.35, 0.676], [0.52, 0.31, 0.91, 0.667], [0.08, 0.70, 0.30, 0.712]] + labels += ["Text", "Picture", "Section-header"] # last line of the paragraph (below the figure), figure, header + groups += [1, 2, 3] + return [_box(*box) for box in boxes], labels, groups + + +BESIDE, BESIDE_LABELS, BESIDE_GROUPS = _paragraph_beside_a_figure() +LAST_LINE, FIGURE = len(BESIDE) - 3, len(BESIDE) - 2 + +# Two text columns on the right half and a wide figure on the left that does not vote +VOTERS_PAGE = [ + box for y in np.arange(6) * 0.05 + 0.1 for box in (_box(0.55, y, 0.72, y + 0.02), _box(0.78, y, 0.95, y + 0.02)) +] +VOTERS_PAGE.append(_box(0.02, 0.1, 0.5, 0.4)) + + +def _flat(segments): + return [idx for segment in segments for idx in segment] + + +def _before(*groups): + """Check that each group of indices is read entirely before the next one""" + + def check(order): + return all( + max(order.index(i) for i in group) < min(order.index(i) for i in nxt) + for group, nxt in zip(groups, groups[1:]) + ) + + return check @pytest.mark.parametrize( @@ -37,296 +168,263 @@ def test_detect_text_direction(texts, language, expected): assert detect_text_direction(texts, language=language) == expected -def test_normalize_layout_label(): - assert normalize_layout_label("Page-header") == "page_header" - assert normalize_layout_label(" Section header ") == "section_header" - assert normalize_layout_label(None) == "" - - @pytest.mark.parametrize( - "label, role", + "label, normalized, role", [ - ("Page-header", "header"), - ("Page-footer", "footer"), - ("Footnote", "footnote"), - ("Caption", "caption"), - ("Table", "float"), - ("Picture", "float"), - ("Text", "body"), - ("Title", "body"), - (None, "body"), + ("Page-header", "page_header", "header"), + (" Section header ", "section_header", "body"), + ("Page-footer", "page_footer", "footer"), + ("Footnote", "footnote", "footnote"), + ("Caption", "caption", "caption"), + ("Table", "table", "float"), + ("Picture", "picture", "float"), + ("Text", "text", "body"), + ("Title", "title", "body"), + (None, "", "body"), ], ) -def test_layout_label_role(label, role): +def test_layout_labels(label, normalized, role): + assert normalize_layout_label(label) == normalized assert layout_label_role(label) == role -def _two_columns(num_lines: int = 4): - left = [((0.1, 0.1 + 0.2 * i), (0.45, 0.25 + 0.2 * i)) for i in range(num_lines)] - right = [((0.55, 0.1 + 0.2 * i), (0.9, 0.25 + 0.2 * i)) for i in range(num_lines)] - return left + right - - -def test_sort_reading_order_basic(): - # Degenerate inputs - assert sort_reading_order([]) == [] - assert sort_reading_order([((0.1, 0.1), (0.2, 0.2))]) == [0] - # Single column, top to bottom - geoms = [((0.1, 0.5), (0.9, 0.6)), ((0.1, 0.1), (0.9, 0.2)), ((0.1, 0.3), (0.9, 0.4))] - assert sort_reading_order(geoms) == [1, 2, 0] - # A title spanning two columns - geoms = [((0.55, 0.2), (0.9, 0.8)), ((0.1, 0.05), (0.9, 0.15)), ((0.1, 0.2), (0.45, 0.8))] - assert sort_reading_order(geoms) == [1, 2, 0] +@pytest.mark.parametrize( + "geoms, kwargs, expected", + [ + pytest.param([], {}, [], id="empty"), + pytest.param([_box(0.1, 0.1, 0.2, 0.2)], {}, [0], id="single"), + pytest.param([_box(0.1, 0.5, 0.9, 0.6), _box(0.1, 0.1, 0.9, 0.2), _box(0.1, 0.3, 0.9, 0.4)], {}, [1, 2, 0]), + pytest.param( + [_box(0.55, 0.2, 0.9, 0.8), _box(0.1, 0.05, 0.9, 0.15), _box(0.1, 0.2, 0.45, 0.8)], + {}, + [1, 2, 0], + id="title-over-columns", + ), + pytest.param(TWO_COLUMNS, {}, list(range(8)), id="columns"), + pytest.param(TWO_COLUMNS, {"direction": "rtl"}, [4, 5, 6, 7, 0, 1, 2, 3], id="columns-rtl"), + pytest.param( + [_box(0.1, 0.02, 0.9, 0.06), _box(0.05, 0.1, 0.95, 0.6)] + + _column(0.1, 0.4, 0.7, 0.05, 3) + + _column(0.6, 0.9, 0.7, 0.05, 3), + {"column_voters": [True, False] + [True] * 6}, + list(range(8)), + id="non-voting-page-wide-figure", + ), + pytest.param( + # A centered title crossing the gutter does not merge the columns: the column of a figure is read to its end + [_box(0.3, 0.02, 0.7, 0.06), _box(0.1, 0.1, 0.45, 0.13), _box(0.1, 0.16, 0.13, 0.19)] + + [_box(0.25, 0.22, 0.3, 0.25)] # figure labels on 2 sides + + _column(0.1, 0.45, 0.28, 0.05, 2) + + _column(0.55, 0.9, 0.1, 0.07, 4), + {}, + list(range(10)), + id="title-crossing-the-gutter", + ), + pytest.param( + [_box(0.8 - 0.15 * i, 0.1, 0.9 - 0.15 * i, 0.9) for i in range(4)], + {"direction": "ttb-rtl"}, + [0, 1, 2, 3], + id="vertical-rtl", + ), + pytest.param( + [_box(0.8 - 0.15 * i, 0.1, 0.9 - 0.15 * i, 0.9) for i in range(4)], + {"direction": "ttb-ltr"}, + [3, 2, 1, 0], + id="vertical-ltr", + ), + pytest.param( + [_box(0.8, 0.5, 0.9, 0.9), _box(0.8, 0.1, 0.9, 0.45), _box(0.6, 0.1, 0.7, 0.9)], + {"direction": "ttb-rtl"}, + [1, 0, 2], + id="vertical-stacked", + ), + pytest.param(LABELED_PAGE, {"labels": PAGE_LABELS}, [1, 2, 3, 4, 5, 6, 0], id="furniture-and-caption"), + pytest.param( + [*LABELED_PAGE[:4], _box(0.1, 0.41, 0.5, 0.44), *LABELED_PAGE[5:]], + {"labels": PAGE_LABELS}, + [1, 2, 4, 3, 5, 6, 0], + id="caption-above", + ), + pytest.param( + # A caption too far from any (shrunk) float keeps its natural position in the body + [*LABELED_PAGE[:3], _box(0.1, 0.45, 0.5, 0.5), _box(0.55, 0.02, 0.9, 0.06), *LABELED_PAGE[5:]], + {"labels": PAGE_LABELS}, + _before([4], [3]), + id="caption-out-of-reach", + ), + pytest.param( + # The text inside a figure is read right after it, before its caption + FIGURES_PAGE, + {"labels": FIGURES_LABELS}, + [0, 1, 2, 3, 4, 5, 14, 15, 16, *range(6, 14)], + id="float-members", + ), + pytest.param( + FIGURES_PAGE, + {"labels": [*FIGURES_LABELS[:2], "Text", *FIGURES_LABELS[3:]]}, + lambda order: order[:4] == [0, 1, 2, 3], + id="float-members-labeled-as-text", + ), + pytest.param(FIGURES_PAGE, {}, lambda order: True, id="float-members-without-labels"), + pytest.param( + [_box(0.1, 0.1, 0.1, 0.1)] * 3 + [_box(0.5, 0.5, 0.5, 0.5)], + {}, + lambda order: sorted(order) == list(range(4)), + id="degenerate", + ), + pytest.param( + # A line split into a narrow part and a stray fragment does not interrupt its column + _column(0.10, 0.45, 0.10, 0.04, 1) + + [_box(0.10, 0.14, 0.25, 0.17), _box(0.34, 0.14, 0.45, 0.17)] + + _column(0.10, 0.45, 0.18, 0.04, 3) + + _column(0.55, 0.90, 0.10, 0.04, 6), + {}, + _before(range(6), range(6, 12)), + id="fragmented-columns", + ), + pytest.param( + [_box(0.35, 0.05, 0.65, 0.10), _box(0.10, 0.15, 0.45, 0.20)] # element bridging both columns + + [_box(0.10, 0.22, 0.16, 0.27), _box(0.17, 0.22, 0.24, 0.27), _box(0.25, 0.22, 0.45, 0.27)] # 3 fragments + + [_box(0.10, 0.29, 0.45, 0.34), _box(0.55, 0.15, 0.90, 0.20), _box(0.55, 0.22, 0.90, 0.27)], + {}, + list(range(8)), + id="fragmented-row", + ), + pytest.param( + _column(0.05, 0.95, 0.02, 0.06, 3, 0.04) + + [ + box + for y in (0.20, 0.26, 0.32) + for box in (_box(0.05, y, 0.30, y + 0.04), _box(0.65, y, 0.95, y + 0.04)) + ] + + _column(0.05, 0.95, 0.38, 0.06, 2, 0.04), + {}, + list(range(11)), + id="key-value-rows", + ), + pytest.param(CRUMPLED_TOTALS, {}, list(range(12)), id="crumpled-receipt-totals"), + pytest.param( + # The cells of a single row are still read left to right when the middle one sits a little lower + [_box(100, 100, 200, 120), _box(300, 103, 400, 123), _box(500, 100, 600, 120)], + {}, + [0, 1, 2], + id="shifted-cell", + ), + pytest.param( + [_box(100, 100, 200, 120), _box(300, 103, 400, 123), _box(500, 100, 600, 120)] + + [_box(100, 140, 200, 160), _box(300, 140, 400, 160)], + {}, + list(range(5)), + id="shifted-cell-next-row", + ), + pytest.param(RECEIPT, {}, list(range(13)), id="receipt-amounts"), + pytest.param( + [_box(1 - x1, y0, 1 - x0, y1) for (x0, y0), (x1, y1) in RECEIPT], + {"direction": "rtl"}, + list(range(13)), + id="receipt-amounts-rtl", + ), + pytest.param( + # Text columns with the same gutter are still read column after column: their lines are long + _column(0.05, 0.40, 0.10, 0.03, 4, 0.02) + _column(0.55, 0.90, 0.10, 0.03, 4, 0.02), + {}, + list(range(8)), + id="columns-with-a-wide-gutter", + ), + pytest.param( + # So is a column of short lines on a page with many lines (e.g. a list beside a text column) + _column(0.05, 0.40, 0.02, 0.012, 40, 0.01) + _column(0.55, 0.60, 0.02, 0.012, 40, 0.01), + {}, + list(range(80)), + id="short-lines-column", + ), + pytest.param( + [_box(0.05, 0.02, 0.95, 0.06)] + + _column(0.05, 0.45, 0.10, 0.06, 3, 0.04) + + _column(0.55, 0.95, 0.10, 0.06, 3, 0.04), + {}, + list(range(7)), + id="header-over-columns", + ), + pytest.param( + SPLIT_CAPTION, + {"labels": SPLIT_CAPTION_LABELS}, + lambda order: order.index(4) == order.index(1) + 1, + id="split-caption", + ), + pytest.param( + # Grouped, the caption lines follow the float closest to their union + SPLIT_CAPTION, + {"labels": SPLIT_CAPTION_LABELS, "region_groups": [-1, -1, 7, 7, 7]}, + lambda order: order[order.index(0) + 1 : order.index(0) + 4] == [2, 3, 4], + id="grouped-caption", + ), + pytest.param( + # Without region groups, the figure splits the paragraph beside it... + BESIDE, + {"labels": BESIDE_LABELS}, + _before([FIGURE], [LAST_LINE]), + id="float-splits-a-region", + ), + pytest.param( + # ... but not with them + BESIDE, + {"labels": BESIDE_LABELS, "region_groups": BESIDE_GROUPS}, + lambda order: ( + _before([LAST_LINE], [FIGURE], [len(BESIDE) - 1])(order) and sorted(order) == list(range(len(BESIDE))) + ), + id="float-beside-a-region", + ), + pytest.param( + # A text region enclosing a figure: the figure stays between the lines around it + _column(0.1, 0.9, 0.1, 0.04, 3) + [_box(0.2, 0.3, 0.8, 0.6)] + _column(0.1, 0.9, 0.7, 0.04, 3), + {"labels": ["Text"] * 3 + ["Picture"] + ["Text"] * 3, "region_groups": [0, 0, 0, 1, 0, 0, 0]}, + list(range(7)), + id="float-stacked-in-a-loose-region", + ), + pytest.param( + # A caption out of reach of any float is read in its own column + _column(0.08, 0.46, 0.1, 0.05, 8) + + [_box(0.54, 0.1, 0.92, 0.3), _box(0.54, 0.45, 0.92, 0.48)] + + _column(0.54, 0.92, 0.55, 0.05, 3), + {"labels": ["Text"] * 8 + ["Picture", "Caption"] + ["Text"] * 3}, + [*range(8), 8, 9, 10, 11, 12], + id="caption-stays-in-its-column", + ), + pytest.param( + VOTERS_PAGE, + {"labels": ["Text"] * 12 + ["Picture"], "column_voters": [True] * 12 + [False]}, + [12, 0, 2, 4, 6, 8, 10, 1, 3, 5, 7, 9, 11], + id="voters-define-the-search-window", + ), + ], +) +def test_sort_reading_order(geoms, kwargs, expected): + order = sort_reading_order(geoms, **kwargs) + assert expected(order) if callable(expected) else order == expected + assert sorted(order) == list(range(len(geoms))) -def test_sort_reading_order_columns(): - boxes = _two_columns() - # The left column must be read entirely before the right one - assert sort_reading_order(boxes) == list(range(8)) - # Right-to-left: the right column comes first - assert sort_reading_order(boxes, direction="rtl") == [4, 5, 6, 7, 0, 1, 2, 3] +def test_sort_reading_order_inputs(): # The result is independent of the input ordering rng = np.random.default_rng(42) for _ in range(5): perm = rng.permutation(8).tolist() - order = sort_reading_order([boxes[idx] for idx in perm]) + order = sort_reading_order([TWO_COLUMNS[idx] for idx in perm]) assert [perm[idx] for idx in order] == list(range(8)) - # A page-wide figure that does not vote keeps the columns below it - geoms = [((0.1, 0.02), (0.9, 0.06)), ((0.05, 0.1), (0.95, 0.6))] - 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(): - geoms = [ - ((0.10, 0.05), (0.90, 0.30)), # 0: figure (page wide) - ((0.15, 0.10), (0.30, 0.12)), # 1: text inside the figure, left side - ((0.70, 0.10), (0.85, 0.12)), # 2: text inside the figure, right side - ((0.10, 0.32), (0.90, 0.35)), # 3: its caption - ((0.10, 0.40), (0.45, 0.70)), # 4: figure in the left column - ((0.10, 0.72), (0.45, 0.80)), # 5: its caption - ] - labels = ["Picture", "Picture", "Picture", "Caption", "Picture", "Caption"] - geoms += [((0.55, 0.40 + 0.05 * i), (0.90, 0.43 + 0.05 * i)) for i in range(8)] # 6-13: right column - geoms += [((0.10, 0.82 + 0.04 * i), (0.45, 0.85 + 0.04 * i)) for i in range(3)] # 14-16: left column - labels += ["Text"] * 11 - # The text inside a figure is read right after it, before its caption - assert sort_reading_order(geoms, labels=labels) == [0, 1, 2, 3, 4, 5, 14, 15, 16, *range(6, 14)] - # Also for text labeled as body text - labels[2] = "Text" - assert sort_reading_order(geoms, labels=labels)[:4] == [0, 1, 2, 3] - assert sorted(sort_reading_order(geoms)) == list(range(len(geoms))) - - -def test_sort_reading_order_input_formats(): - boxes = _two_columns(2) - expected = sort_reading_order(boxes) - # (N, 4) array + # (N, 4) arrays, (N, 4, 2) polygons (straight polygons are untouched) and absolute coordinates + boxes = _column(0.1, 0.45, 0.1, 0.2, 2, 0.15) + _column(0.55, 0.9, 0.1, 0.2, 2, 0.15) as_array = np.asarray([(x0, y0, x1, y1) for ((x0, y0), (x1, y1)) in boxes]) - assert sort_reading_order(as_array) == expected - # (N, 4, 2) rotated polygons as_polys = np.asarray([[(x0, y0), (x1, y0), (x1, y1), (x0, y1)] for ((x0, y0), (x1, y1)) in boxes]) - assert sort_reading_order(as_polys) == expected - # absolute coordinates - assert sort_reading_order(as_array * 1000) == expected - - -def test_sort_reading_order_vertical(): - # 4 vertical columns, from right to left (traditional Japanese/Chinese) - cols = [((0.8 - 0.15 * i, 0.1), (0.9 - 0.15 * i, 0.9)) for i in range(4)] - assert sort_reading_order(cols, direction="ttb-rtl") == [0, 1, 2, 3] - assert sort_reading_order(cols, direction="ttb-ltr") == [3, 2, 1, 0] - # Two stacked elements within a column are read top to bottom - cols = [((0.8, 0.5), (0.9, 0.9)), ((0.8, 0.1), (0.9, 0.45)), ((0.6, 0.1), (0.7, 0.9))] - assert sort_reading_order(cols, direction="ttb-rtl") == [1, 0, 2] - - -def test_sort_reading_order_labels(): - geoms = [ - ((0.1, 0.92), (0.9, 0.97)), # 0: page footer - ((0.1, 0.02), (0.9, 0.06)), # 1: page header - ((0.1, 0.1), (0.9, 0.4)), # 2: body text - ((0.1, 0.45), (0.5, 0.7)), # 3: figure - ((0.1, 0.71), (0.5, 0.75)), # 4: caption below the figure - ((0.55, 0.45), (0.9, 0.88)), # 5: body on the right of the figure - ((0.1, 0.8), (0.5, 0.84)), # 6: footnote - ] - labels = ["Page-footer", "Page-header", "Text", "Picture", "Caption", "Text", "Footnote"] - assert sort_reading_order(geoms, labels=labels) == [1, 2, 3, 4, 5, 6, 0] - # A caption above its figure is read before it - labels_above = list(labels) - geoms_above = list(geoms) - geoms_above[4] = ((0.1, 0.41), (0.5, 0.44)) - assert sort_reading_order(geoms_above, labels=labels_above) == [1, 2, 4, 3, 5, 6, 0] - # A caption too far from any float keeps its natural position in the body - geoms_far = list(geoms) - geoms_far[3] = ((0.1, 0.45), (0.5, 0.5)) # shrink the figure - geoms_far[4] = ((0.55, 0.02), (0.9, 0.06)) # move the caption to the top right - order = sort_reading_order(geoms_far, labels=labels) - assert order.index(4) < order.index(3) - with pytest.raises(ValueError): - sort_reading_order(geoms, labels=labels[:-1]) - with pytest.raises(ValueError): - sort_reading_order(geoms, direction="ttb") - - -def test_sort_reading_order_degenerate_geometries(): - # Identical & zero-area boxes should not crash nor loop - geoms = [((0.1, 0.1), (0.1, 0.1))] * 3 + [((0.5, 0.5), (0.5, 0.5))] - order = sort_reading_order(geoms) - assert sorted(order) == list(range(4)) - - -def test_resolve_reading_segments(): - # 3 lines of a paragraph, a vertical gap, then 2 more lines - geoms = [((0.1, 0.1 + 0.05 * i), (0.9, 0.13 + 0.05 * i)) for i in range(3)] - geoms += [((0.1, 0.4 + 0.05 * i), (0.9, 0.43 + 0.05 * i)) for i in range(2)] - assert resolve_reading_segments(geoms) == [[0, 1, 2], [3, 4]] - # A label change breaks the segment - labels = ["Title", "Text", "Text", "Text", "Text"] - assert resolve_reading_segments(geoms, labels=labels) == [[0], [1, 2], [3, 4]] - # Floats are never merged with their neighbors - labels = ["Table", "Table", "Text", "Text", "Text"] - assert resolve_reading_segments(geoms, labels=labels) == [[0], [1], [2], [3, 4]] - # Two columns end up in distinct segments - boxes = _two_columns() - assert resolve_reading_segments(boxes) == [[0, 1, 2, 3], [4, 5, 6, 7]] - assert resolve_reading_segments([]) == [] - # 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]] - # 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))] - labels = assign_layout_labels(geoms, regions, ["Title", "Text"]) - assert labels == ["Title", "Text", None] - # Rotated region geometries are supported - poly_regions = np.asarray([[(0.05, 0.05), (0.45, 0.05), (0.45, 0.25), (0.05, 0.25)]]) - assert assign_layout_labels(geoms[:1], poly_regions, ["Table"]) == ["Table"] - assert assign_layout_labels([], regions, ["Title", "Text"]) == [] - # The smallest nested region wins - nested = [((0.0, 0.0), (1.0, 1.0)), ((0.08, 0.08), (0.5, 0.3))] - assert assign_layout_labels(geoms[:1], nested, ["Picture", "Caption"]) == ["Caption"] - with pytest.raises(ValueError): - assign_layout_labels(geoms, regions, ["Title"]) - - -def test_reading_order_predictor(): - predictor = ReadingOrderPredictor() - assert predictor.direction == "auto" - assert "auto" in repr(predictor) - geoms = [((0.55, 0.2), (0.9, 0.8)), ((0.1, 0.05), (0.9, 0.15)), ((0.1, 0.2), (0.45, 0.8))] - # Automatic direction detection from the text content - assert predictor(geoms, texts=["right column", "the title", "left column"]) == [1, 2, 0] - assert predictor(geoms, texts=["العمود الأيسر", "العنوان", "العمود الأيمن"]) == [1, 0, 2] - # Language hint fallback when no text is available - assert predictor(geoms, language="ar") == [1, 0, 2] - assert predictor.resolve_direction(["hello"]) == "ltr" - assert ReadingOrderPredictor(direction="rtl").resolve_direction(["hello"]) == "rtl" - with pytest.raises(ValueError): - ReadingOrderPredictor(direction="bottom-up") - - -def _rotated_box(box, deg, width=800, height=1000): - angle = np.deg2rad(deg) - rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) - center = np.array([width / 2, height / 2]) - (x0, y0), (x1, y1) = box - pts = np.array([ - [x0 * width, y0 * height], - [x1 * width, y0 * height], - [x1 * width, y1 * height], - [x0 * width, y1 * height], - ]) - return ((pts - center) @ rot.T + center) / [width, height] - - -def test_sort_reading_order_rotated_pages(): - title = [((0.1, 0.06), (0.9, 0.09))] - left = [((0.1, 0.12 + 0.05 * idx), (0.47, 0.15 + 0.05 * idx)) for idx in range(5)] - right = [((0.53, 0.12 + 0.05 * idx), (0.9, 0.15 + 0.05 * idx)) for idx in range(5)] - geoms = title + left + right - expected = list(range(11)) - for deg in (-35, -15, 15, 35): + for geoms in (as_array, as_polys, as_array * 1000): + assert sort_reading_order(geoms) == [0, 1, 2, 3] + # Rotated pages are de-skewed + geoms = [_box(0.1, 0.06, 0.9, 0.09)] + _column(0.1, 0.47, 0.12, 0.05, 5) + _column(0.53, 0.9, 0.12, 0.05, 5) + for deg in (-35, -15, 0, 15, 35): rotated = [_rotated_box(box, deg) for box in geoms] - assert sort_reading_order(rotated) == expected - assert sort_reading_order(rotated, page_shape=(1000, 800)) == expected - # straight polygons are untouched (angle below the threshold, behavior identical to 2-point boxes) - straight = np.asarray([[(x0, y0), (x1, y0), (x1, y1), (x0, y1)] for ((x0, y0), (x1, y1)) in geoms]) - assert sort_reading_order(straight) == expected - - -def test_deskew_reading_geometries(): - geoms = [((0.1, 0.12), (0.47, 0.15)), ((0.53, 0.12), (0.9, 0.15))] - rotated = [_rotated_box(box, 25) for box in geoms] - # straight 2-point boxes are returned unchanged - out, regions = deskew_reading_geometries(geoms, [((0.0, 0.0), (1.0, 0.5))]) - assert out == list(geoms) and len(regions) == 1 - # rotated polygons are de-skewed: the two boxes end up on the same visual row - out, _ = deskew_reading_geometries(rotated, page_shape=(1000, 800)) - y_centers = [np.asarray(poly)[:, 1].mean() for poly in out] - assert abs(y_centers[0] - y_centers[1]) < 0.005 - # a straight region is expanded to its corners and rotated with the elements - out, regions = deskew_reading_geometries(rotated, [((0.0, 0.1), (1.0, 0.2))], page_shape=(1000, 800)) - assert np.asarray(regions[0]).shape == (4, 2) - # the operation is idempotent - again, _ = deskew_reading_geometries(out, page_shape=(1000, 800)) - assert all(np.allclose(a, b) for a, b in zip(out, again)) - # angle_geoms as the estimation source - out, _ = deskew_reading_geometries(rotated, page_shape=(1000, 800), angle_geoms=np.stack(rotated)) - y_centers = [np.asarray(poly)[:, 1].mean() for poly in out] - assert abs(y_centers[0] - y_centers[1]) < 0.005 - - -def test_reading_order_predictor_rotated(): - left = [_rotated_box(((0.1, 0.1 + 0.2 * idx), (0.3, 0.2 + 0.2 * idx)), 25) for idx in range(3)] - right = [_rotated_box(((0.6, 0.1 + 0.2 * idx), (0.8, 0.2 + 0.2 * idx)), 25) for idx in range(3)] - order = ReadingOrderPredictor()(left + right, page_shape=(1000, 800)) - assert order == [0, 1, 2, 3, 4, 5] - - -def test_deskew_strong_rotation_non_square_page(): + assert sort_reading_order(rotated) == list(range(11)) + assert sort_reading_order(rotated, page_shape=(1000, 800)) == list(range(11)) + # Up to 44 degrees on non-square pages layout = [(80, 40, 670, 110), (80, 150, 360, 900), (390, 150, 670, 900)] # title + 2 columns for height, width in [(1000, 750), (700, 2000)]: sx, sy = width / 750, height / 1000 @@ -339,144 +437,208 @@ def test_deskew_strong_rotation_non_square_page(): pts = np.array([[x0 * sx, y0 * sy], [x1 * sx, y0 * sy], [x1 * sx, y1 * sy], [x0 * sx, y1 * sy]]) polys.append(((pts - center) @ rot.T + center) / np.array([width, height])) assert sort_reading_order(polys, page_shape=(height, width)) == [0, 1, 2], (height, width, angle) + # Invalid inputs + for kwargs in ({"labels": PAGE_LABELS[:-1]}, {"direction": "ttb"}, {"region_groups": [0, 1]}): + with pytest.raises(ValueError): + sort_reading_order(LABELED_PAGE, **kwargs) -# Auto generated regression tests for known failures of the reading order algorithm +@pytest.mark.parametrize( + "geoms, kwargs, expected", + [ + pytest.param([], {}, [], id="empty"), + # 3 lines of a paragraph, a vertical gap, then 2 more lines + pytest.param(_column(0.1, 0.9, 0.1, 0.05, 3) + _column(0.1, 0.9, 0.4, 0.05, 2), {}, [[0, 1, 2], [3, 4]]), + pytest.param( + _column(0.1, 0.9, 0.1, 0.05, 3) + _column(0.1, 0.9, 0.4, 0.05, 2), + {"labels": ["Title", "Text", "Text", "Text", "Text"]}, + [[0], [1, 2], [3, 4]], + id="label-change", + ), + pytest.param( + _column(0.1, 0.9, 0.1, 0.05, 3) + _column(0.1, 0.9, 0.4, 0.05, 2), + {"labels": ["Table", "Table", "Text", "Text", "Text"]}, + [[0], [1], [2], [3, 4]], + id="floats-never-merge", + ), + pytest.param(TWO_COLUMNS, {}, [[0, 1, 2, 3], [4, 5, 6, 7]], id="columns"), + pytest.param( + # The fragments of a line split at a wide word gap belong to the paragraph of that line + [_box(0.10, 0.10, 0.40, 0.13), _box(0.43, 0.10, 0.90, 0.13), _box(0.10, 0.14, 0.90, 0.17)], + {}, + [[0, 1, 2]], + id="split-line", + ), + pytest.param([_box(0.10, 0.10, 0.40, 0.13), _box(0.48, 0.10, 0.90, 0.13)], {}, [[0], [1]], id="gutter"), + # The gap is measured in the unit of the element heights, relative coordinates of a wide page being stretched + pytest.param([_box(0.10, 0.10, 0.40, 0.13), _box(0.44, 0.10, 0.90, 0.13)], {}, [[0, 1]], id="word-gap"), + pytest.param( + [_box(0.10, 0.10, 0.40, 0.13), _box(0.44, 0.10, 0.90, 0.13)], + {"page_shape": (500, 1000)}, + [[0], [1]], + id="word-gap-on-a-wide-page", + ), + pytest.param(_column(0.1, 0.9, 0.1, 0.04, 2), {"region_groups": [0, 0]}, [[0, 1]], id="same-region"), + pytest.param(_column(0.1, 0.9, 0.1, 0.04, 2), {"region_groups": [0, 1]}, [[0], [1]], id="distinct-regions"), + pytest.param( + # A line split into 3 fragments (words missed by the detection), then the short last line of its paragraph, + # then 2 more paragraphs (pixels) + [ + _box(*box) + for box in [ + [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], + ] + ], + {}, + [[0, 1, 2, 3], [4, 5], [6, 7, 8, 9]], + id="fragmented-paragraphs", + ), + pytest.param( + [_box(0.05, 0.10, 0.30, 0.13), _box(0.70, 0.10, 0.95, 0.13), _box(0.05, 0.14, 0.95, 0.17)], + {}, + [[0], [1, 2]], + id="far-fragment", + ), + pytest.param( + # Within a layout region of body text, the fragments of a row belong together whatever the gap + [_box(0.05, 0.10, 0.30, 0.13), _box(0.70, 0.10, 0.95, 0.13), _box(0.05, 0.14, 0.95, 0.17)], + {"labels": ["Text"] * 3, "region_groups": [0, 0, 0]}, + [[0, 1, 2]], + id="far-fragment-in-a-region", + ), + pytest.param( + # A degenerate page shape is ignored + [_box(0.05, 0.10, 0.30, 0.13), _box(0.70, 0.10, 0.95, 0.13), _box(0.05, 0.14, 0.95, 0.17)], + {"page_shape": (0, 0)}, + [[0], [1, 2]], + id="degenerate-page-shape", + ), + pytest.param( + SPLIT_CAPTION, + {"labels": SPLIT_CAPTION_LABELS, "region_groups": [-1, -1, 7, 7, 7]}, + lambda segments: _flat(segments)[_flat(segments).index(0) + 1 : _flat(segments).index(0) + 4] == [2, 3, 4], + id="grouped-caption", + ), + pytest.param( + BESIDE, + {"labels": BESIDE_LABELS, "region_groups": BESIDE_GROUPS}, + lambda segments: [seg for seg in segments if LAST_LINE in seg][0][-2:] == [LAST_LINE - 1, LAST_LINE], + id="paragraph-beside-a-figure", + ), + ], +) +def test_resolve_reading_segments(geoms, kwargs, expected): + segments = resolve_reading_segments(geoms, **kwargs) + assert expected(segments) if callable(expected) else segments == expected + assert sorted(_flat(segments)) == list(range(len(geoms))) -def _box(x0, y0, x1, y1): - return ((x0, y0), (x1, y1)) +@pytest.mark.parametrize( + "geoms, regions, labels, expected", + [ + pytest.param( + [_box(0.1, 0.1, 0.4, 0.2), _box(0.6, 0.1, 0.9, 0.2), _box(0.1, 0.5, 0.4, 0.6)], + [_box(0.05, 0.05, 0.45, 0.25), _box(0.55, 0.05, 0.95, 0.25)], + ["Title", "Text"], + ["Title", "Text", None], + id="basic", + ), + pytest.param( + [_box(0.1, 0.1, 0.4, 0.2)], + np.asarray([[(0.05, 0.05), (0.45, 0.05), (0.45, 0.25), (0.05, 0.25)]]), + ["Table"], + ["Table"], + id="polygon-regions", + ), + pytest.param([], [_box(0.05, 0.05, 0.45, 0.25)], ["Title"], [], id="empty"), + pytest.param( + [_box(0.1, 0.1, 0.4, 0.2)], + [_box(0.0, 0.0, 1.0, 1.0), _box(0.08, 0.08, 0.5, 0.3)], + ["Picture", "Caption"], + ["Caption"], + ), + pytest.param( + # The smallest region wins, except for a float nested in page furniture + [ + _box(0.30, 0.52, 0.60, 0.55), # caption inside the picture + _box(0.20, 0.30, 0.40, 0.33), # picture text + _box(0.03, 0.015, 0.12, 0.035), # logo text, in the page header + _box(0.30, 0.015, 0.60, 0.035), # header text + _box(0.10, 0.80, 0.90, 0.83), # body text + ], + [ + _box(0.10, 0.20, 0.90, 0.60), + _box(0.25, 0.51, 0.65, 0.56), + _box(0.0, 0.0, 1.0, 0.05), + _box(0.02, 0.01, 0.15, 0.04), + ], + ["Picture", "Caption", "Page-header", "Picture"], + ["Caption", "Picture", "Page-header", "Page-header", None], + id="nested", + ), + pytest.param( + [_box(0.30, 0.52, 0.60, 0.55), _box(0.20, 0.30, 0.40, 0.33)] + + [_box(0.03, 0.015, 0.12, 0.035), _box(0.30, 0.015, 0.60, 0.035), _box(0.10, 0.80, 0.90, 0.83)], + [ + _box(0.02, 0.01, 0.15, 0.04), + _box(0.0, 0.0, 1.0, 0.05), + _box(0.25, 0.51, 0.65, 0.56), + _box(0.10, 0.20, 0.90, 0.60), + ], + ["Picture", "Page-header", "Caption", "Picture"], + ["Caption", "Picture", "Page-header", "Page-header", None], + id="nested-reversed", + ), + pytest.param( + # Only floats yield to the furniture + [_box(0.3, 0.01, 0.7, 0.04)], + [_box(0.0, 0.0, 1.0, 0.05), _box(0.25, 0.005, 0.75, 0.045)], + ["Page-header", "Title"], + ["Title"], + id="furniture", + ), + pytest.param( + # A region covering less than `min_coverage` never wins + [_box(0.0, 0.0, 0.4, 0.1)], + [_box(0.3, 0.0, 0.5, 0.1)], + ["Text"], + [None], + id="min-coverage", + ), + pytest.param( + # A smaller region covering the element partially does not win over the one containing it + [_box(0.1, 0.10, 0.9, 0.12)], + [_box(0.05, 0.05, 0.95, 0.5), _box(0.1, 0.09, 0.52, 0.13)], + ["Text", "Section-header"], + ["Text"], + id="partial", + ), + pytest.param( + # A line slightly overflowing a caption nested in a picture still gets the caption label + [_box(0.30, 0.52, 0.60, 0.55)], + [_box(0.10, 0.20, 0.90, 0.60), _box(0.31, 0.51, 0.65, 0.56)], + ["Picture", "Caption"], + ["Caption"], + id="overflowing-caption", + ), + ], +) +def test_assign_layout_labels(geoms, regions, labels, expected): + assert assign_layout_labels(geoms, regions, labels) == expected -def test_sort_reading_order_fragmented_columns(): - left = [ - _box(0.10, 0.10, 0.45, 0.13), # 0 wide - _box(0.10, 0.14, 0.25, 0.17), # 1 narrow (left part of a split line) - _box(0.34, 0.14, 0.45, 0.17), # 2 stray fragment (right part), same visual row as 1 - _box(0.10, 0.18, 0.45, 0.21), # 3 - _box(0.10, 0.22, 0.45, 0.25), # 4 - _box(0.10, 0.26, 0.45, 0.29), # 5 - ] - right = [_box(0.55, 0.10 + 0.04 * i, 0.90, 0.13 + 0.04 * i) for i in range(6)] # 6..11 - order = sort_reading_order(left + right) - # every left element (0..5) is read before every right element (6..11) - assert max(order.index(i) for i in range(6)) < min(order.index(i) for i in range(6, 12)) - - -def test_fragmented_row_with_merged_column_components(): - geoms = [ - _box(0.35, 0.05, 0.65, 0.10), # 0 gutter-straddling element (bridges both columns) - _box(0.10, 0.15, 0.45, 0.20), # 1 left col, row 1 - _box(0.10, 0.22, 0.16, 0.27), # 2 left col, row 2, fragment A - _box(0.17, 0.22, 0.24, 0.27), # 3 left col, row 2, fragment B - _box(0.25, 0.22, 0.45, 0.27), # 4 left col, row 2, fragment C - _box(0.10, 0.29, 0.45, 0.34), # 5 left col, row 3 - _box(0.55, 0.15, 0.90, 0.20), # 6 right col, row 1 - _box(0.55, 0.22, 0.90, 0.27), # 7 right col, row 2 - ] - assert sort_reading_order(geoms) == [0, 1, 2, 3, 4, 5, 6, 7] - - -def test_sort_reading_order_keeps_key_value_rows_together(): - geoms = [ - _box(0.05, 0.02, 0.95, 0.06), # 0 full-width - _box(0.05, 0.08, 0.95, 0.12), # 1 full-width - _box(0.05, 0.14, 0.95, 0.18), # 2 full-width - _box(0.05, 0.20, 0.30, 0.24), # 3 label - _box(0.65, 0.20, 0.95, 0.24), # 4 value - _box(0.05, 0.26, 0.30, 0.30), # 5 label - _box(0.65, 0.26, 0.95, 0.30), # 6 value - _box(0.05, 0.32, 0.30, 0.36), # 7 label - _box(0.65, 0.32, 0.95, 0.36), # 8 value - _box(0.05, 0.38, 0.95, 0.42), # 9 full-width - _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_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 - _box(0.05, 0.10, 0.45, 0.14), # 1 left column - _box(0.05, 0.16, 0.45, 0.20), # 2 left column - _box(0.05, 0.22, 0.45, 0.26), # 3 left column - _box(0.55, 0.10, 0.95, 0.14), # 4 right column - _box(0.55, 0.16, 0.95, 0.20), # 5 right column - _box(0.55, 0.22, 0.95, 0.26), # 6 right column - ] - assert sort_reading_order(geoms) == [0, 1, 2, 3, 4, 5, 6] - - -def test_sort_reading_order_region_groups(): - # A 3-line caption whose last line is closer to the table of the other column - geoms = [ - ((0.526, 0.115), (0.903, 0.33)), # 0: figure - ((0.102, 0.248), (0.477, 0.443)), # 1: table - ((0.52, 0.352), (0.91, 0.366)), # 2: caption line 1 - ((0.52, 0.37), (0.91, 0.384)), # 3: caption line 2 - ((0.52, 0.419), (0.73, 0.434)), # 4: caption line 3 - ] - labels = ["Picture", "Table", "Caption", "Caption", "Caption"] - split = sort_reading_order(geoms, labels=labels) - assert split.index(4) == split.index(1) + 1 - # Grouped, the lines follow the float closest to their union - order = sort_reading_order(geoms, labels=labels, region_groups=[-1, -1, 7, 7, 7]) - assert order[order.index(0) + 1 : order.index(0) + 4] == [2, 3, 4] +def test_layout_region_helpers(): with pytest.raises(ValueError): - sort_reading_order(geoms, labels=labels, region_groups=[0, 1]) - 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] + assign_layout_labels([_box(0.1, 0.1, 0.4, 0.2)], [_box(0.05, 0.05, 0.45, 0.25)] * 2, ["Title"]) # Overlapping regions of the same label are grouped regions = np.asarray([ [0.0, 0.0, 0.5, 0.5], # 0 @@ -490,170 +652,104 @@ def test_sort_reading_order_region_groups(): assert len(set(_region_groups(regions[:2], ["Text"] * 2, min_overlap=0.5))) == 2 -def _paragraph_beside_a_figure(): - """A single-column page (3-column author grid) with a paragraph ending one line below the figure beside it""" - geoms, labels, groups = [[0.1, 0.05, 0.9, 0.09]], ["Title"], [0] - for row in range(4): - for col, x in enumerate((0.15, 0.41, 0.68)): - geoms.append([x, 0.14 + row * 0.03, x + 0.18, 0.154 + row * 0.03]) - labels.append(None) - groups.append(-1) - for y in np.linspace(0.30, 0.64, 12): - geoms.append([0.08, y, 0.48, y + 0.014]) - labels.append("Text") - groups.append(1) - last = len(geoms) - geoms.append([0.08, 0.662, 0.35, 0.676]) # below the figure - labels.append("Text") - groups.append(1) - figure = len(geoms) - geoms.append([0.52, 0.31, 0.91, 0.667]) - labels.append("Picture") - groups.append(2) - geoms.append([0.08, 0.70, 0.30, 0.712]) - labels.append("Section-header") - groups.append(3) - geoms = [((x0, y0), (x1, y1)) for x0, y0, x1, y1 in geoms] - return geoms, labels, groups, last, figure - - -def test_sort_reading_order_float_does_not_split_a_region(): - geoms, labels, groups, last, figure = _paragraph_beside_a_figure() - # Without region groups, the figure splits the paragraph - order = sort_reading_order(geoms, labels=labels) - assert order.index(figure) < order.index(last) - order = sort_reading_order(geoms, labels=labels, region_groups=groups) - assert order.index(last) < order.index(figure) < order.index(len(geoms) - 1) - assert sorted(order) == list(range(len(geoms))) - segments = resolve_reading_segments(geoms, labels=labels, region_groups=groups) - assert [seg for seg in segments if last in seg][0][-2:] == [last - 1, last] - - -def test_sort_reading_order_caption_out_of_reach_stays_in_its_column(): - # A caption out of reach of any float is read in its own column - geoms = [((0.08, 0.1 + 0.05 * i), (0.46, 0.13 + 0.05 * i)) for i in range(8)] # 0-7: left column - geoms += [((0.54, 0.1), (0.92, 0.3))] # 8: figure, right column - geoms += [((0.54, 0.45), (0.92, 0.48))] # 9: its caption, out of reach (0.15 below) - geoms += [((0.54, 0.55 + 0.05 * i), (0.92, 0.58 + 0.05 * i)) for i in range(3)] # 10-12: right column - labels = ["Text"] * 8 + ["Picture", "Caption"] + ["Text"] * 3 - order = sort_reading_order(geoms, labels=labels) - assert order == [*range(8), 8, 9, 10, 11, 12] - - -def test_defer_floats(): - from doctr.models.reading_order.base import _defer_floats - - labels = ["text", "text", "picture", "text", "text", "table", "picture", "text", "text"] - groups = [0, 0, 1, 0, 3, 4, 5, 3, 6] - # A float interrupting region 0, then two floats interrupting region 3 - assert _defer_floats(list(range(9)), labels, {}, groups) == [0, 1, 3, 2, 4, 7, 5, 6, 8] - # The elements inside a float move with it - labels = ["text", "picture", "picture", "text", "text"] - groups = [0, 1, 1, 0, 2] - assert _defer_floats([0, 1, 2, 3, 4], labels, {1: [2]}, groups) == [0, 3, 1, 2, 4] - # Floats interleaved with the lines of a region all wait for its end - labels = ["text", "table", "text", "table", "text", "table", "text"] - groups = [0, 1, 0, 2, 0, 3, 0] - assert _defer_floats(list(range(7)), labels, {}, groups) == [0, 2, 4, 6, 1, 3, 5] - # No region, a float between two regions or at the edges: nothing moves - for labels, groups in [ - (["text", "picture", "text"], [-1, 1, -1]), - (["text", "picture", "text"], [0, 1, 2]), - (["picture", "text", "text"], [1, 0, 0]), - (["text", "text", "picture"], [0, 0, 1]), - ]: - assert _defer_floats([0, 1, 2], labels, {}, groups) == [0, 1, 2] - - -def test_defer_floats_only_moves_a_float_standing_beside_its_region(): - from doctr.models.reading_order.base import _defer_floats - - labels, groups = ["text", "picture", "text"], [0, 1, 0] - # Figure stacked in a loose region - stacked = np.array([[0.1, 0.1, 0.9, 0.2], [0.2, 0.3, 0.8, 0.6], [0.1, 0.7, 0.9, 0.8]]) - assert _defer_floats([0, 1, 2], labels, {}, groups, stacked) == [0, 1, 2] - # Figure beside the paragraph - beside = np.array([[0.05, 0.3, 0.45, 0.32], [0.55, 0.3, 0.95, 0.6], [0.05, 0.62, 0.45, 0.64]]) - assert _defer_floats([0, 1, 2], labels, {}, groups, beside) == [0, 2, 1] - # Text crossing the figure - crossing = np.array([[0.1, 0.3, 0.6, 0.32], [0.5, 0.3, 0.95, 0.6], [0.1, 0.7, 0.9, 0.8]]) - assert _defer_floats([0, 1, 2], labels, {}, groups, crossing) == [0, 1, 2] - # Without boxes, every interrupting float moves - assert _defer_floats([0, 1, 2], labels, {}, groups) == [0, 2, 1] - - -def test_sort_reading_order_float_stacked_in_a_loose_region(): - # A text region enclosing a figure: the figure stays between the lines around it - geoms = [((0.1, 0.1 + 0.04 * idx), (0.9, 0.13 + 0.04 * idx)) for idx in range(3)] - geoms += [((0.2, 0.3), (0.8, 0.6))] - geoms += [((0.1, 0.7 + 0.04 * idx), (0.9, 0.73 + 0.04 * idx)) for idx in range(3)] - labels = ["Text"] * 3 + ["Picture"] + ["Text"] * 3 - groups = [0, 0, 0, 1, 0, 0, 0] - assert sort_reading_order(geoms, labels=labels, region_groups=groups) == list(range(7)) - # A paragraph beside a figure is still not split - geoms, labels, groups, last, figure = _paragraph_beside_a_figure() - order = sort_reading_order(geoms, labels=labels, region_groups=groups) - assert order.index(last) < order.index(figure) - - -def test_sort_reading_order_column_voters_define_the_search_window(): - # Two text columns on the right half and a wide figure on the left that does not vote - geoms, voters = [], [] - for row in range(6): - y = 0.1 + 0.05 * row - geoms += [((0.55, y), (0.72, y + 0.02)), ((0.78, y), (0.95, y + 0.02))] - voters += [True, True] - geoms.append(((0.02, 0.1), (0.5, 0.4))) - voters.append(False) - order = sort_reading_order(geoms, labels=["Text"] * 12 + ["Picture"], column_voters=voters) - assert order == [12, 0, 2, 4, 6, 8, 10, 1, 3, 5, 7, 9, 11] - - -def test_assign_layout_labels_nested_regions(): - geoms = [ - ((0.30, 0.52), (0.60, 0.55)), # caption inside the picture - ((0.20, 0.30), (0.40, 0.33)), # picture text - ((0.03, 0.015), (0.12, 0.035)), # logo text, in the page header - ((0.30, 0.015), (0.60, 0.035)), # header text - ((0.10, 0.80), (0.90, 0.83)), # body text - ] - regions = [ - ((0.10, 0.20), (0.90, 0.60)), - ((0.25, 0.51), (0.65, 0.56)), - ((0.00, 0.00), (1.00, 0.05)), - ((0.02, 0.01), (0.15, 0.04)), - ] - region_labels = ["Picture", "Caption", "Page-header", "Picture"] - # The smallest region wins, except for a float nested in page furniture - assert assign_layout_labels(geoms, regions, region_labels) == [ - "Caption", - "Picture", - "Page-header", - "Page-header", - None, - ] - assert assign_layout_labels(geoms, regions[::-1], region_labels[::-1]) == [ - "Caption", - "Picture", - "Page-header", - "Page-header", - None, - ] - # Only floats yield to the furniture - assert assign_layout_labels( - [((0.3, 0.01), (0.7, 0.04))], - [((0.0, 0.0), (1.0, 0.05)), ((0.25, 0.005), (0.75, 0.045))], - ["Page-header", "Title"], - ) == ["Title"] - # A region covering less than `min_coverage` never wins - assert assign_layout_labels([((0.0, 0.0), (0.4, 0.1))], [((0.3, 0.0), (0.5, 0.1))], ["Text"]) == [None] - # A smaller region covering the element partially does not win over the one containing it - assert assign_layout_labels( - [((0.1, 0.10), (0.9, 0.12))], - [((0.05, 0.05), (0.95, 0.5)), ((0.1, 0.09), (0.52, 0.13))], - ["Text", "Section-header"], - ) == ["Text"] - # A line slightly overflowing a caption nested in a picture still gets the caption label - assert assign_layout_labels( - [((0.30, 0.52), (0.60, 0.55))], [((0.10, 0.20), (0.90, 0.60)), ((0.31, 0.51), (0.65, 0.56))], region_labels[:2] - ) == ["Caption"] +@pytest.mark.parametrize( + "order, labels, members, groups, boxes, expected", + [ + # A float interrupting region 0, then two floats interrupting region 3 + ( + list(range(9)), + ["text", "text", "picture", "text", "text", "table", "picture", "text", "text"], + {}, + [0, 0, 1, 0, 3, 4, 5, 3, 6], + None, + [0, 1, 3, 2, 4, 7, 5, 6, 8], + ), + # The elements inside a float move with it + ( + [0, 1, 2, 3, 4], + ["text", "picture", "picture", "text", "text"], + {1: [2]}, + [0, 1, 1, 0, 2], + None, + [0, 3, 1, 2, 4], + ), + # Floats interleaved with the lines of a region all wait for its end + (list(range(7)), ["text", "table"] * 3 + ["text"], {}, [0, 1, 0, 2, 0, 3, 0], None, [0, 2, 4, 6, 1, 3, 5]), + # No region, a float between two regions or at the edges: nothing moves + ([0, 1, 2], ["text", "picture", "text"], {}, [-1, 1, -1], None, [0, 1, 2]), + ([0, 1, 2], ["text", "picture", "text"], {}, [0, 1, 2], None, [0, 1, 2]), + ([0, 1, 2], ["picture", "text", "text"], {}, [1, 0, 0], None, [0, 1, 2]), + ([0, 1, 2], ["text", "text", "picture"], {}, [0, 0, 1], None, [0, 1, 2]), + # Only a float standing beside its region moves: not a figure stacked in a loose region... + ( + [0, 1, 2], + ["text", "picture", "text"], + {}, + [0, 1, 0], + np.array([[0.1, 0.1, 0.9, 0.2], [0.2, 0.3, 0.8, 0.6], [0.1, 0.7, 0.9, 0.8]]), + [0, 1, 2], + ), + # ... but a figure beside the paragraph... + ( + [0, 1, 2], + ["text", "picture", "text"], + {}, + [0, 1, 0], + np.array([[0.05, 0.3, 0.45, 0.32], [0.55, 0.3, 0.95, 0.6], [0.05, 0.62, 0.45, 0.64]]), + [0, 2, 1], + ), + # ... and not text crossing the figure + ( + [0, 1, 2], + ["text", "picture", "text"], + {}, + [0, 1, 0], + np.array([[0.1, 0.3, 0.6, 0.32], [0.5, 0.3, 0.95, 0.6], [0.1, 0.7, 0.9, 0.8]]), + [0, 1, 2], + ), + # Without boxes, every interrupting float moves + ([0, 1, 2], ["text", "picture", "text"], {}, [0, 1, 0], None, [0, 2, 1]), + ], +) +def test_defer_floats(order, labels, members, groups, boxes, expected): + assert _defer_floats(order, labels, members, groups, boxes) == expected + + +def test_deskew_reading_geometries(): + geoms = [_box(0.1, 0.12, 0.47, 0.15), _box(0.53, 0.12, 0.9, 0.15)] + rotated = [_rotated_box(box, 25) for box in geoms] + # straight 2-point boxes are returned unchanged + out, regions = deskew_reading_geometries(geoms, [_box(0.0, 0.0, 1.0, 0.5)]) + assert out == list(geoms) and len(regions) == 1 + # rotated polygons are de-skewed (with the page angle estimated from themselves or from `angle_geoms`): the two + # boxes end up on the same visual row + for angle_geoms in (None, np.stack(rotated)): + out, _ = deskew_reading_geometries(rotated, page_shape=(1000, 800), angle_geoms=angle_geoms) + y_centers = [np.asarray(poly)[:, 1].mean() for poly in out] + assert abs(y_centers[0] - y_centers[1]) < 0.005 + # a straight region is expanded to its corners and rotated with the elements + out, regions = deskew_reading_geometries(rotated, [_box(0.0, 0.1, 1.0, 0.2)], page_shape=(1000, 800)) + assert np.asarray(regions[0]).shape == (4, 2) + # the operation is idempotent + again, _ = deskew_reading_geometries(out, page_shape=(1000, 800)) + assert all(np.allclose(a, b) for a, b in zip(out, again)) + + +def test_reading_order_predictor(): + predictor = ReadingOrderPredictor() + assert predictor.direction == "auto" + assert "auto" in repr(predictor) + geoms = [_box(0.55, 0.2, 0.9, 0.8), _box(0.1, 0.05, 0.9, 0.15), _box(0.1, 0.2, 0.45, 0.8)] + # Automatic direction detection from the text content + assert predictor(geoms, texts=["right column", "the title", "left column"]) == [1, 2, 0] + assert predictor(geoms, texts=["العمود الأيسر", "العنوان", "العمود الأيمن"]) == [1, 0, 2] + # Language hint fallback when no text is available + assert predictor(geoms, language="ar") == [1, 0, 2] + assert predictor.resolve_direction(["hello"]) == "ltr" + assert ReadingOrderPredictor(direction="rtl").resolve_direction(["hello"]) == "rtl" + # Rotated page + left = [_rotated_box(_box(0.1, 0.1 + 0.2 * idx, 0.3, 0.2 + 0.2 * idx), 25) for idx in range(3)] + right = [_rotated_box(_box(0.6, 0.1 + 0.2 * idx, 0.8, 0.2 + 0.2 * idx), 25) for idx in range(3)] + assert predictor(left + right, page_shape=(1000, 800)) == [0, 1, 2, 3, 4, 5] + with pytest.raises(ValueError): + ReadingOrderPredictor(direction="bottom-up")