onnxtr 0.5.1__py3-none-any.whl → 0.6.0__py3-none-any.whl

onnxtr/contrib/__init__.py

@@ -0,0 +1 @@
+from .artefacts import ArtefactDetector

onnxtr/contrib/artefacts.py

@@ -3,7 +3,7 @@
 # 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, Dict, List, Optional, Tuple
+from typing import Any
 
 import cv2
 import numpy as np
@@ -14,7 +14,7 @@ from .base import _BasePredictor
 
 __all__ = ["ArtefactDetector"]
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "yolov8_artefact": {
         "input_shape": (3, 1024, 1024),
         "labels": ["bar_code", "qr_code", "logo", "photo"],
@@ -34,7 +34,6 @@ class ArtefactDetector(_BasePredictor):
     >>> results = detector(doc)
 
     Args:
-    ----
         arch: the architecture to use
         batch_size: the batch size to use
         model_path: the path to the model to use
@@ -50,9 +49,9 @@ class ArtefactDetector(_BasePredictor):
         self,
         arch: str = "yolov8_artefact",
         batch_size: int = 2,
-        model_path: Optional[str] = None,
-        labels: Optional[List[str]] = None,
-        input_shape: Optional[Tuple[int, int, int]] = None,
+        model_path: str | None = None,
+        labels: list[str] | None = None,
+        input_shape: tuple[int, int, int] | None = None,
         conf_threshold: float = 0.5,
         iou_threshold: float = 0.5,
         **kwargs: Any,
@@ -66,7 +65,7 @@ class ArtefactDetector(_BasePredictor):
     def preprocess(self, img: np.ndarray) -> np.ndarray:
         return np.transpose(cv2.resize(img, (self.input_shape[2], self.input_shape[1])), (2, 0, 1)) / np.array(255.0)
 
-    def postprocess(self, output: List[np.ndarray], input_images: List[List[np.ndarray]]) -> List[List[Dict[str, Any]]]:
+    def postprocess(self, output: list[np.ndarray], input_images: list[list[np.ndarray]]) -> list[list[dict[str, Any]]]:
         results = []
 
         for batch in zip(output, input_images):
@@ -109,7 +108,6 @@ class ArtefactDetector(_BasePredictor):
         Display the results
 
         Args:
-        ----
             **kwargs: additional keyword arguments to be passed to `plt.show`
         """
         requires_package("matplotlib", "`.show()` requires matplotlib installed")

onnxtr/contrib/base.py

@@ -3,7 +3,7 @@
 # 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, List, Optional
+from typing import Any
 
 import numpy as np
 import onnxruntime as ort
@@ -16,32 +16,29 @@ class _BasePredictor:
     Base class for all predictors
 
     Args:
-    ----
         batch_size: the batch size to use
         url: the url to use to download a model if needed
         model_path: the path to the model to use
         **kwargs: additional arguments to be passed to `download_from_url`
     """
 
-    def __init__(self, batch_size: int, url: Optional[str] = None, model_path: Optional[str] = None, **kwargs) -> None:
+    def __init__(self, batch_size: int, url: str | None = None, model_path: str | None = None, **kwargs) -> None:
         self.batch_size = batch_size
         self.session = self._init_model(url, model_path, **kwargs)
 
-        self._inputs: List[np.ndarray] = []
-        self._results: List[Any] = []
+        self._inputs: list[np.ndarray] = []
+        self._results: list[Any] = []
 
-    def _init_model(self, url: Optional[str] = None, model_path: Optional[str] = None, **kwargs: Any) -> Any:
+    def _init_model(self, url: str | None = None, model_path: str | None = None, **kwargs: Any) -> Any:
         """
         Download the model from the given url if needed
 
         Args:
-        ----
             url: the url to use
             model_path: the path to the model to use
             **kwargs: additional arguments to be passed to `download_from_url`
 
         Returns:
-        -------
             Any: the ONNX loaded model
         """
         if not url and not model_path:
@@ -54,40 +51,34 @@ class _BasePredictor:
         Preprocess the input image
 
         Args:
-        ----
             img: the input image to preprocess
 
         Returns:
-        -------
             np.ndarray: the preprocessed image
         """
         raise NotImplementedError
 
-    def postprocess(self, output: List[np.ndarray], input_images: List[List[np.ndarray]]) -> Any:
+    def postprocess(self, output: list[np.ndarray], input_images: list[list[np.ndarray]]) -> Any:
         """
         Postprocess the model output
 
         Args:
-        ----
             output: the model output to postprocess
             input_images: the input images used to generate the output
 
         Returns:
-        -------
             Any: the postprocessed output
         """
         raise NotImplementedError
 
-    def __call__(self, inputs: List[np.ndarray]) -> Any:
+    def __call__(self, inputs: list[np.ndarray]) -> Any:
         """
         Call the model on the given inputs
 
         Args:
-        ----
             inputs: the inputs to use
 
         Returns:
-        -------
             Any: the postprocessed output
         """
         self._inputs = inputs

onnxtr/file_utils.py

@@ -6,7 +6,6 @@
 import importlib.metadata
 import importlib.util
 import logging
-from typing import Optional
 
 __all__ = ["requires_package"]
 
@@ -14,12 +13,11 @@ ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
 ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
 
 
-def requires_package(name: str, extra_message: Optional[str] = None) -> None:  # pragma: no cover
+def requires_package(name: str, extra_message: str | None = None) -> None:  # pragma: no cover
     """
     package requirement helper
 
     Args:
-    ----
         name: name of the package
         extra_message: additional message to display if the package is not found
     """

onnxtr/io/elements.py

@@ -3,7 +3,7 @@
 # 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, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 from defusedxml import defuse_stdlib
 
@@ -32,8 +32,8 @@ __all__ = ["Element", "Word", "Artefact", "Line", "Block", "Page", "Document"]
 class Element(NestedObject):
     """Implements an abstract document element with exporting and text rendering capabilities"""
 
-    _children_names: List[str] = []
-    _exported_keys: List[str] = []
+    _children_names: list[str] = []
+    _exported_keys: list[str] = []
 
     def __init__(self, **kwargs: Any) -> None:
         for k, v in kwargs.items():
@@ -42,7 +42,7 @@ class Element(NestedObject):
             else:
                 raise KeyError(f"{self.__class__.__name__} object does not have any attribute named '{k}'")
 
-    def export(self) -> Dict[str, Any]:
+    def export(self) -> dict[str, Any]:
         """Exports the object into a nested dict format"""
         export_dict = {k: getattr(self, k) for k in self._exported_keys}
         for children_name in self._children_names:
@@ -51,7 +51,7 @@ class Element(NestedObject):
         return export_dict
 
     @classmethod
-    def from_dict(cls, save_dict: Dict[str, Any], **kwargs):
+    def from_dict(cls, save_dict: dict[str, Any], **kwargs):
         raise NotImplementedError
 
     def render(self) -> str:
@@ -62,7 +62,6 @@ 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
@@ -71,16 +70,16 @@ class Word(Element):
         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] = []
+    _exported_keys: list[str] = ["value", "confidence", "geometry", "objectness_score", "crop_orientation"]
+    _children_names: list[str] = []
 
     def __init__(
         self,
         value: str,
         confidence: float,
-        geometry: Union[BoundingBox, np.ndarray],
+        geometry: BoundingBox | np.ndarray,
         objectness_score: float,
-        crop_orientation: Dict[str, Any],
+        crop_orientation: dict[str, Any],
     ) -> None:
         super().__init__()
         self.value = value
@@ -97,7 +96,7 @@ class Word(Element):
         return f"value='{self.value}', confidence={self.confidence:.2}"
 
     @classmethod
-    def from_dict(cls, save_dict: Dict[str, Any], **kwargs):
+    def from_dict(cls, save_dict: dict[str, Any], **kwargs):
         kwargs = {k: save_dict[k] for k in cls._exported_keys}
         return cls(**kwargs)
 
@@ -106,15 +105,14 @@ 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] = []
+    _exported_keys: list[str] = ["geometry", "type", "confidence"]
+    _children_names: list[str] = []
 
     def __init__(self, artefact_type: str, confidence: float, geometry: BoundingBox) -> None:
         super().__init__()
@@ -130,7 +128,7 @@ class Artefact(Element):
         return f"type='{self.type}', confidence={self.confidence:.2}"
 
     @classmethod
-    def from_dict(cls, save_dict: Dict[str, Any], **kwargs):
+    def from_dict(cls, save_dict: dict[str, Any], **kwargs):
         kwargs = {k: save_dict[k] for k in cls._exported_keys}
         return cls(**kwargs)
 
@@ -139,22 +137,21 @@ 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] = []
+    _exported_keys: list[str] = ["geometry", "objectness_score"]
+    _children_names: list[str] = ["words"]
+    words: list[Word] = []
 
     def __init__(
         self,
-        words: List[Word],
-        geometry: Optional[Union[BoundingBox, np.ndarray]] = None,
-        objectness_score: Optional[float] = None,
+        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:
@@ -174,7 +171,7 @@ class Line(Element):
         return " ".join(w.render() for w in self.words)
 
     @classmethod
-    def from_dict(cls, save_dict: Dict[str, Any], **kwargs):
+    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"]],
@@ -186,7 +183,6 @@ 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
@@ -194,17 +190,17 @@ class Block(Element):
             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] = []
+    _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: Optional[Union[BoundingBox, np.ndarray]] = None,
-        objectness_score: Optional[float] = None,
+        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:
@@ -227,7 +223,7 @@ class Block(Element):
         return line_break.join(line.render() for line in self.lines)
 
     @classmethod
-    def from_dict(cls, save_dict: Dict[str, Any], **kwargs):
+    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"]],
@@ -240,7 +236,6 @@ class Page(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
@@ -249,18 +244,18 @@ class Page(Element):
         language: a dictionary with the language value and confidence of the prediction
     """
 
-    _exported_keys: List[str] = ["page_idx", "dimensions", "orientation", "language"]
-    _children_names: List[str] = ["blocks"]
-    blocks: List[Block] = []
+    _exported_keys: list[str] = ["page_idx", "dimensions", "orientation", "language"]
+    _children_names: list[str] = ["blocks"]
+    blocks: list[Block] = []
 
     def __init__(
         self,
         page: np.ndarray,
-        blocks: List[Block],
+        blocks: list[Block],
         page_idx: int,
-        dimensions: Tuple[int, int],
-        orientation: Optional[Dict[str, Any]] = None,
-        language: Optional[Dict[str, Any]] = None,
+        dimensions: tuple[int, int],
+        orientation: dict[str, Any] | None = None,
+        language: dict[str, Any] | None = None,
     ) -> None:
         super().__init__(blocks=blocks)
         self.page = page
@@ -295,25 +290,21 @@ class Page(Element):
         """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)
 
-    def export_as_xml(self, file_title: str = "OnnxTR - XML export (hOCR)") -> Tuple[bytes, ET.ElementTree]:
+    def export_as_xml(self, file_title: str = "OnnxTR - 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
@@ -411,7 +402,7 @@ class Page(Element):
         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):
+    def from_dict(cls, save_dict: dict[str, Any], **kwargs):
         kwargs = {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"]]})
         return cls(**kwargs)
@@ -421,16 +412,15 @@ class Document(Element):
     """Implements a document element as a collection of pages
 
     Args:
-    ----
         pages: list of page elements
     """
 
-    _children_names: List[str] = ["pages"]
-    pages: List[Page] = []
+    _children_names: list[str] = ["pages"]
+    pages: list[Page] = []
 
     def __init__(
         self,
-        pages: List[Page],
+        pages: list[Page],
     ) -> None:
         super().__init__(pages=pages)
 
@@ -443,34 +433,30 @@ class Document(Element):
         for result in self.pages:
             result.show(**kwargs)
 
-    def synthesize(self, **kwargs) -> List[np.ndarray]:
+    def synthesize(self, **kwargs) -> list[np.ndarray]:
         """Synthesize all pages from their predictions
 
         Args:
-        ----
             **kwargs: keyword arguments passed to the `Page.synthesize` method
 
-        Returns
-        -------
+        Returns:
             list of synthesized pages
         """
         return [page.synthesize(**kwargs) for page in self.pages]
 
-    def export_as_xml(self, **kwargs) -> List[Tuple[bytes, ET.ElementTree]]:
+    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):
+    def from_dict(cls, save_dict: dict[str, Any], **kwargs):
         kwargs = {k: save_dict[k] for k in cls._exported_keys}
         kwargs.update({"pages": [Page.from_dict(page_dict) for page_dict in save_dict["pages"]]})
         return cls(**kwargs)

onnxtr/io/html.py

@@ -15,12 +15,10 @@ def read_html(url: str, **kwargs: Any) -> bytes:
     >>> doc = read_html("https://www.yoursite.com")
 
     Args:
-    ----
         url: URL of the target web page
         **kwargs: keyword arguments from `weasyprint.HTML`
 
     Returns:
-    -------
         decoded PDF file as a bytes stream
     """
     from weasyprint import HTML

onnxtr/io/image.py

@@ -4,7 +4,6 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 from pathlib import Path
-from typing import Optional, Tuple
 
 import cv2
 import numpy as np
@@ -16,7 +15,7 @@ __all__ = ["read_img_as_numpy"]
 
 def read_img_as_numpy(
     file: AbstractFile,
-    output_size: Optional[Tuple[int, int]] = None,
+    output_size: tuple[int, int] | None = None,
     rgb_output: bool = True,
 ) -> np.ndarray:
     """Read an image file into numpy format
@@ -25,13 +24,11 @@ def read_img_as_numpy(
     >>> page = read_img_as_numpy("path/to/your/doc.jpg")
 
     Args:
-    ----
         file: the path to the image file
         output_size: the expected output size of each page in format H x W
         rgb_output: whether the output ndarray channel order should be RGB instead of BGR.
 
     Returns:
-    -------
         the page decoded as numpy ndarray of shape H x W x 3
     """
     if isinstance(file, (str, Path)):

onnxtr/io/pdf.py

@@ -3,7 +3,7 @@
 # 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, List, Optional
+from typing import Any
 
 import numpy as np
 import pypdfium2 as pdfium
@@ -17,16 +17,15 @@ def read_pdf(
     file: AbstractFile,
     scale: int = 2,
     rgb_mode: bool = True,
-    password: Optional[str] = None,
+    password: str | None = None,
     **kwargs: Any,
-) -> List[np.ndarray]:
+) -> list[np.ndarray]:
     """Read a PDF file and convert it into an image in numpy format
 
     >>> from onnxtr.io import read_pdf
     >>> doc = read_pdf("path/to/your/doc.pdf")
 
     Args:
-    ----
         file: the path to the PDF file
         scale: rendering scale (1 corresponds to 72dpi)
         rgb_mode: if True, the output will be RGB, otherwise BGR
@@ -34,7 +33,6 @@ def read_pdf(
         **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
 
     Returns:
-    -------
         the list of pages decoded as numpy ndarray of shape H x W x C
     """
     # Rasterise pages to numpy ndarrays with pypdfium2

onnxtr/io/reader.py

@@ -3,8 +3,8 @@
 # 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 collections.abc import Sequence
 from pathlib import Path
-from typing import List, Sequence, Union
 
 import numpy as np
 
@@ -22,37 +22,33 @@ class DocumentFile:
     """Read a document from multiple extensions"""
 
     @classmethod
-    def from_pdf(cls, file: AbstractFile, **kwargs) -> List[np.ndarray]:
+    def from_pdf(cls, file: AbstractFile, **kwargs) -> list[np.ndarray]:
         """Read a PDF file
 
         >>> from onnxtr.io import DocumentFile
         >>> doc = DocumentFile.from_pdf("path/to/your/doc.pdf")
 
         Args:
-        ----
             file: the path to the PDF file or a binary stream
             **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
 
         Returns:
-        -------
             the list of pages decoded as numpy ndarray of shape H x W x 3
         """
         return read_pdf(file, **kwargs)
 
     @classmethod
-    def from_url(cls, url: str, **kwargs) -> List[np.ndarray]:
+    def from_url(cls, url: str, **kwargs) -> list[np.ndarray]:
         """Interpret a web page as a PDF document
 
         >>> from onnxtr.io import DocumentFile
         >>> doc = DocumentFile.from_url("https://www.yoursite.com")
 
         Args:
-        ----
             url: the URL of the target web page
             **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
 
         Returns:
-        -------
             the list of pages decoded as numpy ndarray of shape H x W x 3
         """
         requires_package(
@@ -64,19 +60,17 @@ class DocumentFile:
         return cls.from_pdf(pdf_stream, **kwargs)
 
     @classmethod
-    def from_images(cls, files: Union[Sequence[AbstractFile], AbstractFile], **kwargs) -> List[np.ndarray]:
+    def from_images(cls, files: Sequence[AbstractFile] | AbstractFile, **kwargs) -> list[np.ndarray]:
         """Read an image file (or a collection of image files) and convert it into an image in numpy format
 
         >>> from onnxtr.io import DocumentFile
         >>> pages = DocumentFile.from_images(["path/to/your/page1.png", "path/to/your/page2.png"])
 
         Args:
-        ----
             files: the path to the image file or a binary stream, or a collection of those
             **kwargs: additional parameters to :meth:`onnxtr.io.image.read_img_as_numpy`
 
         Returns:
-        -------
             the list of pages decoded as numpy ndarray of shape H x W x 3
         """
         if isinstance(files, (str, Path, bytes)):