diff --git a/docs/source/modules/io.rst b/docs/source/modules/io.rst index 56d88e014d..104c969c42 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,30 @@ 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 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`. + +.. currentmodule:: doctr.io + +.. 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 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..caa32a62ec 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 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 @@ -39,37 +36,6 @@ ] -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 Element(NestedObject): """Implements an abstract document element with exporting and text rendering capabilities""" @@ -451,7 +417,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 +459,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,108 +496,6 @@ 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} @@ -647,7 +507,7 @@ def from_dict(cls, save_dict: dict[str, Any], **kwargs): return cls(**kwargs) -class KIEPage(Element): +class KIEPage(KIEPageExportsMixin, Element): """Implements a KIE page element as a collection of predictions Args: @@ -682,12 +542,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,96 +578,6 @@ 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} @@ -827,7 +591,7 @@ def from_dict(cls, save_dict: dict[str, Any], **kwargs): return cls(**kwargs) -class Document(Element): +class Document(DocumentExportsMixin, Element): """Implements a document element as a collection of pages Args: @@ -843,10 +607,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,17 +623,6 @@ 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] - @classmethod def from_dict(cls, save_dict: dict[str, Any], **kwargs): kwargs = {k: save_dict[k] for k in cls._exported_keys} diff --git a/doctr/io/exporters.py b/doctr/io/exporters.py new file mode 100644 index 0000000000..0dd66ec584 --- /dev/null +++ b/doctr/io/exporters.py @@ -0,0 +1,1029 @@ +# 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", + "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) + + +_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 - e.g. a __slots__ subclass + 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 = [] + # 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: + items[-1] = Block(lines=[*items[-1].lines, elements[idx]]) + else: + items.append(Block(lines=[elements[idx]])) + labels.append(seg_label) + open_list_region = region + else: + items.append(Block(lines=[elements[idx] for idx in segment])) + labels.append(seg_label) + open_list_region = None + _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]) + + +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] = "- " + 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.""" + words = line.words + if direction in ("ttb-rtl", "ttb-ltr"): + words = sorted(words, key=lambda word: float(np.asarray(word.geometry, dtype=np.float64)[..., 1].mean())) + elif direction == "rtl": + words = sorted(words, key=lambda word: -float(np.asarray(word.geometry, dtype=np.float64)[..., 0].mean())) + text = " ".join(word.render() for word in words) + 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 + ) -> 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 + + 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 "\n\n".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 + """ + from doctr.models.reading_order import ReadingOrderPredictor + + language = page.language.get("value") if isinstance(page.language, dict) else None + predictor = ReadingOrderPredictor(direction=direction) + parts: list[str] = [] + for class_name, predictions in page.predictions.items(): + if len(predictions) == 0: + continue + order = predictor( + [prediction.geometry for prediction in predictions], + texts=[prediction.value for prediction in predictions], + language=language, + page_shape=page.dimensions, + ) + values = "\n".join( + self.bullet + + (self.finalize_line(self.escape_text(predictions[idx].value)) if escape else predictions[idx].value) + for idx in order + ) + 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 MarkdownExporter(_PageTextExporter): + """Export OCR results to Markdown, with the content sorted in reading order. + + >>> from doctr.io import MarkdownExporter + >>> markdown = MarkdownExporter().export_page(page) # doctest: +SKIP + """ + + 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) # doctest: +SKIP + """ + + 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"} + 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 + ) -> 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 "\n".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: + from doctr.models.reading_order import ReadingOrderPredictor + + language = page.language.get("value") if isinstance(page.language, dict) else None + predictor = ReadingOrderPredictor(direction=direction) + parts: list[str] = [] + for class_name, predictions in page.predictions.items(): + if len(predictions) == 0: + continue + order = predictor( + [prediction.geometry for prediction in predictions], + texts=[prediction.value for prediction in predictions], + language=language, + page_shape=page.dimensions, + ) + values = "\n".join( + f"
  • {self.escape_text(predictions[idx].value) if escape else predictions[idx].value}
  • " + for idx in order + ) + 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 export_page(self, page: "Page", file_title: str = "docTR - XML export (hOCR)") -> tuple[bytes, ET.ElementTree]: + """Export a page as hOCR XML. + + Args: + page: the page to export + file_title: the title of the XML file + + Returns: + a tuple of the XML byte string, and its ElementTree + """ + block_count: int = 1 + line_count: int = 1 + word_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", + }, + ) + # iterate over the blocks / lines / words and create the XML elements line by line with the attributes + for block in page.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: + 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(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)" + ) -> tuple[bytes, ET.ElementTree]: + """Export a KIE page as hOCR XML. + + Args: + page: the KIE page to export + file_title: the title of the XML file + + 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(): + 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) + + 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 export(self) -> dict[str, Any]: ... + + def render(self, block_break: str = "\n\n") -> str: + """Renders the full text of the page. + + Args: + block_break: the string inserted between two blocks + + Returns: + the text of the page + """ + return block_break.join(block.render() for block in self.blocks) + + 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 + """ + return XMLExporter().export_page(cast("Page", self), file_title=file_title) + + 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) -> dict[str, Any]: ... + + def render(self, prediction_break: str = "\n\n") -> 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 + + Returns: + the text of the page, one section per detection class with its predictions in reading order + """ + from doctr.models.reading_order import ReadingOrderPredictor + + predictor = ReadingOrderPredictor() + language = self.language.get("value") if isinstance(self.language, dict) else None + parts: list[str] = [] + for class_name, predictions in self.predictions.items(): + ordered = predictions + if len(ordered) > 1: + order = predictor( + [prediction.geometry for prediction in ordered], + texts=[prediction.value for prediction in ordered], + language=language, + page_shape=self.dimensions, + ) + ordered = [ordered[idx] for idx in order] + parts.extend(f"{class_name}: {prediction.render()}" for prediction in ordered) + return prediction_break.join(parts) + + 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 + """ + return XMLExporter().export_kie_page(cast("KIEPage", self), file_title=file_title) + + 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] + + def export(self) -> dict[str, Any]: ... + + 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 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_exporters.py b/tests/common/test_io_exporters.py new file mode 100644 index 0000000000..c97bcd2d8a --- /dev/null +++ b/tests/common/test_io_exporters.py @@ -0,0 +1,426 @@ +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, XMLExporter, page_reading_order + + +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_preserves_block_order(): + left = elements.Block( + lines=[_line_at("left top", 0.08, 0.1, 0.45, 0.13), _line_at("left low", 0.08, 0.2, 0.45, 0.23)] + ) + right = elements.Block( + lines=[_line_at("right top", 0.55, 0.1, 0.92, 0.13), _line_at("right low", 0.55, 0.2, 0.92, 0.23)] + ) + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [right, left], 0, (1000, 800)) + assert page.render() == "right top\nright low\n\nleft top\nleft low" + # Single-block and empty pages + assert elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [left], 0, (1000, 800)).render() == "left top\nleft low" + assert elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [], 0, (1000, 800)).render() == "" + + +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] diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index 2c1084b8a1..14ca1d6dc1 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -96,7 +96,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 +185,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 +581,221 @@ 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) + assert doc.pages[0].render(block_break=" ").split() == ["L0", "L1", "L2", "R0", "R1", "R2"] + # Without the flag, the blocks keep their original (interleaved) order + doc = builder.DocumentBuilder(resolve_blocks=True, keep_reading_order=False)(*args) + assert doc.pages[0].render(block_break=" ").split() != ["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 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]