# Copyright (C) 2021-2026, Mindee.
# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
from typing import Any
import numpy as np
from doctr.file_utils import requires_package
from doctr.io.exporters import DocumentExportsMixin, KIEPageExportsMixin, PageExportsMixin, to_json_safe
from doctr.utils.common_types import BoundingBox
from doctr.utils.geometry import resolve_enclosing_bbox, resolve_enclosing_rbbox
from doctr.utils.reconstitution import synthesize_kie_page, synthesize_page
from doctr.utils.repr import NestedObject
try: # optional dependency for visualization
from doctr.utils.visualization import visualize_kie_page, visualize_page
except ModuleNotFoundError: # pragma: no cover
pass
__all__ = [
"Element",
"Word",
"Artefact",
"Line",
"Prediction",
"Block",
"Page",
"KIEPage",
"KIEDocument",
"Document",
"LayoutElement",
"TableCell",
"Table",
]
def _empty_page_image(page: np.ndarray | None) -> np.ndarray:
"""Return the given page image, or an empty placeholder when the page was restored from an export."""
return page if page is not None else np.zeros((0, 0, 3), dtype=np.uint8)
class Element(NestedObject):
"""Implements an abstract document element with exporting and text rendering capabilities"""
_children_names: list[str] = []
_exported_keys: list[str] = []
def __init__(self, **kwargs: Any) -> None:
for k, v in kwargs.items():
if k in self._children_names:
setattr(self, k, v)
else:
raise KeyError(f"{self.__class__.__name__} object does not have any attribute named '{k}'")
def export(self) -> dict[str, Any]:
"""Exports the object into a nested dict format"""
export_dict = {k: to_json_safe(getattr(self, k)) for k in self._exported_keys}
for children_name in self._children_names:
if children_name in ["predictions"]:
export_dict[children_name] = {
k: [item.export() for item in c] for k, c in getattr(self, children_name).items()
}
else:
export_dict[children_name] = [c.export() for c in getattr(self, children_name)]
return export_dict
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
raise NotImplementedError
def render(self) -> str:
raise NotImplementedError
[docs]
class Word(Element):
"""Implements a word element
Args:
value: the text string of the word
confidence: the confidence associated with the text prediction
geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to
the page's size
objectness_score: the objectness score of the detection
crop_orientation: the general orientation of the crop in degrees and its confidence
"""
_exported_keys: list[str] = ["value", "confidence", "geometry", "objectness_score", "crop_orientation"]
_children_names: list[str] = []
def __init__(
self,
value: str,
confidence: float,
geometry: BoundingBox | np.ndarray,
objectness_score: float,
crop_orientation: dict[str, Any],
) -> None:
super().__init__()
self.value = value
self.confidence = confidence
self.geometry = geometry
self.objectness_score = objectness_score
self.crop_orientation = crop_orientation
def render(self) -> str:
"""Renders the full text of the element"""
return self.value
def extra_repr(self) -> str:
return f"value='{self.value}', confidence={self.confidence:.2}"
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
kwargs = {k: save_dict[k] for k in cls._exported_keys}
return cls(**kwargs)
[docs]
class Artefact(Element):
"""Implements a non-textual element
Args:
artefact_type: the type of artefact
confidence: the confidence of the type prediction
geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to
the page's size.
"""
_exported_keys: list[str] = ["geometry", "type", "confidence"]
_children_names: list[str] = []
def __init__(self, artefact_type: str, confidence: float, geometry: BoundingBox) -> None:
super().__init__()
self.geometry = geometry
self.type = artefact_type
self.confidence = confidence
def render(self) -> str:
"""Renders the region as a tag"""
return f"<[{self.type.upper()}]>"
def extra_repr(self) -> str:
return f"type='{self.type}', confidence={self.confidence:.2}"
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
return cls(artefact_type=save_dict["type"], confidence=save_dict["confidence"], geometry=save_dict["geometry"])
[docs]
class LayoutElement(Element):
"""Implements a layout region predicted by a layout detection model
Args:
layout_type: the predicted region class (e.g. 'Title', 'Text', 'Table', 'Page-header')
confidence: the confidence of the region prediction
geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to
the page's size
"""
_exported_keys: list[str] = ["geometry", "type", "confidence"]
_children_names: list[str] = []
def __init__(self, layout_type: str, confidence: float, geometry: BoundingBox | np.ndarray) -> None:
super().__init__()
self.geometry = geometry
self.type = layout_type
self.confidence = confidence
def render(self) -> str:
"""Renders the region as a tag"""
return f"<[{self.type.upper()}]>"
def extra_repr(self) -> str:
return f"type='{self.type}', confidence={self.confidence:.2}"
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
kwargs = {k: save_dict[k] for k in cls._exported_keys}
return cls(layout_type=kwargs["type"], confidence=kwargs["confidence"], geometry=kwargs["geometry"])
class TableCell(Element):
"""Implements a single cell of a recognized table
Args:
value: the text content of the cell (words assigned to the cell, joined together)
confidence: the mean recognition confidence of the words assigned to the cell
geometry: bounding box of the cell in format ((xmin, ymin), (xmax, ymax)) or a (4, 2) polygon,
with coordinates relative to the page's size
row_start: index of the first row spanned by the cell (0-indexed)
row_end: index of the last row spanned by the cell (0-indexed, inclusive)
col_start: index of the first column spanned by the cell (0-indexed)
col_end: index of the last column spanned by the cell (0-indexed, inclusive)
"""
_exported_keys: list[str] = [
"geometry",
"value",
"confidence",
"row_start",
"row_end",
"col_start",
"col_end",
]
_children_names: list[str] = []
def __init__(
self,
value: str,
confidence: float,
geometry: BoundingBox | np.ndarray,
row_start: int,
row_end: int,
col_start: int,
col_end: int,
) -> None:
super().__init__()
self.value = value
self.confidence = confidence
self.geometry = geometry
self.row_start = row_start
self.row_end = row_end
self.col_start = col_start
self.col_end = col_end
@property
def row_span(self) -> int:
"""Number of rows spanned by the cell"""
return self.row_end - self.row_start + 1
@property
def col_span(self) -> int:
"""Number of columns spanned by the cell"""
return self.col_end - self.col_start + 1
def render(self) -> str:
"""Renders the cell text"""
return self.value
def extra_repr(self) -> str:
return f"value='{self.value}', rows=({self.row_start}, {self.row_end}), cols=({self.col_start}, {self.col_end})"
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
kwargs = {k: save_dict[k] for k in cls._exported_keys}
return cls(**kwargs)
class Table(Element):
"""Implements a table recognized on a page as a grid of cells
The recognized text of the words falling inside the table is regrouped here and removed from the
regular `blocks` output of the page, so it is not duplicated. The structured content can be loaded
directly into pandas, e.g. `pd.DataFrame(table.to_grid())`.
Args:
cells: list of `TableCell` objects composing the table
num_rows: number of rows of the table
num_cols: number of columns of the table
geometry: bounding box enclosing the whole table, with coordinates relative to the page's size
confidence: the confidence of the table structure prediction
"""
_exported_keys: list[str] = ["geometry", "num_rows", "num_cols", "confidence"]
_children_names: list[str] = ["cells"]
cells: list[TableCell] = []
def __init__(
self,
cells: list[TableCell],
num_rows: int,
num_cols: int,
geometry: BoundingBox | np.ndarray,
confidence: float = 1.0,
) -> None:
super().__init__(cells=cells)
self.num_rows = num_rows
self.num_cols = num_cols
self.geometry = geometry
self.confidence = confidence
def to_grid(self) -> list[list[str]]:
"""Return the table content as a dense `num_rows` x `num_cols` grid of strings.
Cells spanning several rows/columns have their value placed at their top-left position; the
remaining positions they span are left empty. The result is directly loadable into pandas via
`pd.DataFrame(table.to_grid())`.
Returns:
a list of `num_rows` lists, each of length `num_cols`
"""
grid = [["" for _ in range(self.num_cols)] for _ in range(self.num_rows)]
for cell in self.cells:
if 0 <= cell.row_start < self.num_rows and 0 <= cell.col_start < self.num_cols:
grid[cell.row_start][cell.col_start] = cell.value
return grid
def render(self, row_break: str = "\n", col_break: str = "\t") -> str:
"""Renders the table as plain text (tab-separated values)"""
return row_break.join(col_break.join(row) for row in self.to_grid())
def extra_repr(self) -> str:
return f"num_rows={self.num_rows}, num_cols={self.num_cols}, confidence={self.confidence:.2}"
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
kwargs = {k: save_dict[k] for k in cls._exported_keys}
kwargs["cells"] = [TableCell.from_dict(cell) for cell in save_dict["cells"]]
return cls(**kwargs)
[docs]
class Line(Element):
"""Implements a line element as a collection of words
Args:
words: list of word elements
geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to
the page's size. If not specified, it will be resolved by default to the smallest bounding box enclosing
all words in it.
"""
_exported_keys: list[str] = ["geometry", "objectness_score"]
_children_names: list[str] = ["words"]
words: list[Word] = []
def __init__(
self,
words: list[Word],
geometry: BoundingBox | np.ndarray | None = None,
objectness_score: float | None = None,
) -> None:
# Compute the objectness score of the line
if objectness_score is None:
objectness_score = float(np.mean([w.objectness_score for w in words]))
# Resolve the geometry using the smallest enclosing bounding box
if geometry is None:
# Check whether this is a rotated or straight box
box_resolution_fn = resolve_enclosing_rbbox if len(words[0].geometry) == 4 else resolve_enclosing_bbox
geometry = box_resolution_fn([w.geometry for w in words]) # type: ignore[misc]
super().__init__(words=words)
self.geometry = geometry
self.objectness_score = objectness_score
def render(self) -> str:
"""Renders the full text of the element"""
return " ".join(w.render() for w in self.words)
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
kwargs = {k: save_dict[k] for k in cls._exported_keys}
kwargs.update({
"words": [Word.from_dict(_dict) for _dict in save_dict["words"]],
})
return cls(**kwargs)
[docs]
class Prediction(Word):
"""Implements a prediction element"""
def render(self) -> str:
"""Renders the full text of the element"""
return self.value
def extra_repr(self) -> str:
return f"value='{self.value}', confidence={self.confidence:.2}, bounding_box={self.geometry}"
[docs]
class Block(Element):
"""Implements a block element as a collection of lines and artefacts
Args:
lines: list of line elements
artefacts: list of artefacts
geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to
the page's size. If not specified, it will be resolved by default to the smallest bounding box enclosing
all lines and artefacts in it.
"""
_exported_keys: list[str] = ["geometry", "objectness_score"]
_children_names: list[str] = ["lines", "artefacts"]
lines: list[Line] = []
artefacts: list[Artefact] = []
def __init__(
self,
lines: list[Line] = [],
artefacts: list[Artefact] = [],
geometry: BoundingBox | np.ndarray | None = None,
objectness_score: float | None = None,
) -> None:
# Compute the objectness score of the line
if objectness_score is None:
objectness_score = float(np.mean([w.objectness_score for line in lines for w in line.words]))
# Resolve the geometry using the smallest enclosing bounding box
if geometry is None:
line_boxes = [word.geometry for line in lines for word in line.words]
artefact_boxes = [artefact.geometry for artefact in artefacts]
box_resolution_fn = (
resolve_enclosing_rbbox if isinstance(lines[0].geometry, np.ndarray) else resolve_enclosing_bbox
)
geometry = box_resolution_fn(line_boxes + artefact_boxes) # type: ignore
super().__init__(lines=lines, artefacts=artefacts)
self.geometry = geometry
self.objectness_score = objectness_score
def render(self, line_break: str = "\n") -> str:
"""Renders the full text of the element"""
return line_break.join(line.render() for line in self.lines)
@classmethod
def from_dict(cls, save_dict: dict[str, Any], **kwargs):
kwargs = {k: save_dict[k] for k in cls._exported_keys}
kwargs.update({
"lines": [Line.from_dict(_dict) for _dict in save_dict["lines"]],
"artefacts": [Artefact.from_dict(_dict) for _dict in save_dict["artefacts"]],
})
return cls(**kwargs)
[docs]
class Page(PageExportsMixin, Element):
"""Implements a page element as a collection of blocks
Args:
page: image encoded as a numpy array in uint8
blocks: list of block elements
page_idx: the index of the page in the input raw document
dimensions: the page size in pixels in format (height, width)
orientation: a dictionary with the value of the rotation angle in degress and confidence of the prediction
language: a dictionary with the language value and confidence of the prediction
layout: optional list of layout regions detected on the page
tables: optional list of tables recognized on the page. Words assigned to a table are removed from `blocks`.
"""
_exported_keys: list[str] = ["page_idx", "dimensions", "orientation", "language"]
_children_names: list[str] = ["blocks", "layout", "tables"]
blocks: list[Block] = []
layout: list[LayoutElement] = []
tables: list[Table] = []
def __init__(
self,
page: np.ndarray,
blocks: list[Block],
page_idx: int,
dimensions: tuple[int, int],
orientation: dict[str, Any] | None = None,
language: dict[str, Any] | None = None,
layout: list[LayoutElement] | None = None,
tables: list[Table] | None = None,
) -> None:
super().__init__(
blocks=blocks,
layout=layout if layout is not None else [],
tables=tables if tables is not None else [],
)
self.page = page
self.page_idx = page_idx
self.dimensions = dimensions
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 extra_repr(self) -> str:
return f"dimensions={self.dimensions}"
[docs]
def show(self, interactive: bool = True, preserve_aspect_ratio: bool = False, **kwargs) -> None:
"""Overlay the result on a given image
Args:
interactive: whether the display should be interactive
preserve_aspect_ratio: pass True if you passed True to the predictor
**kwargs: additional keyword arguments passed to the matplotlib.pyplot.show method
(e.g. `display_layout=False` to hide detected layout regions)
"""
requires_package("matplotlib", "`.show()` requires matplotlib & mplcursors installed")
requires_package("mplcursors", "`.show()` requires matplotlib & mplcursors installed")
import matplotlib.pyplot as plt
show_kwargs = {k: kwargs.pop(k) for k in ("words_only", "display_artefacts", "display_layout") if k in kwargs}
visualize_page(
self.export(),
self.page,
interactive=interactive,
preserve_aspect_ratio=preserve_aspect_ratio,
**show_kwargs,
)
plt.show(**kwargs)
def synthesize(self, **kwargs) -> np.ndarray:
"""Synthesize the page from the predictions
Args:
**kwargs: keyword arguments passed to the `synthesize_page` method
Returns:
synthesized page
"""
return synthesize_page(self.export(), **kwargs)
@classmethod
def from_dict(cls, save_dict: dict[str, Any], page: np.ndarray | None = None, **kwargs):
_kwargs: dict[str, Any] = {k: save_dict[k] for k in cls._exported_keys}
_kwargs.update({
"blocks": [Block.from_dict(block_dict) for block_dict in save_dict["blocks"]],
"layout": [LayoutElement.from_dict(region_dict) for region_dict in save_dict.get("layout", [])],
"tables": [Table.from_dict(table_dict) for table_dict in save_dict.get("tables", [])],
})
# The page image is not part of the export: pass it back explicitly to restore a fully usable page
return cls(page=_empty_page_image(page), **_kwargs)
[docs]
class KIEPage(KIEPageExportsMixin, Element):
"""Implements a KIE page element as a collection of predictions
Args:
predictions: Dictionary with list of block elements for each detection class
page: image encoded as a numpy array in uint8
page_idx: the index of the page in the input raw document
dimensions: the page size in pixels in format (height, width)
orientation: a dictionary with the value of the rotation angle in degress and confidence of the prediction
language: a dictionary with the language value and confidence of the prediction
layout: optional list of layout regions detected on the page
"""
_exported_keys: list[str] = ["page_idx", "dimensions", "orientation", "language"]
_children_names: list[str] = ["predictions", "layout"]
predictions: dict[str, list[Prediction]] = {}
layout: list[LayoutElement] = []
def __init__(
self,
page: np.ndarray,
predictions: dict[str, list[Prediction]],
page_idx: int,
dimensions: tuple[int, int],
orientation: dict[str, Any] | None = None,
language: dict[str, Any] | None = None,
layout: list[LayoutElement] | None = None,
) -> None:
super().__init__(predictions=predictions, layout=layout if layout is not None else [])
self.page = page
self.page_idx = page_idx
self.dimensions = dimensions
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 extra_repr(self) -> str:
return f"dimensions={self.dimensions}"
[docs]
def show(self, interactive: bool = True, preserve_aspect_ratio: bool = False, **kwargs) -> None:
"""Overlay the result on a given image
Args:
interactive: whether the display should be interactive
preserve_aspect_ratio: pass True if you passed True to the predictor
**kwargs: keyword arguments passed to the matplotlib.pyplot.show method
"""
requires_package("matplotlib", "`.show()` requires matplotlib & mplcursors installed")
requires_package("mplcursors", "`.show()` requires matplotlib & mplcursors installed")
import matplotlib.pyplot as plt
show_kwargs = {k: kwargs.pop(k) for k in ("words_only", "display_artefacts", "display_layout") if k in kwargs}
visualize_kie_page(
self.export(),
self.page,
interactive=interactive,
preserve_aspect_ratio=preserve_aspect_ratio,
**show_kwargs,
)
plt.show(**kwargs)
def synthesize(self, **kwargs) -> np.ndarray:
"""Synthesize the page from the predictions
Args:
**kwargs: keyword arguments passed to the `synthesize_kie_page` method
Returns:
synthesized page
"""
return synthesize_kie_page(self.export(), **kwargs)
@classmethod
def from_dict(cls, save_dict: dict[str, Any], page: np.ndarray | None = None, **kwargs):
_kwargs: dict[str, Any] = {k: save_dict[k] for k in cls._exported_keys}
_kwargs.update({
"predictions": {
class_name: [Prediction.from_dict(pred) for pred in preds]
for class_name, preds in save_dict["predictions"].items()
},
"layout": [LayoutElement.from_dict(region_dict) for region_dict in save_dict.get("layout", [])],
})
# The page image is not part of the export: pass it back explicitly to restore a fully usable page
return cls(page=_empty_page_image(page), **_kwargs)
[docs]
class Document(DocumentExportsMixin, Element):
"""Implements a document element as a collection of pages
Args:
pages: list of page elements
"""
_children_names: list[str] = ["pages"]
pages: list[Page] = []
def __init__(
self,
pages: list[Page],
) -> None:
super().__init__(pages=pages)
[docs]
def show(self, **kwargs) -> None:
"""Overlay the result on a given image"""
for result in self.pages:
result.show(**kwargs)
def synthesize(self, **kwargs) -> list[np.ndarray]:
"""Synthesize all pages from their predictions
Args:
**kwargs: keyword arguments passed to the `Page.synthesize` method
Returns:
list of synthesized pages
"""
return [page.synthesize(**kwargs) for page in self.pages]
_page_cls: Any = Page
@classmethod
def from_dict(cls, save_dict: dict[str, Any], pages: list[np.ndarray] | None = None, **kwargs):
_kwargs: dict[str, Any] = {k: save_dict[k] for k in cls._exported_keys}
_kwargs.update({
"pages": [
cls._page_cls.from_dict(page_dict, page=None if pages is None else pages[idx])
for idx, page_dict in enumerate(save_dict["pages"])
]
})
return cls(**_kwargs)
[docs]
class KIEDocument(Document):
"""Implements a document element as a collection of pages
Args:
pages: list of page elements
"""
_children_names: list[str] = ["pages"]
_page_cls: Any = KIEPage
pages: list[KIEPage] = [] # type: ignore[assignment]
def __init__(
self,
pages: list[KIEPage],
) -> None:
super().__init__(pages=pages) # type: ignore[arg-type]