diff --git a/docs/source/modules/io.rst b/docs/source/modules/io.rst index 56d88e014d..f57e42ba15 100644 --- a/docs/source/modules/io.rst +++ b/docs/source/modules/io.rst @@ -60,6 +60,14 @@ A Page is a collection of Blocks that were on the same physical page. .. autoclass:: Page .. automethod:: show + .. automethod:: items_in_reading_order + .. automethod:: export_as_markdown + .. automethod:: export_as_asciidoc + .. automethod:: export_as_html + .. automethod:: export_as_xml + .. automethod:: export + .. automethod:: render + .. automethod:: export_as KIEPage @@ -71,6 +79,13 @@ semantic class rather than by spatial layout. .. autoclass:: KIEPage .. automethod:: show + .. automethod:: export_as_markdown + .. automethod:: export_as_asciidoc + .. automethod:: export_as_html + .. automethod:: export_as_xml + .. automethod:: export + .. automethod:: render + .. automethod:: export_as Document @@ -81,6 +96,13 @@ A Document is a collection of Pages. .. autoclass:: Document .. automethod:: show + .. automethod:: export_as_markdown + .. automethod:: export_as_asciidoc + .. automethod:: export_as_xml + .. automethod:: export_as_html + .. automethod:: export + .. automethod:: render + .. automethod:: export_as KIEDocument @@ -117,3 +139,41 @@ High-performance file reading and conversion to processable structured data. .. automethod:: from_url .. automethod:: from_images + + +.. _reading_order: + +Reading order +------------- + +The reading-order-aware export of a :class:`Document` / :class:`Page` to Markdown, AsciiDoc, HTML, or XML is available +through the ``export_as_markdown`` / ``export_as_asciidoc`` / ``export_as_html`` / ``export_as_xml`` / ``export_as`` / ``export`` / ``render`` methods documented above, which +delegate to the exporters of :mod:`doctr.io.exporters`. +The underlying ordering primitives live in :mod:`doctr.models.reading_order`. + +Every export path shares the same linearization, so ``render()``, ``export()``, ``export_as_xml()`` and the +Markdown / AsciiDoc / HTML exports all present the content in the same order. The result is memoized on the +page, so exporting one page to several formats orders it only once. + +.. currentmodule:: doctr.io + +.. autoclass:: TextExporter + :members: export_page, export_kie_page, export_document + +.. autoclass:: MarkdownExporter + :members: export_page, export_kie_page, export_document + +.. autoclass:: AsciiDocExporter + :members: export_page, export_kie_page, export_document + +.. autoclass:: HTMLExporter + :members: export_page, export_kie_page, export_document + +.. autoclass:: XMLExporter + :members: export_page, export_kie_page, export_document + +.. autofunction:: doctr.io.exporters.page_reading_order + +.. autofunction:: doctr.io.exporters.predictions_in_reading_order + +.. autofunction:: doctr.io.exporters.to_json_safe diff --git a/docs/source/using_doctr/using_models.rst b/docs/source/using_doctr/using_models.rst index b00d75a86a..0ba368e7e2 100644 --- a/docs/source/using_doctr/using_models.rst +++ b/docs/source/using_doctr/using_models.rst @@ -221,7 +221,7 @@ For a comprehensive comparison, we have compiled a detailed benchmark: +--------------------------------------------------+-----------------+---------------+------------------+-------------+--------------+--------------------+ | lw_detr_s | (1024, 1024, 3) | 15.1 M | | | | 0.5 | +--------------------------------------------------+-----------------+---------------+------------------+-------------+--------------+--------------------+ -| lw_detr_m | (1024, 1024, 3) | 29.5 M | | | | 0.7 | +| lw_detr_m | (1024, 1024, 3) | 29.5 M | soon | soon | soon | 0.7 | +--------------------------------------------------+-----------------+---------------+------------------+-------------+--------------+--------------------+ @@ -628,6 +628,35 @@ For reference, here is a sample XML byte string output: +To export the output in reading order as Markdown, AsciiDoc, HTML or XML, you can use the `export_as_markdown`, +`export_as_asciidoc`, `export_as_html` and `export_as_xml` methods (also available on a single :class:`Page`). +Reading order is always applied by these exporters: the content is linearized column by column, the reading direction is inferred from the +recognized text (e.g. right-to-left for Arabic or Hebrew documents), and - when the predictor is run with +layout detection (``detect_layout=True``) - the layout regions are used to render headings, list items and +recognized tables, and to place the page furniture (headers, footers, footnotes): + +.. code-block:: python + + markdown_output = result.export_as_markdown() + asciidoc_output = result.export_as_asciidoc() + html_output = result.export_as_html() + xml_output = result.export_as_xml() + raw_text_output = result.render() # same as result.export_as("text") + dict_output = result.export() # same as result.export_as("json") + +The `export_as` method is a convenience dispatcher over all the formats above:: + + result.export_as("markdown") # or "md" + result.export_as("asciidoc") # or "adoc" + result.export_as("html") + result.export_as("text") # same as render() + result.export_as("json") # same as export() + result.export_as("xml") # same as export_as_xml() + +The document structure itself can also be produced in reading order by building the predictor with +``keep_reading_order=True``, which sorts the blocks of every page in reading order:: + + predictor = ocr_predictor(pretrained=True, keep_reading_order=True) Advanced options ^^^^^^^^^^^^^^^^ diff --git a/doctr/io/__init__.py b/doctr/io/__init__.py index 6eab8c2406..be8e911358 100644 --- a/doctr/io/__init__.py +++ b/doctr/io/__init__.py @@ -1,4 +1,5 @@ from .elements import * +from .exporters import * from .html import * from .image import * from .pdf import * diff --git a/doctr/io/elements.py b/doctr/io/elements.py index e76e2f8c38..cdf03c26d9 100644 --- a/doctr/io/elements.py +++ b/doctr/io/elements.py @@ -4,14 +4,11 @@ # See LICENSE or go to for full license details. from typing import Any -from xml.etree import ElementTree as ET -from xml.etree.ElementTree import Element as ETElement -from xml.etree.ElementTree import SubElement import numpy as np -import doctr from doctr.file_utils import requires_package +from doctr.io.exporters import DocumentExportsMixin, KIEPageExportsMixin, PageExportsMixin, to_json_safe from doctr.utils.common_types import BoundingBox from doctr.utils.geometry import resolve_enclosing_bbox, resolve_enclosing_rbbox from doctr.utils.reconstitution import synthesize_kie_page, synthesize_page @@ -19,7 +16,7 @@ try: # optional dependency for visualization from doctr.utils.visualization import visualize_kie_page, visualize_page -except ModuleNotFoundError: +except ModuleNotFoundError: # pragma: no cover pass __all__ = [ @@ -39,35 +36,9 @@ ] -def _resolve_hocr_language(language: dict[str, Any]) -> str: - """Resolve the language code to use in the hOCR export, falling back to 'en'. - - Args: - language: the page language dictionary `{"value": str | None, "confidence": float | None}` - - Returns: - the detected language code when available, 'en' otherwise - """ - lang_value = language.get("value") if isinstance(language, dict) else None - return lang_value if isinstance(lang_value, str) and len(lang_value) > 0 else "en" - - -def _hocr_bbox(geometry: BoundingBox, width: int, height: int) -> str: - """Format a relative straight bounding box as an absolute hOCR `bbox` property string. - - Args: - geometry: the relative bounding box ((xmin, ymin), (xmax, ymax)) - width: the page width in pixels - height: the page height in pixels - - Returns: - the hOCR `bbox` property string - """ - (xmin, ymin), (xmax, ymax) = geometry - return ( - f"bbox {int(round(xmin * width))} {int(round(ymin * height))} " - f"{int(round(xmax * width))} {int(round(ymax * height))}" - ) +def _empty_page_image(page: np.ndarray | None) -> np.ndarray: + """Return the given page image, or an empty placeholder when the page was restored from an export.""" + return page if page is not None else np.zeros((0, 0, 3), dtype=np.uint8) class Element(NestedObject): @@ -85,7 +56,7 @@ def __init__(self, **kwargs: Any) -> None: def export(self) -> dict[str, Any]: """Exports the object into a nested dict format""" - export_dict = {k: getattr(self, k) for k in self._exported_keys} + export_dict = {k: to_json_safe(getattr(self, k)) for k in self._exported_keys} for children_name in self._children_names: if children_name in ["predictions"]: export_dict[children_name] = { @@ -175,8 +146,7 @@ def extra_repr(self) -> str: @classmethod def from_dict(cls, save_dict: dict[str, Any], **kwargs): - kwargs = {k: save_dict[k] for k in cls._exported_keys} - return cls(**kwargs) + return cls(artefact_type=save_dict["type"], confidence=save_dict["confidence"], geometry=save_dict["geometry"]) class LayoutElement(Element): @@ -451,7 +421,7 @@ def from_dict(cls, save_dict: dict[str, Any], **kwargs): return cls(**kwargs) -class Page(Element): +class Page(PageExportsMixin, Element): """Implements a page element as a collection of blocks Args: @@ -493,10 +463,6 @@ def __init__( self.orientation = orientation if isinstance(orientation, dict) else dict(value=None, confidence=None) self.language = language if isinstance(language, dict) else dict(value=None, confidence=None) - def render(self, block_break: str = "\n\n") -> str: - """Renders the full text of the element""" - return block_break.join(b.render() for b in self.blocks) - def extra_repr(self) -> str: return f"dimensions={self.dimensions}" @@ -534,120 +500,19 @@ def synthesize(self, **kwargs) -> np.ndarray: """ return synthesize_page(self.export(), **kwargs) - def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[bytes, ET.ElementTree]: - """Export the page as XML (hOCR-format) - convention: https://github.com/kba/hocr-spec/blob/master/1.2/spec.md - - Args: - file_title: the title of the XML file - - Returns: - a tuple of the XML byte string, and its ElementTree - """ - p_idx = self.page_idx - block_count: int = 1 - line_count: int = 1 - word_count: int = 1 - height, width = self.dimensions - language = _resolve_hocr_language(self.language) - # Create the XML root element - page_hocr = ETElement("html", attrib={"xmlns": "http://www.w3.org/1999/xhtml", "xml:lang": str(language)}) - # Create the header / SubElements of the root element - head = SubElement(page_hocr, "head") - SubElement(head, "title").text = file_title - SubElement(head, "meta", attrib={"http-equiv": "Content-Type", "content": "text/html; charset=utf-8"}) - SubElement( - head, - "meta", - attrib={"name": "ocr-system", "content": f"python-doctr {doctr.__version__}"}, # type: ignore[attr-defined] - ) - SubElement( - head, - "meta", - attrib={"name": "ocr-capabilities", "content": "ocr_page ocr_carea ocr_par ocr_line ocrx_word"}, - ) - # Create the body - body = SubElement(page_hocr, "body") - page_div = SubElement( - body, - "div", - attrib={ - "class": "ocr_page", - "id": f"page_{p_idx + 1}", - "title": f"image; bbox 0 0 {width} {height}; ppageno 0", - }, - ) - # iterate over the blocks / lines / words and create the XML elements in body line by line with the attributes - for block in self.blocks: - if len(block.geometry) != 2: - raise TypeError("XML export is only available for straight bounding boxes for now.") - block_bbox = _hocr_bbox(block.geometry, width, height) # type: ignore[arg-type] - block_div = SubElement( - page_div, - "div", - attrib={ - "class": "ocr_carea", - "id": f"block_{block_count}", - "title": block_bbox, - }, - ) - paragraph = SubElement( - block_div, - "p", - attrib={ - "class": "ocr_par", - "id": f"par_{block_count}", - "title": block_bbox, - }, - ) - block_count += 1 - for line in block.lines: - # NOTE: baseline, x_size, x_descenders, x_ascenders is currently initalized to 0 - line_span = SubElement( - paragraph, - "span", - attrib={ - "class": "ocr_line", - "id": f"line_{line_count}", - "title": ( - f"{_hocr_bbox(line.geometry, width, height)}; " # type: ignore[arg-type] - "baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0" - ), - }, - ) - line_count += 1 - for word in line.words: - conf = word.confidence - word_div = SubElement( - line_span, - "span", - attrib={ - "class": "ocrx_word", - "id": f"word_{word_count}", - "title": ( - f"{_hocr_bbox(word.geometry, width, height)}; " # type: ignore[arg-type] - f"x_wconf {int(round(conf * 100))}" - ), - }, - ) - # set the text - word_div.text = word.value - word_count += 1 - - return (ET.tostring(page_hocr, encoding="utf-8", method="xml"), ET.ElementTree(page_hocr)) - @classmethod - def from_dict(cls, save_dict: dict[str, Any], **kwargs): - kwargs = {k: save_dict[k] for k in cls._exported_keys} - kwargs.update({ + def from_dict(cls, save_dict: dict[str, Any], page: np.ndarray | None = None, **kwargs): + _kwargs: dict[str, Any] = {k: save_dict[k] for k in cls._exported_keys} + _kwargs.update({ "blocks": [Block.from_dict(block_dict) for block_dict in save_dict["blocks"]], "layout": [LayoutElement.from_dict(region_dict) for region_dict in save_dict.get("layout", [])], "tables": [Table.from_dict(table_dict) for table_dict in save_dict.get("tables", [])], }) - return cls(**kwargs) + # The page image is not part of the export: pass it back explicitly to restore a fully usable page + return cls(page=_empty_page_image(page), **_kwargs) -class KIEPage(Element): +class KIEPage(KIEPageExportsMixin, Element): """Implements a KIE page element as a collection of predictions Args: @@ -682,12 +547,6 @@ def __init__( self.orientation = orientation if isinstance(orientation, dict) else dict(value=None, confidence=None) self.language = language if isinstance(language, dict) else dict(value=None, confidence=None) - def render(self, prediction_break: str = "\n\n") -> str: - """Renders the full text of the element""" - return prediction_break.join( - f"{class_name}: {p.render()}" for class_name, predictions in self.predictions.items() for p in predictions - ) - def extra_repr(self) -> str: return f"dimensions={self.dimensions}" @@ -724,110 +583,21 @@ def synthesize(self, **kwargs) -> np.ndarray: """ return synthesize_kie_page(self.export(), **kwargs) - def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[bytes, ET.ElementTree]: - """Export the page as XML (hOCR-format) - convention: https://github.com/kba/hocr-spec/blob/master/1.2/spec.md - - Args: - file_title: the title of the XML file - - Returns: - a tuple of the XML byte string, and its ElementTree - """ - p_idx = self.page_idx - prediction_count: int = 1 - height, width = self.dimensions - language = _resolve_hocr_language(self.language) - # Create the XML root element - page_hocr = ETElement("html", attrib={"xmlns": "http://www.w3.org/1999/xhtml", "xml:lang": str(language)}) - # Create the header / SubElements of the root element - head = SubElement(page_hocr, "head") - SubElement(head, "title").text = file_title - SubElement(head, "meta", attrib={"http-equiv": "Content-Type", "content": "text/html; charset=utf-8"}) - SubElement( - head, - "meta", - attrib={"name": "ocr-system", "content": f"python-doctr {doctr.__version__}"}, # type: ignore[attr-defined] - ) - SubElement( - head, - "meta", - attrib={"name": "ocr-capabilities", "content": "ocr_page ocr_carea ocr_par ocr_line ocrx_word"}, - ) - # Create the body - body = SubElement(page_hocr, "body") - SubElement( - body, - "div", - attrib={ - "class": "ocr_page", - "id": f"page_{p_idx + 1}", - "title": f"image; bbox 0 0 {width} {height}; ppageno 0", - }, - ) - # iterate over the blocks / lines / words and create the XML elements in body line by line with the attributes - for class_name, predictions in self.predictions.items(): - for prediction in predictions: - if len(prediction.geometry) != 2: - raise TypeError("XML export is only available for straight bounding boxes for now.") - prediction_bbox = _hocr_bbox(prediction.geometry, width, height) # type: ignore[arg-type] - prediction_div = SubElement( - body, - "div", - attrib={ - "class": "ocr_carea", - "id": f"{class_name}_prediction_{prediction_count}", - "title": prediction_bbox, - }, - ) - # NOTE: ocr_par, ocr_line and ocrx_word are the same because the KIE predictions contain only words - # This is a workaround to make it PDF/A compatible - par_div = SubElement( - prediction_div, - "p", - attrib={ - "class": "ocr_par", - "id": f"{class_name}_par_{prediction_count}", - "title": prediction_bbox, - }, - ) - line_span = SubElement( - par_div, - "span", - attrib={ - "class": "ocr_line", - "id": f"{class_name}_line_{prediction_count}", - "title": f"{prediction_bbox}; baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0", - }, - ) - word_div = SubElement( - line_span, - "span", - attrib={ - "class": "ocrx_word", - "id": f"{class_name}_word_{prediction_count}", - "title": f"{prediction_bbox}; x_wconf {int(round(prediction.confidence * 100))}", - }, - ) - word_div.text = prediction.value - prediction_count += 1 - - return ET.tostring(page_hocr, encoding="utf-8", method="xml"), ET.ElementTree(page_hocr) - @classmethod - def from_dict(cls, save_dict: dict[str, Any], **kwargs): - kwargs = {k: save_dict[k] for k in cls._exported_keys} - kwargs.update({ + def from_dict(cls, save_dict: dict[str, Any], page: np.ndarray | None = None, **kwargs): + _kwargs: dict[str, Any] = {k: save_dict[k] for k in cls._exported_keys} + _kwargs.update({ "predictions": { class_name: [Prediction.from_dict(pred) for pred in preds] for class_name, preds in save_dict["predictions"].items() }, "layout": [LayoutElement.from_dict(region_dict) for region_dict in save_dict.get("layout", [])], }) - return cls(**kwargs) + # The page image is not part of the export: pass it back explicitly to restore a fully usable page + return cls(page=_empty_page_image(page), **_kwargs) -class Document(Element): +class Document(DocumentExportsMixin, Element): """Implements a document element as a collection of pages Args: @@ -843,10 +613,6 @@ def __init__( ) -> None: super().__init__(pages=pages) - def render(self, page_break: str = "\n\n\n\n") -> str: - """Renders the full text of the element""" - return page_break.join(p.render() for p in self.pages) - def show(self, **kwargs) -> None: """Overlay the result on a given image""" for result in self.pages: @@ -863,22 +629,18 @@ def synthesize(self, **kwargs) -> list[np.ndarray]: """ return [page.synthesize(**kwargs) for page in self.pages] - def export_as_xml(self, **kwargs) -> list[tuple[bytes, ET.ElementTree]]: - """Export the document as XML (hOCR-format) - - Args: - **kwargs: additional keyword arguments passed to the Page.export_as_xml method - - Returns: - list of tuple of (bytes, ElementTree) - """ - return [page.export_as_xml(**kwargs) for page in self.pages] + _page_cls: Any = Page @classmethod - def from_dict(cls, save_dict: dict[str, Any], **kwargs): - kwargs = {k: save_dict[k] for k in cls._exported_keys} - kwargs.update({"pages": [Page.from_dict(page_dict) for page_dict in save_dict["pages"]]}) - return cls(**kwargs) + def from_dict(cls, save_dict: dict[str, Any], pages: list[np.ndarray] | None = None, **kwargs): + _kwargs: dict[str, Any] = {k: save_dict[k] for k in cls._exported_keys} + _kwargs.update({ + "pages": [ + cls._page_cls.from_dict(page_dict, page=None if pages is None else pages[idx]) + for idx, page_dict in enumerate(save_dict["pages"]) + ] + }) + return cls(**_kwargs) class KIEDocument(Document): @@ -889,6 +651,7 @@ class KIEDocument(Document): """ _children_names: list[str] = ["pages"] + _page_cls: Any = KIEPage pages: list[KIEPage] = [] # type: ignore[assignment] def __init__( diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py new file mode 100644 index 0000000000..6f47789652 --- /dev/null +++ b/doctr/io/exporters.py @@ -0,0 +1,1275 @@ +# Copyright (C) 2021-2026, Mindee. + +# This program is licensed under the Apache License 2.0. +# See LICENSE or go to for full license details. + +from html import escape as _html_escape +from typing import TYPE_CHECKING, Any, ClassVar, cast +from xml.etree import ElementTree as ET +from xml.etree.ElementTree import Element as ETElement +from xml.etree.ElementTree import SubElement + +import numpy as np + +import doctr +from doctr.utils.common_types import BoundingBox + +if TYPE_CHECKING: # pragma: no cover + from doctr.io.elements import Block, KIEPage, Line, Page, Table + +__all__ = [ + "AsciiDocExporter", + "DocumentExportsMixin", + "HTMLExporter", + "KIEPageExportsMixin", + "MarkdownExporter", + "PageExportsMixin", + "TextExporter", + "XMLExporter", + "page_reading_order", +] + + +def _export_as(exporters: dict[str, Any], format: str, **kwargs: Any) -> Any: + fmt = format.strip().lower() + if fmt not in exporters: + raise ValueError(f"unsupported export format '{format}', should be one of {sorted(exporters)}") + return exporters[fmt](**kwargs) + + +def to_json_safe(value: Any) -> Any: + """Recursively convert NumPy containers and scalars into built-in Python types. + + Args: + value: any exported value + + Returns: + the same value with every NumPy array converted to nested tuples and every NumPy scalar to its + Python equivalent + """ + if isinstance(value, np.ndarray): + return value.item() if value.ndim == 0 else tuple(to_json_safe(item) for item in value) + if isinstance(value, np.generic): # np.float32, np.int64, np.bool_, ... + return value.item() + if isinstance(value, dict): + return {str(key): to_json_safe(item) for key, item in value.items()} + if isinstance(value, tuple): + return tuple(to_json_safe(item) for item in value) + if isinstance(value, (list, set, frozenset)): + return [to_json_safe(item) for item in value] + return value + + +_LIST_LABELS = {"list_item"} +# Characters / line markers that carry a structural meaning and are escaped to preserve the raw OCR text +_MD_SPECIAL_CHARS = "\\`*_[]|#<>" +_MD_LINE_MARKERS = "-+>#=`" +_ADOC_SPECIAL_CHARS = "\\`*_#^~|+{}<>" +_ADOC_LINE_MARKERS = "=*.-/+" + + +def _covering_region_indices(geoms: list[Any], region_geoms: list[Any], min_coverage: float = 0.5) -> list[int]: + """For each element geometry, the index of the layout region covering the largest share of its area. + + Uses the same area-coverage criterion as :func:`doctr.models.reading_order.assign_layout_labels`, and + returns -1 when no region covers the element by at least `min_coverage`. The geometries are expected to + be in the same (upright) frame. + """ + from doctr.models.reading_order.base import _to_boxes + + if len(region_geoms) == 0 or len(geoms) == 0: + return [-1] * len(geoms) + boxes, regions = _to_boxes(geoms), _to_boxes(region_geoms) + 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] + best = coverage.argmax(axis=1) + return [int(reg) if coverage[i, reg] >= min_coverage else -1 for i, reg in enumerate(best)] + + +def _reading_order_signature(page: "Page", direction: str) -> tuple[Any, ...]: + """A cheap structural fingerprint of a page, used to invalidate the reading-order cache. + + Covers the requested direction and the identity (plus line count) of every block and table, so + replacing or re-grouping the page content invalidates the cache. In-place edits to a `Line`'s words + are not detected; callers mutating a page that deeply should drop `_reading_order_cache` themselves. + """ + return ( + direction, + tuple((id(block), len(block.lines)) for block in page.blocks), + tuple(id(table) for table in getattr(page, "tables", ()) or ()), + ) + + +def _store_reading_order(page: "Page", signature: tuple[Any, ...], result: tuple[Any, ...]) -> None: + """Memoize a reading-order result on the page, ignoring pages that reject attribute assignment.""" + try: + page._reading_order_cache = (signature, result) # type: ignore[attr-defined] + except AttributeError: # pragma: no cover + pass + + +def page_reading_order(page: "Page", direction: str = "auto") -> tuple[list[Any], list[str | None], str]: + """Linearize the content of a page (blocks & tables) in reading order. + + The result is memoized on the page: every exporter calls this, so a page exported to several formats + (or built with `keep_reading_order=True` and then exported) orders its content once. + + Args: + page: the page to linearize + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + + Returns: + a tuple with the ordered items (blocks & tables), their layout label (None without layout) and the + effective reading direction + """ + from doctr.io.elements import Block, Table + from doctr.models.reading_order import ( + ReadingOrderPredictor, + assign_layout_labels, + deskew_reading_geometries, + normalize_layout_label, + 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 list(items), list(labels), resolved + + 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] + + lines = [line for block in page.blocks for line in block.lines] + elements: list[Any] = [*lines, *page.tables] + 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], + ) + elt_labels: list[str | None] = [None] * len(elements) + if len(region_geoms) > 0: + elt_labels = assign_layout_labels(elt_geoms, region_geoms, region_labels) + elt_labels = ["Table" if isinstance(elt, Table) else label for elt, label in zip(elements, elt_labels)] + segments = resolve_reading_segments(elt_geoms, direction=direction, labels=elt_labels) + + 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 = _covering_region_indices(elt_geoms, region_geoms) if len(region_geoms) > 0 else [-1] * len(elements) + 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]] + if isinstance(first, Table): + items.append(first) + labels.append("Table") + 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]])]) + else: + items.append(Block(lines=[elements[idx]], artefacts=_claim_artefacts([elements[idx]]))) + labels.append(seg_label) + open_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: + last_block = next((item for item in reversed(items) if isinstance(item, Block)), None) + if last_block is not None: + last_block.artefacts = [*last_block.artefacts, *leftover] + _store_reading_order(page, signature, (items, labels, direction)) + return items, labels, direction + + +def _line_render_direction(line: "Line", page_direction: str, auto: bool) -> str: + """Resolve the direction used to order the words of a line. + + For vertical pages the words are always read top to bottom. For horizontal pages, when the page direction + was inferred automatically, the base direction of each line is detected from its own text so that an + embedded left-to-right run (e.g. a Latin quotation on an Arabic page) keeps its natural word order; when + the direction is set explicitly, it is applied uniformly to every line. + """ + if page_direction in ("ttb-rtl", "ttb-ltr") or not auto or len(line.words) <= 1: + return page_direction + from doctr.models.reading_order import detect_text_direction + + return detect_text_direction([word.render() for word in line.words]) + + +def ordered_line_words(line: "Line", direction: str = "ltr", auto: bool = False) -> list[Any]: + """Return the words of a line in reading order. + + Args: + line: the line whose words should be ordered + direction: the reading direction resolved for the page + auto: whether the page direction was inferred (each line then gets its own base direction) + + Returns: + the words of the line, ordered logically + """ + direction = _line_render_direction(line, direction, auto) + if direction in ("ttb-rtl", "ttb-ltr"): + return sorted(line.words, key=lambda word: float(np.asarray(word.geometry, dtype=np.float64)[..., 1].mean())) + if direction == "rtl": + return sorted(line.words, key=lambda word: -float(np.asarray(word.geometry, dtype=np.float64)[..., 0].mean())) + return list(line.words) + + +def predictions_in_reading_order(page: "KIEPage", predictions: list[Any], direction: str = "auto") -> list[Any]: + """Sort the predictions of a single KIE detection class in reading order. + + Args: + page: the KIE page the predictions belong to (used for its dimensions and detected language) + predictions: the predictions of one detection class + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + + Returns: + the predictions, ordered logically + """ + from doctr.models.reading_order import ReadingOrderPredictor + + if len(predictions) < 2: + return list(predictions) + language = page.language.get("value") if isinstance(page.language, dict) else None + order = ReadingOrderPredictor(direction=direction)( + [prediction.geometry for prediction in predictions], + texts=[prediction.value for prediction in predictions], + language=language, + page_shape=page.dimensions, + ) + return [predictions[idx] for idx in order] + + +class _PageTextExporter: + """Shared logic of the reading-order-aware text exporters. + + Subclasses define the format specifics: heading prefixes (per normalized layout label), the bullet + prefix, character escaping, line finalization (neutralizing markers a line must not start with) and the + table rendering. + """ + + headings: ClassVar[dict[str, str]] = {} + bullet: ClassVar[str] = "- " + block_break: ClassVar[str] = "\n\n" + page_break: ClassVar[str] = "\n\n" + + def escape_text(self, text: str) -> str: + """Escape the characters carrying a structural meaning in the target format""" + return text + + def finalize_line(self, line: str) -> str: + """Neutralize the block-level markers a line must not start with in the target format""" + return line + + def render_table(self, table: "Table", escape: bool = True) -> str: + """Render a recognized table in the target format""" + raise NotImplementedError + + def class_header(self, class_name: str, escape: bool = True) -> str: + """Render the header of a detection class in a KIE export""" + raise NotImplementedError + + def _line_text(self, line: "Line", direction: str, escape: bool) -> str: + """Render the text of a line, ordering the words according to the reading direction.""" + text = " ".join(word.render() for word in ordered_line_words(line, direction)) + return self.escape_text(text) if escape else text + + def export_page( + self, + page: "Page", + direction: str = "auto", + escape: bool = True, + include_furniture: bool = True, + block_break: str | None = None, + ) -> str: + """Export a page, with its content sorted in reading order. + + Args: + page: the page to export + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + escape: whether the characters or markers carrying a structural meaning should be neutralized + include_furniture: whether page headers, page footers and footnotes should be included + block_break: the string inserted between two blocks (the format-specific default when None) + + Returns: + the exported page as a string + """ + from doctr.io.elements import Table + from doctr.models.reading_order import layout_label_role, normalize_layout_label + + auto = direction == "auto" + items, labels, direction = page_reading_order(page, direction) + parts: list[str] = [] + list_group: list[str] = [] + + def _flush_list() -> None: + if list_group: + parts.append("\n".join(list_group)) + list_group.clear() + + for item, label in zip(items, labels): + if not include_furniture and layout_label_role(label) in ("header", "footer", "footnote"): + continue + if isinstance(item, Table): + _flush_list() + rendered = self.render_table(item, escape=escape) + if rendered: + parts.append(rendered) + continue + item_lines = [ + self._line_text(line, _line_render_direction(line, direction, auto), escape) for line in item.lines + ] + item_lines = [line for line in item_lines if line.strip()] + if len(item_lines) == 0: + continue + norm_label = normalize_layout_label(label) + if norm_label in self.headings: + _flush_list() + parts.append(self.headings[norm_label] + " ".join(item_lines)) + elif norm_label in _LIST_LABELS: + # A list item (possibly wrapped over several lines) renders as a single bullet + text = " ".join(item_lines) + list_group.append(self.bullet + (self.finalize_line(text) if escape else text)) + else: + _flush_list() + parts.append("\n".join(self.finalize_line(line) if escape else line for line in item_lines)) + _flush_list() + return (self.block_break if block_break is None else block_break).join(parts) + + def export_kie_page(self, page: "KIEPage", direction: str = "auto", escape: bool = True) -> str: + """Export a KIE page, with the predictions of each class sorted in reading order. + + Args: + page: the KIE page to export + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + escape: whether the characters or markers carrying a structural meaning should be neutralized + + Returns: + the exported page as a string, with one section per detection class + """ + parts: list[str] = [] + for class_name, predictions in page.predictions.items(): + if len(predictions) == 0: + continue + values = "\n".join( + self.bullet + (self.finalize_line(self.escape_text(prediction.value)) if escape else prediction.value) + for prediction in predictions_in_reading_order(page, predictions, direction) + ) + parts.append(f"{self.class_header(class_name, escape)}\n\n{values}") + return "\n\n".join(parts) + + def export_document(self, document: Any, page_break: str | None = None, **kwargs: Any) -> str: + """Export a document page by page. + + Args: + document: the document to export + page_break: the string inserted between two pages (a format-specific default when None) + **kwargs: additional keyword arguments passed to the page export + + Returns: + the exported document as a string + """ + from doctr.io.elements import KIEPage + + page_break = self.page_break if page_break is None else page_break + return page_break.join( + self.export_kie_page(page, **kwargs) if isinstance(page, KIEPage) else self.export_page(page, **kwargs) + for page in document.pages + ) + + +class TextExporter(_PageTextExporter): + """Export OCR results to plain text, with the content sorted in reading order. + + >>> from doctr.io import TextExporter + >>> text = TextExporter().export_page(page) + """ + + headings: ClassVar[dict[str, str]] = {} + bullet: ClassVar[str] = "" + block_break: ClassVar[str] = "\n\n" + page_break: ClassVar[str] = "\n\n\n\n" + + def render_table(self, table: "Table", escape: bool = True) -> str: + """Render a table as tab-separated values, one line per row""" + return table.render() + + def class_header(self, class_name: str, escape: bool = True) -> str: + return f"{class_name}:" + + +class MarkdownExporter(_PageTextExporter): + """Export OCR results to Markdown, with the content sorted in reading order. + + >>> from doctr.io import MarkdownExporter + >>> markdown = MarkdownExporter().export_page(page) + """ + + headings: ClassVar[dict[str, str]] = {"title": "# ", "section_header": "## "} + bullet: ClassVar[str] = "- " + page_break: ClassVar[str] = "\n\n---\n\n" + + def escape_text(self, text: str) -> str: + return "".join(f"\\{char}" if char in _MD_SPECIAL_CHARS else char for char in text) + + def finalize_line(self, line: str) -> str: + stripped = line.lstrip() + if stripped and (stripped[0] in _MD_LINE_MARKERS or stripped.split(" ")[0].rstrip(".").isdigit()): + return f"\\{line}" if line[0] != "\\" else line + return line + + def render_table(self, table: "Table", escape: bool = True) -> str: + """Render a table as a GitHub-flavored Markdown table (first row used as header)""" + grid = table.to_grid() + if len(grid) == 0 or len(grid[0]) == 0: + return "" + + def _cell(value: str) -> str: + value = self.escape_text(value) if escape else value.replace("|", "\\|") + return value.replace("\n", " ").strip() + + rows = ["| " + " | ".join(_cell(value) for value in row) + " |" for row in grid] + separator = "| " + " | ".join("---" for _ in grid[0]) + " |" + return "\n".join([rows[0], separator, *rows[1:]]) + + def class_header(self, class_name: str, escape: bool = True) -> str: + return f"**{self.escape_text(class_name) if escape else class_name}**" + + +class AsciiDocExporter(_PageTextExporter): + """Export OCR results to AsciiDoc, with the content sorted in reading order. + + >>> from doctr.io import AsciiDocExporter + >>> asciidoc = AsciiDocExporter().export_page(page) + """ + + headings: ClassVar[dict[str, str]] = {"title": "== ", "section_header": "=== "} + bullet: ClassVar[str] = "* " + page_break: ClassVar[str] = "\n\n<<<\n\n" + + def escape_text(self, text: str) -> str: + return "".join(f"\\{char}" if char in _ADOC_SPECIAL_CHARS else char for char in text) + + def finalize_line(self, line: str) -> str: + stripped = line.lstrip() + if stripped and stripped[0] in _ADOC_LINE_MARKERS: + return f"{{empty}}{line}" + return line + + def render_table(self, table: "Table", escape: bool = True) -> str: + """Render a table as an AsciiDoc table (first row used as header)""" + grid = table.to_grid() + if len(grid) == 0 or len(grid[0]) == 0: + return "" + + def _row(row: list[str]) -> str: + return " ".join( + "|" + (self.escape_text(value) if escape else value.replace("|", "\\|")).replace("\n", " ").strip() + for value in row + ) + + return "\n".join(["|===", _row(grid[0]), "", *[_row(row) for row in grid[1:]], "|==="]) + + def class_header(self, class_name: str, escape: bool = True) -> str: + return f"*{self.escape_text(class_name) if escape else class_name}*" + + +class HTMLExporter(_PageTextExporter): + """Export OCR results to semantic HTML, with the content sorted in reading order. + + Headings map to `

`/`

`, list items to `
  • `, recognized tables to `` and + paragraphs to `

    ` (with `
    ` between the visual lines of a paragraph). The output is a + fragment, not a full document: it carries no doctype, `` or charset declaration. + + .. warning:: + The recognized text is HTML-escaped by default. Passing ``escape=False`` interpolates the OCR + output into the markup verbatim, so a document containing markup yields active HTML. + Only disable escaping for output that is never rendered in a browser. + + >>> from doctr.io import HTMLExporter + >>> html = HTMLExporter().export_page(page) + """ + + headings: ClassVar[dict[str, str]] = {"title": "h1", "section_header": "h2"} + block_break: ClassVar[str] = "\n" + page_break: ClassVar[str] = "\n

    \n" + + def escape_text(self, text: str) -> str: + return _html_escape(text, quote=False) + + def export_page( + self, + page: "Page", + direction: str = "auto", + escape: bool = True, + include_furniture: bool = True, + block_break: str | None = None, + ) -> str: + from doctr.io.elements import Table + from doctr.models.reading_order import layout_label_role, normalize_layout_label + + auto = direction == "auto" + items, labels, direction = page_reading_order(page, direction) + parts: list[str] = [] + list_group: list[str] = [] + + def _flush_list() -> None: + if list_group: + parts.append("
      \n" + "\n".join(list_group) + "\n
    ") + list_group.clear() + + for item, label in zip(items, labels): + if not include_furniture and layout_label_role(label) in ("header", "footer", "footnote"): + continue + if isinstance(item, Table): + _flush_list() + rendered = self.render_table(item, escape=escape) + if rendered: + parts.append(rendered) + continue + item_lines = [ + self._line_text(line, _line_render_direction(line, direction, auto), escape) for line in item.lines + ] + item_lines = [line for line in item_lines if line.strip()] + if len(item_lines) == 0: + continue + norm_label = normalize_layout_label(label) + if norm_label in self.headings: + _flush_list() + tag = self.headings[norm_label] + parts.append(f"<{tag}>{' '.join(item_lines)}") + elif norm_label in _LIST_LABELS: + list_group.append(f"
  • {' '.join(item_lines)}
    • ") + else: + _flush_list() + parts.append("

    " + "
    \n".join(item_lines) + "

    ") + _flush_list() + return (self.block_break if block_break is None else block_break).join(parts) + + def render_table(self, table: "Table", escape: bool = True) -> str: + """Render a table as an HTML table (first row used as header)""" + grid = table.to_grid() + if len(grid) == 0 or len(grid[0]) == 0: + return "" + + def _cell(value: str, tag: str) -> str: + content = self.escape_text(value) if escape else value + return f"<{tag}>{content.strip()}" + + head = "" + "".join(_cell(value, "th") for value in grid[0]) + "" + body = "\n".join("" + "".join(_cell(value, "td") for value in row) + "" for row in grid[1:]) + return f"
    \n{head}\n{body}\n
    " if body else f"\n{head}\n
    " + + def export_kie_page(self, page: "KIEPage", direction: str = "auto", escape: bool = True) -> str: + parts: list[str] = [] + for class_name, predictions in page.predictions.items(): + if len(predictions) == 0: + continue + values = "\n".join( + f"
  • {self.escape_text(prediction.value) if escape else prediction.value}
    • " + for prediction in predictions_in_reading_order(page, predictions, direction) + ) + header = self.escape_text(class_name) if escape else class_name + parts.append(f"

    {header}

    \n
      \n{values}\n
    ") + return "\n".join(parts) + + +def _resolve_hocr_language(language: dict[str, Any]) -> str: + """Resolve the language code to use in the hOCR export, falling back to 'en'. + + Args: + language: the page language dictionary `{"value": str | None, "confidence": float | None}` + + Returns: + the detected language code when available, 'en' otherwise + """ + lang_value = language.get("value") if isinstance(language, dict) else None + return lang_value if isinstance(lang_value, str) and len(lang_value) > 0 else "en" + + +def _hocr_bbox(geometry: BoundingBox, width: int, height: int) -> str: + """Format a relative straight bounding box as an absolute hOCR `bbox` property string. + + Args: + geometry: the relative bounding box ((xmin, ymin), (xmax, ymax)) + width: the page width in pixels + height: the page height in pixels + + Returns: + the hOCR `bbox` property string + """ + (xmin, ymin), (xmax, ymax) = geometry + return ( + f"bbox {int(round(xmin * width))} {int(round(ymin * height))} " + f"{int(round(xmax * width))} {int(round(ymax * height))}" + ) + + +class XMLExporter: + """hOCR (XML) exporter for pages, KIE pages and documents. + See the hOCR 1.2 specification for the XML convention: https://github.com/kba/hocr-spec/blob/master/1.2/spec.md + + >>> from doctr.io import XMLExporter + >>> xml_bytes, xml_tree = XMLExporter().export_page(page) + """ + + ocr_capabilities: ClassVar[str] = "ocr_page ocr_carea ocr_par ocr_line ocrx_word" + + def _new_document(self, file_title: str, language: str) -> tuple[ETElement, ETElement]: + """Create the hOCR root element with its , returning the root and its element.""" + root = ETElement("html", attrib={"xmlns": "http://www.w3.org/1999/xhtml", "xml:lang": str(language)}) + head = SubElement(root, "head") + SubElement(head, "title").text = file_title + SubElement(head, "meta", attrib={"http-equiv": "Content-Type", "content": "text/html; charset=utf-8"}) + SubElement( + head, + "meta", + attrib={"name": "ocr-system", "content": f"python-doctr {doctr.__version__}"}, # type: ignore[attr-defined] + ) + SubElement(head, "meta", attrib={"name": "ocr-capabilities", "content": self.ocr_capabilities}) + return root, SubElement(root, "body") + + def _add_table(self, page_div: ETElement, table: "Table", width: int, height: int, table_count: int) -> int: + """Serialize a recognized table as an hOCR text area, with one `ocr_line` per row. + + Args: + page_div: the `ocr_page` element the table is appended to + table: the table to serialize + width: the page width in pixels + height: the page height in pixels + table_count: the 1-based index of the table on the page + + Returns: + the index of the next table + """ + if len(table.geometry) != 2 or any(len(cell.geometry) != 2 for cell in table.cells): + raise TypeError("XML export is only available for straight bounding boxes for now.") + table_bbox = _hocr_bbox(table.geometry, width, height) # type: ignore[arg-type] + table_div = SubElement( + page_div, "div", attrib={"class": "ocr_carea", "id": f"table_{table_count}", "title": table_bbox} + ) + paragraph = SubElement( + table_div, "p", attrib={"class": "ocr_par", "id": f"table_par_{table_count}", "title": table_bbox} + ) + rows: dict[int, list[Any]] = {} + for cell in table.cells: + rows.setdefault(cell.row_start, []).append(cell) + for row_idx in sorted(rows): + cells = sorted(rows[row_idx], key=lambda cell: cell.col_start) + xs = [coord for cell in cells for coord in (cell.geometry[0][0], cell.geometry[1][0])] + ys = [coord for cell in cells for coord in (cell.geometry[0][1], cell.geometry[1][1])] + row_bbox = _hocr_bbox(((min(xs), min(ys)), (max(xs), max(ys))), width, height) + line_span = SubElement( + paragraph, + "span", + attrib={ + "class": "ocr_line", + "id": f"table_{table_count}_row_{row_idx + 1}", + "title": f"{row_bbox}; baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0", + }, + ) + for col_idx, cell in enumerate(cells): + cell_span = SubElement( + line_span, + "span", + attrib={ + "class": "ocrx_word", + "id": f"table_{table_count}_cell_{row_idx + 1}_{col_idx + 1}", + "title": ( + f"{_hocr_bbox(cell.geometry, width, height)}; x_wconf {int(round(cell.confidence * 100))}" + ), + }, + ) + cell_span.text = cell.value + return table_count + 1 + + def export_page( + self, + page: "Page", + file_title: str = "docTR - XML export (hOCR)", + direction: str = "auto", + reading_order: bool = True, + ) -> tuple[bytes, ET.ElementTree]: + """Export a page as hOCR XML, with its content sorted in reading order. + + Args: + page: the page to export + file_title: the title of the XML file + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + reading_order: whether the content should be linearized in reading order. Pass False to serialize + `page.blocks` then `page.tables` in their raw order. + + Returns: + a tuple of the XML byte string, and its ElementTree + """ + from doctr.io.elements import Table + + block_count: int = 1 + line_count: int = 1 + word_count: int = 1 + table_count: int = 1 + height, width = page.dimensions + page_hocr, body = self._new_document(file_title, _resolve_hocr_language(page.language)) + page_div = SubElement( + body, + "div", + attrib={ + "class": "ocr_page", + "id": f"page_{page.page_idx + 1}", + "title": f"image; bbox 0 0 {width} {height}; ppageno 0", + }, + ) + auto = direction == "auto" + if reading_order: + items, _, direction = page_reading_order(page, direction) + else: + items = [*page.blocks, *page.tables] + # iterate over the blocks / lines / words and create the XML elements line by line with the attributes + for item in items: + if isinstance(item, Table): + table_count = self._add_table(page_div, item, width, height, table_count) + continue + block = item + if len(block.geometry) != 2: + raise TypeError("XML export is only available for straight bounding boxes for now.") + block_bbox = _hocr_bbox(block.geometry, width, height) + block_div = SubElement( + page_div, + "div", + attrib={"class": "ocr_carea", "id": f"block_{block_count}", "title": block_bbox}, + ) + paragraph = SubElement( + block_div, + "p", + attrib={"class": "ocr_par", "id": f"par_{block_count}", "title": block_bbox}, + ) + block_count += 1 + for line in block.lines: + # NOTE: baseline, x_size, x_descenders, x_ascenders is currently initalized to 0 + line_span = SubElement( + paragraph, + "span", + attrib={ + "class": "ocr_line", + "id": f"line_{line_count}", + "title": ( + f"{_hocr_bbox(line.geometry, width, height)}; " + "baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0" + ), + }, + ) + line_count += 1 + for word in ordered_line_words(line, direction, auto): + word_div = SubElement( + line_span, + "span", + attrib={ + "class": "ocrx_word", + "id": f"word_{word_count}", + "title": ( + f"{_hocr_bbox(word.geometry, width, height)}; " + f"x_wconf {int(round(word.confidence * 100))}" + ), + }, + ) + word_div.text = word.value + word_count += 1 + return ET.tostring(page_hocr, encoding="utf-8", method="xml"), ET.ElementTree(page_hocr) + + def export_kie_page( + self, + page: "KIEPage", + file_title: str = "docTR - XML export (hOCR)", + direction: str = "auto", + reading_order: bool = True, + ) -> tuple[bytes, ET.ElementTree]: + """Export a KIE page as hOCR XML, with the predictions of each class sorted in reading order. + + Args: + page: the KIE page to export + file_title: the title of the XML file + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + reading_order: whether the predictions of each class should be sorted in reading order + + Returns: + a tuple of the XML byte string, and its ElementTree + """ + prediction_count: int = 1 + height, width = page.dimensions + page_hocr, body = self._new_document(file_title, _resolve_hocr_language(page.language)) + SubElement( + body, + "div", + attrib={ + "class": "ocr_page", + "id": f"page_{page.page_idx + 1}", + "title": f"image; bbox 0 0 {width} {height}; ppageno 0", + }, + ) + # iterate over the predictions and create the XML elements line by line with the attributes + for class_name, predictions in page.predictions.items(): + ordered = predictions_in_reading_order(page, predictions, direction) if reading_order else predictions + for prediction in ordered: + if len(prediction.geometry) != 2: + raise TypeError("XML export is only available for straight bounding boxes for now.") + prediction_bbox = _hocr_bbox(prediction.geometry, width, height) # type: ignore[arg-type] + prediction_div = SubElement( + body, + "div", + attrib={ + "class": "ocr_carea", + "id": f"{class_name}_prediction_{prediction_count}", + "title": prediction_bbox, + }, + ) + # NOTE: ocr_par, ocr_line and ocrx_word are the same because the KIE predictions contain only words + # This is a workaround to make it PDF/A compatible + par_div = SubElement( + prediction_div, + "p", + attrib={ + "class": "ocr_par", + "id": f"{class_name}_par_{prediction_count}", + "title": prediction_bbox, + }, + ) + line_span = SubElement( + par_div, + "span", + attrib={ + "class": "ocr_line", + "id": f"{class_name}_line_{prediction_count}", + "title": f"{prediction_bbox}; baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0", + }, + ) + word_div = SubElement( + line_span, + "span", + attrib={ + "class": "ocrx_word", + "id": f"{class_name}_word_{prediction_count}", + "title": f"{prediction_bbox}; x_wconf {int(round(prediction.confidence * 100))}", + }, + ) + word_div.text = prediction.value + prediction_count += 1 + return ET.tostring(page_hocr, encoding="utf-8", method="xml"), ET.ElementTree(page_hocr) + + def export_document(self, document: Any, **kwargs: Any) -> list[tuple[bytes, ET.ElementTree]]: + """Export a document as a list of hOCR pages. + + Args: + document: the document to export + **kwargs: additional keyword arguments passed to the page export + + Returns: + list of tuple of (bytes, ElementTree), one per page + """ + from doctr.io.elements import KIEPage + + return [ + self.export_kie_page(page, **kwargs) if isinstance(page, KIEPage) else self.export_page(page, **kwargs) + for page in document.pages + ] + + +class PageExportsMixin: + """Export functionality of a :class:`~doctr.io.elements.Page`""" + + if TYPE_CHECKING: # structural attributes provided by the element class + page: np.ndarray + blocks: list["Block"] + page_idx: int + dimensions: tuple[int, int] + orientation: dict[str, Any] + language: dict[str, Any] + layout: list[Any] + tables: list["Table"] + + def render(self, block_break: str = "\n\n", direction: str = "auto", include_furniture: bool = True) -> str: + """Renders the full text of the page, with its content sorted in reading order. + + Args: + block_break: the string inserted between two blocks + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + include_furniture: whether page headers, page footers and footnotes should be included + + Returns: + the text of the page + """ + return TextExporter().export_page( + cast("Page", self), direction=direction, include_furniture=include_furniture, block_break=block_break + ) + + def export(self, reading_order: bool = True) -> dict[str, Any]: + """Export the page into a nested dict, with its content sorted in reading order. + + Args: + reading_order: whether the blocks should be linearized in reading order, exactly like the + Markdown / HTML / AsciiDoc / hOCR exports. Pass False to serialize `page.blocks` as stored. + + Returns: + a JSON-serializable dict + """ + from doctr.io.elements import Element, Table + + export_dict = Element.export(cast("Element", self)) + if reading_order: + blocks = [item for item in page_reading_order(cast("Page", self))[0] if not isinstance(item, Table)] + if blocks: # an empty linearization (no line on the page) leaves the stored blocks untouched + export_dict["blocks"] = [block.export() for block in blocks] + return export_dict + + def export_as_xml( + self, + file_title: str = "docTR - XML export (hOCR)", + direction: str = "auto", + reading_order: bool = True, + ) -> tuple[bytes, ET.ElementTree]: + """Export the page as XML (hOCR-format), with its content sorted in reading order + convention: https://github.com/kba/hocr-spec/blob/master/1.2/spec.md + + Args: + file_title: the title of the XML file + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + reading_order: whether the content should be linearized in reading order + + Returns: + a tuple of the XML byte string, and its ElementTree + """ + return XMLExporter().export_page( + cast("Page", self), file_title=file_title, direction=direction, reading_order=reading_order + ) + + def items_in_reading_order(self, direction: str = "auto") -> list["Block | Table"]: + """Return the content of the page (blocks & tables) sorted in reading order. + + Args: + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + + Returns: + list of blocks & tables in reading order + """ + return page_reading_order(cast("Page", self), direction)[0] + + def export_as_markdown(self, direction: str = "auto", escape: bool = True, include_furniture: bool = True) -> str: + """Export the page as Markdown, with its content sorted in reading order. + + Args: + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + escape: whether the characters carrying a structural meaning in Markdown should be escaped + include_furniture: whether page headers, page footers and footnotes should be included + + Returns: + a Markdown string + """ + return MarkdownExporter().export_page( + cast("Page", self), direction=direction, escape=escape, include_furniture=include_furniture + ) + + def export_as_asciidoc(self, direction: str = "auto", escape: bool = True, include_furniture: bool = True) -> str: + """Export the page as AsciiDoc, with its content sorted in reading order. + + Args: + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + escape: whether the characters and line markers carrying a structural meaning in AsciiDoc should + be escaped + include_furniture: whether page headers, page footers and footnotes should be included + + Returns: + an AsciiDoc string + """ + return AsciiDocExporter().export_page( + cast("Page", self), direction=direction, escape=escape, include_furniture=include_furniture + ) + + def export_as_html(self, direction: str = "auto", include_furniture: bool = True) -> str: + """Export the page as semantic HTML, with its content sorted in reading order. + + Args: + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + include_furniture: whether page headers, page footers and footnotes should be included + + Returns: + an HTML string + """ + return HTMLExporter().export_page(cast("Page", self), direction=direction, include_furniture=include_furniture) + + def export_as(self, format: str, **kwargs: Any) -> Any: + """Export the page in the requested format. + + Args: + format: one of 'markdown'/'md', 'asciidoc'/'adoc', 'html', 'text'/'txt', 'json'/'dict', + 'xml'/'hocr' + **kwargs: additional keyword arguments passed to the format-specific export method + + Returns: + the exported page + """ + exporters: dict[str, Any] = { + "markdown": self.export_as_markdown, + "md": self.export_as_markdown, + "asciidoc": self.export_as_asciidoc, + "adoc": self.export_as_asciidoc, + "html": self.export_as_html, + "text": self.render, + "txt": self.render, + "json": self.export, + "dict": self.export, + "xml": self.export_as_xml, + "hocr": self.export_as_xml, + } + return _export_as(exporters, format, **kwargs) + + +class KIEPageExportsMixin: + """Export functionality of a :class:`~doctr.io.elements.KIEPage`""" + + if TYPE_CHECKING: # structural attributes provided by the element class + page: np.ndarray + predictions: dict[str, list[Any]] + page_idx: int + dimensions: tuple[int, int] + orientation: dict[str, Any] + language: dict[str, Any] + + def export(self, reading_order: bool = True) -> dict[str, Any]: + """Export the KIE page into a nested dict, with the predictions of each class in reading order. + + Args: + reading_order: whether the predictions of each class should be sorted in reading order + + Returns: + a JSON-serializable dict + """ + from doctr.io.elements import Element + + export_dict = Element.export(cast("Element", self)) + if reading_order: + export_dict["predictions"] = { + class_name: [ + prediction.export() + for prediction in predictions_in_reading_order(cast("KIEPage", self), predictions) + ] + for class_name, predictions in self.predictions.items() + } + return export_dict + + def render(self, prediction_break: str = "\n\n", direction: str = "auto") -> str: + """Renders the full text of the page, with the predictions of each class sorted in reading order. + + Args: + prediction_break: the string inserted between two predictions + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + + Returns: + the text of the page, one section per detection class with its predictions in reading order + """ + parts: list[str] = [] + for class_name, predictions in self.predictions.items(): + parts.extend( + f"{class_name}: {prediction.render()}" + for prediction in predictions_in_reading_order(cast("KIEPage", self), predictions, direction) + ) + return prediction_break.join(parts) + + def export_as_xml( + self, + file_title: str = "docTR - XML export (hOCR)", + direction: str = "auto", + reading_order: bool = True, + ) -> tuple[bytes, ET.ElementTree]: + """Export the page as XML (hOCR-format), with the predictions of each class in reading order + convention: https://github.com/kba/hocr-spec/blob/master/1.2/spec.md + + Args: + file_title: the title of the XML file + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + reading_order: whether the predictions of each class should be sorted in reading order + + Returns: + a tuple of the XML byte string, and its ElementTree + """ + return XMLExporter().export_kie_page( + cast("KIEPage", self), file_title=file_title, direction=direction, reading_order=reading_order + ) + + def export_as_markdown(self, direction: str = "auto", escape: bool = True) -> str: + """Export the KIE page as Markdown, with the predictions of each class sorted in reading order. + + Args: + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + escape: whether the characters carrying a structural meaning in Markdown should be escaped + + Returns: + a Markdown string with one section per detection class + """ + return MarkdownExporter().export_kie_page(cast("KIEPage", self), direction=direction, escape=escape) + + def export_as_asciidoc(self, direction: str = "auto", escape: bool = True) -> str: + """Export the KIE page as AsciiDoc, with the predictions of each class sorted in reading order. + + Args: + direction: reading direction, one of 'auto', 'ltr', 'rtl', 'ttb-rtl' or 'ttb-ltr' + escape: whether the characters and line markers carrying a structural meaning in AsciiDoc should + be escaped + + Returns: + an AsciiDoc string with one section per detection class + """ + return AsciiDocExporter().export_kie_page(cast("KIEPage", self), direction=direction, escape=escape) + + def export_as_html(self, direction: str = "auto") -> str: + """Export the KIE page as semantic HTML, with the predictions of each class sorted in reading order""" + return HTMLExporter().export_kie_page(cast("KIEPage", self), direction=direction) + + def export_as(self, format: str, **kwargs: Any) -> Any: + """Export the KIE page in the requested format ('markdown'/'md', 'asciidoc'/'adoc', 'html', + 'text'/'txt', 'json'/'dict', 'xml'/'hocr').""" + exporters: dict[str, Any] = { + "markdown": self.export_as_markdown, + "md": self.export_as_markdown, + "asciidoc": self.export_as_asciidoc, + "adoc": self.export_as_asciidoc, + "html": self.export_as_html, + "text": self.render, + "txt": self.render, + "json": self.export, + "dict": self.export, + "xml": self.export_as_xml, + "hocr": self.export_as_xml, + } + return _export_as(exporters, format, **kwargs) + + +class DocumentExportsMixin: + """Export functionality of a :class:`~doctr.io.elements.Document` (also used by `KIEDocument`)""" + + if TYPE_CHECKING: # structural attributes provided by the element class + pages: list[Any] + _exported_keys: list[str] + + def render(self, page_break: str = "\n\n\n\n", **kwargs: Any) -> str: + """Renders the full text of the document, with the content of each page sorted in reading order. + + Args: + page_break: the string inserted between two pages + **kwargs: additional keyword arguments passed to the `Page.render` / `KIEPage.render` method + + Returns: + the text of the document + """ + return page_break.join(page.render(**kwargs) for page in self.pages) + + def export(self, reading_order: bool = True) -> dict[str, Any]: + """Export the document into a nested dict, with the content of each page sorted in reading order. + + Args: + reading_order: whether the content of each page should be linearized in reading order + + Returns: + a JSON-serializable dict + """ + export_dict: dict[str, Any] = {key: to_json_safe(getattr(self, key)) for key in self._exported_keys} + export_dict["pages"] = [page.export(reading_order=reading_order) for page in self.pages] + return export_dict + + def export_as_xml(self, **kwargs: Any) -> list[tuple[bytes, ET.ElementTree]]: + """Export the document as XML (hOCR-format) + + Args: + **kwargs: additional keyword arguments passed to the XML page export + + Returns: + list of tuple of (bytes, ElementTree) + """ + return XMLExporter().export_document(self, **kwargs) + + def export_as_markdown(self, page_break: str = "\n\n---\n\n", **kwargs: Any) -> str: + """Export the document as Markdown, with the content of each page sorted in reading order. + + Args: + page_break: the string inserted between two pages (a thematic break by default) + **kwargs: additional keyword arguments passed to the `Page.export_as_markdown` method + + Returns: + a Markdown string + """ + return page_break.join(page.export_as_markdown(**kwargs) for page in self.pages) + + def export_as_asciidoc(self, page_break: str = "\n\n<<<\n\n", **kwargs: Any) -> str: + """Export the document as AsciiDoc, with the content of each page sorted in reading order. + + Args: + page_break: the string inserted between two pages (an AsciiDoc page break by default) + **kwargs: additional keyword arguments passed to the `Page.export_as_asciidoc` method + + Returns: + an AsciiDoc string + """ + return page_break.join(page.export_as_asciidoc(**kwargs) for page in self.pages) + + def export_as_html(self, page_break: str = "
    ", **kwargs: Any) -> str: + """Export the document as semantic HTML, with the content of each page sorted in reading order. + + Args: + page_break: the HTML snippet inserted between two pages + **kwargs: additional keyword arguments passed to the page export + + Returns: + an HTML string + """ + return page_break.join(page.export_as_html(**kwargs) for page in self.pages) + + def export_as(self, format: str, **kwargs: Any) -> Any: + """Export the document in the requested format ('markdown'/'md', 'asciidoc'/'adoc', 'html', + 'text'/'txt', 'json'/'dict', 'xml'/'hocr').""" + exporters: dict[str, Any] = { + "markdown": self.export_as_markdown, + "md": self.export_as_markdown, + "asciidoc": self.export_as_asciidoc, + "adoc": self.export_as_asciidoc, + "html": self.export_as_html, + "text": self.render, + "txt": self.render, + "json": self.export, + "dict": self.export, + "xml": self.export_as_xml, + "hocr": self.export_as_xml, + } + return _export_as(exporters, format, **kwargs) diff --git a/doctr/models/builder.py b/doctr/models/builder.py index 149fc3c7d8..5067c500fb 100644 --- a/doctr/models/builder.py +++ b/doctr/models/builder.py @@ -23,7 +23,13 @@ TableCell, Word, ) -from doctr.utils.geometry import estimate_page_angle, resolve_enclosing_bbox, resolve_enclosing_rbbox, rotate_boxes +from doctr.utils.geometry import ( + estimate_page_angle, + order_points, + resolve_enclosing_bbox, + resolve_enclosing_rbbox, + rotate_boxes, +) from doctr.utils.repr import NestedObject __all__ = ["DocumentBuilder"] @@ -38,6 +44,10 @@ class DocumentBuilder(NestedObject): paragraph_break: relative length of the minimum space separating paragraphs export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. + keep_reading_order: if True, arrange the content of every page in reading order (cf. + :mod:`doctr.models.reading_order`). The reading direction is inferred from the recognized text. + When a layout is available it is always preferred: the text blocks and the page furniture + (headers, footers, ...) are taken from the layout regions. """ def __init__( @@ -46,18 +56,21 @@ def __init__( resolve_blocks: bool = False, paragraph_break: float = 0.035, export_as_straight_boxes: bool = False, + keep_reading_order: bool = False, ) -> None: self.resolve_lines = resolve_lines self.resolve_blocks = resolve_blocks self.paragraph_break = paragraph_break self.export_as_straight_boxes = export_as_straight_boxes + self.keep_reading_order = keep_reading_order @staticmethod - def _sort_boxes(boxes: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + def _sort_boxes(boxes: np.ndarray, shape: tuple[int, int] | None = None) -> tuple[np.ndarray, np.ndarray]: """Sort bounding boxes from top to bottom, left to right Args: - boxes: bounding boxes of shape (N, 4) or (N, 4, 2) (in case of rotated bbox) + boxes: bounding boxes of shape (N, 4) or (N, 4, 2) + shape: the page dimensions (height, width) Returns: tuple: indices of ordered boxes of shape (N,), boxes @@ -66,28 +79,45 @@ def _sort_boxes(boxes: np.ndarray) -> tuple[np.ndarray, np.ndarray]: so that we fit the lines afterwards to the straigthened page """ if boxes.ndim == 3: - boxes = rotate_boxes( + height, width = shape if shape is not None else (1024, 1024) + scale = np.array([width, height], dtype=boxes.dtype) + angle = estimate_page_angle(boxes * scale) + rotated = rotate_boxes( loc_preds=boxes, - angle=-estimate_page_angle(boxes), - orig_shape=(1024, 1024), - min_angle=5.0, + angle=-angle, + orig_shape=(height, width), + min_angle=1.0, ) - boxes = np.concatenate((boxes.min(1), boxes.max(1)), -1) + # On rotated pages, detectors can output a mix of properly rotated polygons and axis-aligned + # enclosing boxes. A box that is not itself rotated carries no rotation to remove + if abs(angle) >= 1.0: + edges = (boxes[:, 1] - boxes[:, 0]) * scale + own_angle = np.rad2deg(np.arctan2(-edges[:, 1], edges[:, 0])) + keep = np.abs(own_angle) < abs(angle) / 2 + if keep.any(): + centers = boxes.mean(axis=1, keepdims=True) + new_centers = rotated.mean(axis=1, keepdims=True) + rotated[keep] = boxes[keep] - centers[keep] + new_centers[keep] + boxes = np.concatenate((rotated.min(1), rotated.max(1)), -1) med_height = float(np.median(boxes[:, 3] - boxes[:, 1])) if not np.isfinite(med_height) or med_height <= 0: med_height = 1.0 return (boxes[:, 0] + 2 * boxes[:, 3] / med_height).argsort(), boxes - def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: list[int]) -> list[list[int]]: + def _resolve_sub_lines( + self, boxes: np.ndarray, word_idcs: list[int], break_dist: float | None = None + ) -> list[list[int]]: """Split a line in sub_lines Args: boxes: bounding boxes of shape (N, 4) word_idcs: list of indexes for the words of the line + break_dist: horizontal distance above which the line is split (defaults to `self.paragraph_break`) Returns: 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()] @@ -103,8 +133,8 @@ def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: list[int]) -> list[li prev_box = boxes[sub_line[-1]] # Compute distance between boxes dist = boxes[i, 0] - prev_box[2] - # If distance between boxes is lower than paragraph break, same sub-line - if dist < self.paragraph_break: + # If distance between boxes is lower than the break distance, same sub-line + if dist < break_dist: horiz_break = False if horiz_break: @@ -116,22 +146,24 @@ def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: list[int]) -> list[li return lines - def _resolve_lines(self, boxes: np.ndarray) -> list[list[int]]: + def _resolve_lines(self, boxes: np.ndarray, shape: tuple[int, int] | None = None) -> list[list[int]]: """Order boxes to group them in lines Args: boxes: bounding boxes of shape (N, 4) or (N, 4, 2) in case of rotated bbox + shape: the page dimensions (height, width), used to de-skew rotated pages exactly Returns: nested list of box indices """ # Sort boxes, and straighten the boxes if they are rotated - idxs, boxes = self._sort_boxes(boxes) + idxs, boxes = self._sort_boxes(boxes, shape) # Compute median for boxes heights y_med = np.median(boxes[:, 3] - boxes[:, 1]) - lines = [] + # Group the words into visual rows + rows = [] words = [idxs[0]] # Assign the top-left word to the first line # Define a mean y-center for the line y_center_sum = boxes[idxs[0]][[1, 3]].mean() @@ -146,18 +178,43 @@ def _resolve_lines(self, boxes: np.ndarray) -> list[list[int]]: vert_break = False if vert_break: - # Compute sub-lines (horizontal split) - lines.extend(self._resolve_sub_lines(boxes, words)) + rows.append(words) words = [] y_center_sum = 0 words.append(idx) y_center_sum += boxes[idx][[1, 3]].mean() - # Use the remaining words to form the last(s) line(s) + # Use the remaining words to form the last row if len(words) > 0: - # Compute sub-lines (horizontal split) - lines.extend(self._resolve_sub_lines(boxes, words)) + rows.append(words) + + # Split the rows horizontally, using a break distance adapted to the page spacing. + # The gaps are collected with one vectorized pass per row (the loop is over rows, not word pairs). + gap_chunks = [] + n_pairs = 0 + for row in rows: + if len(row) < 2: + continue + row_idcs = np.asarray(row) + row_idcs = row_idcs[np.argsort(boxes[row_idcs, 0], kind="stable")] + n_pairs += row_idcs.shape[0] - 1 + gap_chunks.append(boxes[row_idcs[1:], 0] - boxes[row_idcs[:-1], 2]) + all_gaps = np.concatenate(gap_chunks) if gap_chunks else np.empty(0, dtype=boxes.dtype) + pos_gaps = all_gaps[all_gaps > 0] + # Median word height, converted to width units + aspect = (shape[0] / shape[1]) if shape is not None else 1.0 + floor = float(y_med) * aspect + if pos_gaps.shape[0] >= 5 and pos_gaps.shape[0] >= 0.5 * n_pairs: + break_dist = min(self.paragraph_break, max(3.0 * float(np.median(pos_gaps)), floor)) + elif n_pairs >= 5: + break_dist = min(self.paragraph_break, floor) + else: + break_dist = self.paragraph_break + + lines = [] + for row in rows: + lines.extend(self._resolve_sub_lines(boxes, row, break_dist)) return lines @@ -285,7 +342,7 @@ def _as_cell_polygon(geometry: Any) -> np.ndarray: if arr.ndim == 1: # straight box (xmin, ymin, xmax, ymax) x0, y0, x1, y1 = arr return np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype=np.float32) - return arr.reshape(-1, 2) + return order_points(arr.reshape(-1, 2)) @staticmethod def _points_in_polygons(points: np.ndarray, polys: np.ndarray) -> np.ndarray: @@ -310,6 +367,33 @@ def _points_in_polygons(points: np.ndarray, polys: np.ndarray) -> np.ndarray: # A point is inside when a ray crosses the polygon boundary an odd number of times return (crossing.sum(axis=-1) % 2).astype(bool) + @staticmethod + def _order_cell_words(w_idcs: list[int], centers: np.ndarray, heights: np.ndarray) -> list[int]: + """Order the words of a table cell in reading order: rows top to bottom, words left to right. + + Args: + w_idcs: indices of the cell's words + centers: word centers of shape (N, 2), de-skewed for rotated pages + heights: per-word heights of shape (N,) (rotation-invariant) + + Returns: + the cell's word indices in reading order + """ + idcs = sorted(w_idcs, key=lambda i: float(centers[i][1])) + med_height = float(np.median([heights[i] for i in idcs])) + if not np.isfinite(med_height) or med_height <= 0: + med_height = 1.0 + rows: list[list[int]] = [[idcs[0]]] + y_sum = float(centers[idcs[0]][1]) + for idx in idcs[1:]: + if float(centers[idx][1]) - y_sum / len(rows[-1]) < med_height / 2: + rows[-1].append(idx) + y_sum += float(centers[idx][1]) + else: + rows.append([idx]) + y_sum = float(centers[idx][1]) + return [idx for row in rows for idx in sorted(row, key=lambda i: float(centers[i][0]))] + @staticmethod def _localize_logic(cells: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], int, int]: """Re-index a table's logical coordinates to a local 0-based grid. @@ -345,12 +429,6 @@ def _build_tables( ) -> tuple[list[Table], np.ndarray]: """Assign detected words to table cells and build the page tables. - A page may contain several tables; each one is provided as its own grid (the OCR pipeline detects table - regions with the layout model, then runs the table model on every cropped region). Both a single grid and - a list of grids are accepted. Each word whose center falls inside a cell polygon is assigned to (at most) - one cell, across all tables, and flagged so it can be removed from the regular `blocks` output. Words - are joined per cell in reading order (top to bottom, then left to right). - Args: boxes: word boxes of the page, of shape (N, 4) or (N, 4, 2), in relative coordinates word_preds: list of (text, confidence) for each of the N words @@ -382,6 +460,13 @@ def _build_tables( # 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). 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] + else: + word_heights = np.linalg.norm(boxes[:, 3] - boxes[:, 0], axis=1) tables_out: list[Table] = [] for table_dict in table_dicts: @@ -389,25 +474,66 @@ def _build_tables( 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: the first cell (in cell - # order) whose polygon contains the word center + # 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 and len(cell_polys) > 0: - inside = self._points_in_polygons(centers[free_idcs], np.stack(cell_polys)) # (F, C) - first_cell = np.where(inside.any(axis=1), inside.argmax(axis=1), -1) + 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): if len(w_idcs) > 0: - # Reading order inside the cell: top to bottom, then left to right - ordered = sorted(w_idcs, key=lambda i: (round(float(centers[i][1]), 3), float(centers[i][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) value = " ".join(word_preds[i][0] for i in ordered) confidence = float(np.mean([word_preds[i][1] for i in ordered])) else: @@ -459,6 +585,7 @@ def _build_blocks( objectness_scores: np.ndarray, word_preds: list[tuple[str, float]], crop_orientations: list[dict[str, Any]], + shape: tuple[int, int] | None = None, ) -> list[Block]: """Gather independent words in structured blocks @@ -468,6 +595,8 @@ def _build_blocks( word_preds: list of all detected words of the page, of shape N crop_orientations: list of dictoinaries containing the general orientation (orientations + confidences) of the crops + shape: the page dimensions (height, width), used to de-skew rotated pages exactly when + resolving lines Returns: list of block elements @@ -481,7 +610,7 @@ def _build_blocks( # Decide whether we try to form lines _boxes = boxes if self.resolve_lines: - lines = self._resolve_lines(_boxes if _boxes.ndim == 3 else _boxes[:, :4]) + lines = self._resolve_lines(_boxes if _boxes.ndim == 3 else _boxes[:, :4], shape) # 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) @@ -489,7 +618,7 @@ def _build_blocks( _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])[0]] # type: ignore[list-item] + lines = [self._sort_boxes(_boxes if _boxes.ndim == 3 else _boxes[:, :4], shape)[0]] # type: ignore[list-item] _blocks = [lines] blocks = [ @@ -517,11 +646,54 @@ def _build_blocks( return blocks + @staticmethod + def _words_to_boxes(words: list[Word]) -> np.ndarray: + """Rebuild a localization array from word geometries (rotated polygons or straight boxes).""" + geoms = [np.asarray(word.geometry, dtype=np.float32) for word in words] + if len(geoms) > 0 and geoms[0].shape == (4, 2): + return np.stack(geoms) # (N, 4, 2) rotated polygons + return np.asarray( + [[g[0][0], g[0][1], g[1][0], g[1][1]] for g in geoms], dtype=np.float32 + ) # (N, 4) straight boxes + + def _apply_reading_order(self, page: Page) -> None: + """Arrange the blocks of a page in reading order, in place. + + Args: + page: the page whose blocks should be reordered in place + """ + from doctr.io.exporters import _reading_order_signature, _store_reading_order, page_reading_order + + if not page.blocks: + return + + collapse = not self.resolve_lines + if collapse: + words = [word for block in page.blocks for line in block.lines for word in line.words] + if len(words) < 2: + return + groups = self._resolve_lines(self._words_to_boxes(words), page.dimensions) + page.blocks = [Block([Line([words[idx] for idx in group]) for group in groups])] + + items, labels, direction = page_reading_order(page) + blocks = [item for item in items if isinstance(item, Block)] + page.blocks = ( + [Block(lines=[Line([word for block in blocks for line in block.lines for word in line.words])])] + if collapse + else blocks + ) + if not collapse: + # The page now *is* the reading order, so an export would recompute the very same linearization. + # Re-key the cached result to the new block list so the first export reuses it. + # Skipped when collapsing: the single flattened line no longer matches the line-level items. + _store_reading_order(page, _reading_order_signature(page, "auto"), (items, labels, direction)) + def extra_repr(self) -> str: return ( f"resolve_lines={self.resolve_lines}, resolve_blocks={self.resolve_blocks}, " f"paragraph_break={self.paragraph_break}, " - f"export_as_straight_boxes={self.export_as_straight_boxes}" + f"export_as_straight_boxes={self.export_as_straight_boxes}, " + f"keep_reading_order={self.keep_reading_order}" ) def __call__( @@ -614,23 +786,27 @@ def __call__( word_preds = [wp for wp, k in zip(word_preds, keep) if k] word_crop_orientations = [co for co, k in zip(word_crop_orientations, keep) if k] - _pages.append( - Page( - page, - self._build_blocks( - page_boxes, - loc_scores, - word_preds, - word_crop_orientations, - ), - _idx, - shape, - orientation, - language, - self._build_layout_elements(page_regions), - page_tables, - ) + page_blocks = self._build_blocks( + page_boxes, + loc_scores, + word_preds, + word_crop_orientations, + shape, ) + _page = Page( + page, + page_blocks, + _idx, + shape, + orientation, + language, + self._build_layout_elements(page_regions), + page_tables, + ) + if self.keep_reading_order: + self._apply_reading_order(_page) + + _pages.append(_page) return Document(_pages) @@ -644,6 +820,10 @@ class KIEDocumentBuilder(DocumentBuilder): paragraph_break: relative length of the minimum space separating paragraphs export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. + keep_reading_order: if True, arrange the content of every page in reading order (cf. + :mod:`doctr.models.reading_order`). The reading direction is inferred from the recognized text. + When a layout is available it is always preferred: the text blocks and the page furniture + (headers, footers, ...) are taken from the layout regions. """ def __call__( # type: ignore[override] @@ -714,6 +894,7 @@ def __call__( # type: ignore[override] loc_scores[k], word_preds[k], word_crop_orientations[k], + shape, ) for k in page_boxes.keys() }, @@ -745,6 +926,7 @@ def _build_blocks( # type: ignore[override] objectness_scores: np.ndarray, word_preds: list[tuple[str, float]], crop_orientations: list[dict[str, Any]], + shape: tuple[int, int] | None = None, ) -> list[Prediction]: """Gather independent words in structured blocks @@ -753,6 +935,7 @@ def _build_blocks( # type: ignore[override] objectness_scores: objectness scores of all detected words of the page word_preds: list of all detected words of the page, of shape N crop_orientations: list of orientations for each word crop + shape: the page dimensions (height, width), used to de-skew rotated pages exactly Returns: list of block elements @@ -765,7 +948,7 @@ def _build_blocks( # type: ignore[override] # Decide whether we try to form lines _boxes = boxes - idxs, _ = self._sort_boxes(_boxes if _boxes.ndim == 3 else _boxes[:, :4]) + idxs, _ = self._sort_boxes(_boxes if _boxes.ndim == 3 else _boxes[:, :4], shape) predictions = [ Prediction( value=word_preds[idx][0], diff --git a/doctr/models/layout/lw_detr/pytorch.py b/doctr/models/layout/lw_detr/pytorch.py index f8165f8aad..ce55853400 100644 --- a/doctr/models/layout/lw_detr/pytorch.py +++ b/doctr/models/layout/lw_detr/pytorch.py @@ -45,8 +45,6 @@ "Table", "Text", "Title", - "Checkbox-Selected", - "Checkbox-Unselected", ], "url": None, }, @@ -66,8 +64,6 @@ "Table", "Text", "Title", - "Checkbox-Selected", - "Checkbox-Unselected", ], "url": None, }, diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index ba3bf82feb..41dc228320 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -251,12 +251,18 @@ def _tables_from_regions( crop_h = int(round(max(np.linalg.norm(src[3] - src[0]), np.linalg.norm(src[2] - src[1])))) if crop_w < 2 or crop_h < 2: continue - # Full-extent destination corners so crop-relative [0, 1] maps exactly to [0, crop] pixels - dst = np.array([[0, 0], [crop_w, 0], [crop_w, crop_h], [0, crop_h]], dtype=np.float32) + # Map the region onto an inset rectangle so the warp samples `pad` extra pixels of page + # context on each side. Crop-relative [0, 1] still maps exactly to [0, out_w] / [0, out_h]. + pad = max(20, round(0.01 * max(crop_w, crop_h))) + out_w, out_h = crop_w + 2 * pad, crop_h + 2 * pad + dst = np.array( + [[pad, pad], [pad + crop_w, pad], [pad + crop_w, pad + crop_h], [pad, pad + crop_h]], + dtype=np.float32, + ) transform = cv2.getPerspectiveTransform(src, dst) inverse = cv2.getPerspectiveTransform(dst, src) - crops.append(cv2.warpPerspective(page, transform, (crop_w, crop_h))) - crop_meta.append((p_idx, inverse, crop_w, crop_h, w, h)) + crops.append(cv2.warpPerspective(page, transform, (out_w, out_h), borderMode=cv2.BORDER_REPLICATE)) + crop_meta.append((p_idx, inverse, out_w, out_h, w, h)) tables_per_page: list[list[dict[str, Any]]] = [[] for _ in pages] if len(crops) == 0: diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index 8ffa55e437..bafba48a27 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -155,9 +155,16 @@ def _to_canonical_ltr(boxes: np.ndarray, direction: str) -> np.ndarray: def _overlap_ratios(starts: np.ndarray, ends: np.ndarray) -> np.ndarray: """Compute the pairwise 1D overlap of intervals, normalized by the length of the smaller interval""" - inter = np.minimum(ends[:, None], ends[None, :]) - np.maximum(starts[:, None], starts[None, :]) - min_len = np.minimum(ends - starts, (ends - starts)[:, None]) - return np.clip(inter, 0, None) / np.clip(min_len, 1e-9, None) + starts = starts.astype(np.float32, copy=False) + ends = ends.astype(np.float32, copy=False) + lengths = ends - starts + inter = np.minimum(ends[:, None], ends[None, :]) + inter -= np.maximum(starts[:, None], starts[None, :]) + np.clip(inter, 0, None, out=inter) + min_len = np.minimum(lengths, lengths[:, None]) + np.clip(min_len, 1e-9, None, out=min_len) + inter /= min_len + return inter def _strict_rank(primary: np.ndarray, secondary: np.ndarray) -> np.ndarray: @@ -193,13 +200,17 @@ def _topological_order(boxes: np.ndarray, x_overlap_threshold: float, y_overlap_ # 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) - is_above = y_rank[:, None] < y_rank[None, :] - is_left = x_rank[:, None] < x_rank[None, :] - # i -> j edges: i must be read before j - edges = ((x_overlap > x_overlap_threshold) & is_above) | ( - (x_overlap <= x_overlap_threshold) & (y_overlap > y_overlap_threshold) & is_left - ) + # 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. + x_linked = x_overlap > x_overlap_threshold + edges = y_rank[:, None] < y_rank[None, :] # i is above j + edges &= x_linked + 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 + edges |= same_row + del same_row np.fill_diagonal(edges, False) in_degree = edges.sum(axis=0) @@ -220,13 +231,33 @@ def _find(node: int) -> int: node = int(parent[node]) return node - col_edges = (x_overlap > x_overlap_threshold) & ~spanning[:, None] & ~spanning[None, :] - for i, j in np.argwhere(np.triu(col_edges, 1)): + # Reuse the horizontal-overlap matrix; keep the upper triangle only, so each pair is visited once + col_edges = np.triu(x_linked, 1) + col_edges &= ~spanning[:, None] + col_edges &= ~spanning[None, :] + for i, j in np.argwhere(col_edges): ri, rj = _find(int(i)), _find(int(j)) if ri != rj: parent[ri] = rj component = np.array([_find(i) for i in range(num_boxes)]) + # 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 + if num_boxes >= 3: + span = page_width + tolerance = max(1, int(0.05 * num_boxes)) + centers = (x0 + x1) / 2 + lo, hi = x0.min() + 0.25 * span, x0.min() + 0.75 * span + for split in np.unique(x1[(x1 >= lo) & (x1 <= hi)]): + crossing = int(np.count_nonzero(np.minimum(x1 - split, split - x0) > 0.02 * span)) + left = int(np.count_nonzero(centers <= split)) + if crossing <= tolerance and left >= 0.25 * num_boxes and num_boxes - left >= 0.25 * num_boxes: + multi_column = True + break + 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) @@ -237,7 +268,7 @@ def _find(node: int) -> int: # last emitted one with a horizontal overlap candidates = ( ready[(x_overlap[last, ready] > x_overlap_threshold) & (y0[ready] >= y0[last])] - if last >= 0 + if last >= 0 and multi_column else np.empty(0, dtype=int) ) if candidates.size == 0 and last >= 0: diff --git a/tests/common/test_io_elements.py b/tests/common/test_io_elements.py index 2d8f00b7c3..122d7a4786 100644 --- a/tests/common/test_io_elements.py +++ b/tests/common/test_io_elements.py @@ -1,3 +1,4 @@ +import json from xml.etree.ElementTree import ElementTree import numpy as np @@ -453,10 +454,10 @@ def test_page(): assert page.language == language # Render - assert page.render() == "hello world\nhello world\n\nhello world\nhello world" + assert page.render() == "hello world\n\nhello world\n\nhello world\n\nhello world" - # Export - assert page.export() == { + # Export - `reading_order=False` serializes the blocks exactly as stored + assert page.export(reading_order=False) == { "blocks": [b.export() for b in blocks], "page_idx": page_idx, "dimensions": page_size, @@ -465,6 +466,7 @@ def test_page(): "layout": [r.export() for r in layout], "tables": [], } + assert json.dumps(page.export()) and json.dumps(page.export(reading_order=False)) # Export XML xml_bytes, xml_tree = page.export_as_xml() @@ -570,11 +572,12 @@ def test_document(): assert all(isinstance(p, elements.Page) for p in doc.pages) # Render - page_export = "hello world\nhello world\n\nhello world\nhello world" + page_export = "hello world\n\nhello world\n\nhello world\n\nhello world" assert doc.render() == f"{page_export}\n\n\n\n{page_export}" # Export assert doc.export() == {"pages": [p.export() for p in pages]} + assert json.dumps(doc.export()) # a full document export can be written to a JSON file as is # Export XML xml_output = doc.export_as_xml() @@ -645,3 +648,59 @@ def test_kie_document(): # Synthesize img_list = doc.synthesize() assert isinstance(img_list, list) and len(img_list) == len(pages) + + +def test_element_is_abstract(): + with pytest.raises(NotImplementedError): + elements.Element.from_dict({}) + with pytest.raises(NotImplementedError): + elements.Element().render() + with pytest.raises(KeyError): + elements.Element(unknown_child=[]) + + +def test_artefact_from_dict_round_trip(): + artefact = elements.Artefact("qr_code", 0.8, ((0.1, 0.1), (0.2, 0.2))) + assert elements.Artefact.from_dict(artefact.export()).export() == artefact.export() + # Artefacts survive a Block round trip + block = _mock_blocks()[0] + assert elements.Block.from_dict(block.export()).export() == block.export() + + +def test_page_from_dict_restores_the_image(): + page = _mock_pages()[0] + exported = page.export() + # The page image is not part of the export, so a placeholder is used unless it is passed back + restored = elements.Page.from_dict(exported) + assert isinstance(restored, elements.Page) and restored.page.shape == (0, 0, 3) + assert restored.export() == exported + restored = elements.Page.from_dict(exported, page=page.page) + assert np.array_equal(restored.page, page.page) + assert elements.Page.from_dict(json.loads(json.dumps(exported))).render() == restored.render() + + +def test_kie_page_and_kie_document_from_dict(): + kie_page = _mock_kie_pages()[0] + exported = kie_page.export() + restored = elements.KIEPage.from_dict(exported) + assert isinstance(restored, elements.KIEPage) and restored.page.shape == (0, 0, 3) + assert restored.export() == exported + assert np.array_equal(elements.KIEPage.from_dict(exported, page=kie_page.page).page, kie_page.page) + + # A KIEDocument must rebuild KIEPage objects, not plain Page ones + doc = elements.KIEDocument(_mock_kie_pages()) + restored_doc = elements.KIEDocument.from_dict(doc.export()) + assert isinstance(restored_doc, elements.KIEDocument) + assert all(isinstance(page, elements.KIEPage) for page in restored_doc.pages) + assert restored_doc.export() == doc.export() + assert restored_doc.render() == doc.render() + # The page images can be handed back positionally, page by page + images = [page.page for page in doc.pages] + restored_doc = elements.KIEDocument.from_dict(doc.export(), pages=images) + assert all(np.array_equal(page.page, image) for page, image in zip(restored_doc.pages, images)) + + # plain Document keeps rebuilding Page objects + plain = elements.Document(_mock_pages()) + restored_plain = elements.Document.from_dict(plain.export()) + assert all(isinstance(page, elements.Page) for page in restored_plain.pages) + assert restored_plain.export() == plain.export() diff --git a/tests/common/test_io_exporters.py b/tests/common/test_io_exporters.py new file mode 100644 index 0000000000..51f6ca8d7b --- /dev/null +++ b/tests/common/test_io_exporters.py @@ -0,0 +1,834 @@ +import json + +import numpy as np +import pytest + +from doctr.file_utils import CLASS_NAME +from doctr.io import elements +from doctr.io.exporters import ( + AsciiDocExporter, + HTMLExporter, + MarkdownExporter, + TextExporter, + XMLExporter, + ordered_line_words, + page_reading_order, + to_json_safe, +) + + +def _word_at(text, x0, y0, x1, y1): + return elements.Word(text, 0.95, ((x0, y0), (x1, y1)), 0.9, {"value": 0, "confidence": None}) + + +def _line_at(text, x0, y0, x1, y1, rtl=False): + """Build a line whose words are laid out geometrically (leftmost word = last logical word when rtl)""" + words = text.split() + step = (x1 - x0) / max(len(words), 1) + geo_words = words[::-1] if rtl else words + return elements.Line([ + _word_at(word, x0 + idx * step, y0, x0 + (idx + 0.9) * step, y1) for idx, word in enumerate(geo_words) + ]) + + +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)] + lines += [_line_at(f"left line {idx}", 0.08, 0.14 + 0.05 * idx, 0.46, 0.17 + 0.05 * idx) for idx in range(3)] + lines += [_line_at(f"right line {idx}", 0.54, 0.14 + 0.05 * idx, 0.92, 0.17 + 0.05 * idx) for idx in range(3)] + lines += [_line_at("- item one", 0.08, 0.4, 0.46, 0.43), _line_at("Page 3 of 12", 0.4, 0.95, 0.6, 0.97)] + # Shuffle the lines to make sure the export does not rely on the input order + lines = [lines[idx] for idx in [5, 0, 8, 2, 4, 7, 1, 6, 3]] + layout = [ + elements.LayoutElement("Title", 0.99, ((0.15, 0.04), (0.85, 0.1))), + elements.LayoutElement("Text", 0.98, ((0.06, 0.12), (0.48, 0.32))), + elements.LayoutElement("Text", 0.98, ((0.52, 0.12), (0.94, 0.32))), + 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 + ) + + +def test_page_items_in_reading_order(): + page = _reading_order_page() + items = page.items_in_reading_order() + assert all(isinstance(item, elements.Block) for item in items) + rendered = [item.render(line_break=" ") for item in items] + 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" + ) + # 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"] + + +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 + 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)])], + 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() == "" + + +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("الأيسر") + + +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_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_page_export_as(): + 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") + + +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) + ] + 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") + + +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)) + 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() + + +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) + + +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_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) + + 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), + ] + 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) + + 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), + ] + 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" + + +def test_page_export_as_html(): + page = _reading_order_page() + html = page.export_as_html() + # title heading, paragraphs in reading order, escaping + 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
    " + + +def test_export_mixins_carry_full_api(): + # 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", + "export_as_markdown", + "export_as_asciidoc", + "export_as_html", + "export_as", + ): + assert getattr(elements.Page, method) is getattr(PageExportsMixin, method) + 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" + + +def test_xml_exporter_class(): + from xml.etree import ElementTree as ET + + 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 + + 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] + + 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] + + +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 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(): + page = _table_page() + html = page.export_as_html() + assert "" in html and "" in html and "" in html + assert html.index("

    intro line

    ") < html.index("
    headbody
    ") + # A single-row table has a header but no body + single = elements.Table( + cells=[elements.TableCell("only", 0.9, ((0.1, 0.6), (0.4, 0.65)), 0, 0, 0, 0)], + num_rows=1, + num_cols=1, + 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(): + 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], + ) + 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 // ") + + +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_edge_cases(): + from doctr.io.exporters import _covering_region_indices + + 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() + 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)) + + +def test_every_format_is_reachable_from_a_document(): + 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), + ) + 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") + + 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().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() diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index 2c1084b8a1..c96c4487cb 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -1,3 +1,5 @@ +import json + import numpy as np import pytest @@ -96,7 +98,7 @@ def test_documentbuilder(): # Repr assert ( repr(doc_builder) == "DocumentBuilder(resolve_lines=True, " - "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False)" + "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False, keep_reading_order=False)" ) @@ -185,7 +187,7 @@ def test_kiedocumentbuilder(): # Repr assert ( repr(doc_builder) == "KIEDocumentBuilder(resolve_lines=True, " - "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False)" + "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False, keep_reading_order=False)" ) @@ -581,3 +583,393 @@ def test_documentbuilder_tables_empty_cells(): 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"] + + +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"] + # 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 + 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"] + + # 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)] + 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"] + + +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)], + ) + 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")): + for row in range(3): + 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 = { + "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(): + 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"] + + +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(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)] + ] + + 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 + ) + 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"] + 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 d4c6ffc2dd..a5ca20620d 100644 --- a/tests/common/test_models_reading_order.py +++ b/tests/common/test_models_reading_order.py @@ -273,16 +273,20 @@ def test_deskew_strong_rotation_non_square_page(): # Auto generated regression tests for known failures of the reading order algorithm +def _box(x0, y0, x1, y1): + return ((x0, y0), (x1, y1)) + + def test_sort_reading_order_fragmented_columns(): left = [ - ((0.10, 0.10), (0.45, 0.13)), # 0 wide - ((0.10, 0.14), (0.25, 0.17)), # 1 narrow (left part of a split line) - ((0.34, 0.14), (0.45, 0.17)), # 2 stray fragment (right part), same visual row as 1 - ((0.10, 0.18), (0.45, 0.21)), # 3 - ((0.10, 0.22), (0.45, 0.25)), # 4 - ((0.10, 0.26), (0.45, 0.29)), # 5 + _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 = [((0.55, 0.10 + 0.04 * i), (0.90, 0.13 + 0.04 * i)) for i in range(6)] # 6..11 + 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)) @@ -290,13 +294,43 @@ def test_sort_reading_order_fragmented_columns(): def test_fragmented_row_with_merged_column_components(): geoms = [ - ((0.35, 0.05), (0.65, 0.10)), # 0 gutter-straddling element (bridges both columns) - ((0.10, 0.15), (0.45, 0.20)), # 1 left col, row 1 - ((0.10, 0.22), (0.16, 0.27)), # 2 left col, row 2, fragment A - ((0.17, 0.22), (0.24, 0.27)), # 3 left col, row 2, fragment B - ((0.25, 0.22), (0.45, 0.27)), # 4 left col, row 2, fragment C - ((0.10, 0.29), (0.45, 0.34)), # 5 left col, row 3 - ((0.55, 0.15), (0.90, 0.20)), # 6 right col, row 1 - ((0.55, 0.22), (0.90, 0.27)), # 7 right col, row 2 + _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] + + +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]