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

onnxtr/models/recognition/models/sar.py

@@ -4,7 +4,7 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 from copy import deepcopy
-from typing import Any, Dict, Optional
+from typing import Any
 
 import numpy as np
 from scipy.special import softmax
@@ -16,7 +16,7 @@ from ..core import RecognitionPostProcessor
 
 __all__ = ["SAR", "sar_resnet31"]
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "sar_resnet31": {
         "mean": (0.694, 0.695, 0.693),
         "std": (0.299, 0.296, 0.301),
@@ -32,7 +32,6 @@ class SAR(Engine):
     """SAR Onnx loader
 
     Args:
-    ----
         model_path: path to onnx model file
         vocab: vocabulary used for encoding
         engine_cfg: configuration for the inference engine
@@ -44,8 +43,8 @@ class SAR(Engine):
         self,
         model_path: str,
         vocab: str,
-        engine_cfg: Optional[EngineConfig] = None,
-        cfg: Optional[Dict[str, Any]] = None,
+        engine_cfg: EngineConfig | None = None,
+        cfg: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
@@ -59,10 +58,10 @@ class SAR(Engine):
         self,
         x: np.ndarray,
         return_model_output: bool = False,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         logits = self.run(x)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
         if return_model_output:
             out["out_map"] = logits
 
@@ -75,7 +74,6 @@ class SARPostProcessor(RecognitionPostProcessor):
     """Post processor for SAR architectures
 
     Args:
-    ----
         embedding: string containing the ordered sequence of supported characters
     """
 
@@ -105,7 +103,7 @@ def _sar(
     arch: str,
     model_path: str,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> SAR:
     # Patch the config
@@ -124,7 +122,7 @@ def _sar(
 def sar_resnet31(
     model_path: str = default_cfgs["sar_resnet31"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> SAR:
     """SAR with a resnet-31 feature extractor as described in `"Show, Attend and Read:A Simple and Strong
@@ -137,14 +135,12 @@ def sar_resnet31(
     >>> out = model(input_tensor)
 
     Args:
-    ----
         model_path: path to onnx model file, defaults to url in default_cfgs
         load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False
         engine_cfg: configuration for the inference engine
         **kwargs: keyword arguments of the SAR architecture
 
     Returns:
-    -------
         text recognition architecture
     """
     return _sar("sar_resnet31", model_path, load_in_8_bit, engine_cfg, **kwargs)

onnxtr/models/recognition/models/vitstr.py

@@ -4,7 +4,7 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 from copy import deepcopy
-from typing import Any, Dict, Optional
+from typing import Any
 
 import numpy as np
 from scipy.special import softmax
@@ -16,7 +16,7 @@ from ..core import RecognitionPostProcessor
 
 __all__ = ["ViTSTR", "vitstr_small", "vitstr_base"]
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "vitstr_small": {
         "mean": (0.694, 0.695, 0.693),
         "std": (0.299, 0.296, 0.301),
@@ -40,7 +40,6 @@ class ViTSTR(Engine):
     """ViTSTR Onnx loader
 
     Args:
-    ----
         model_path: path to onnx model file
         vocab: vocabulary used for encoding
         engine_cfg: configuration for the inference engine
@@ -52,8 +51,8 @@ class ViTSTR(Engine):
         self,
         model_path: str,
         vocab: str,
-        engine_cfg: Optional[EngineConfig] = None,
-        cfg: Optional[Dict[str, Any]] = None,
+        engine_cfg: EngineConfig | None = None,
+        cfg: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
@@ -67,10 +66,10 @@ class ViTSTR(Engine):
         self,
         x: np.ndarray,
         return_model_output: bool = False,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         logits = self.run(x)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
         if return_model_output:
             out["out_map"] = logits
 
@@ -83,7 +82,6 @@ class ViTSTRPostProcessor(RecognitionPostProcessor):
     """Post processor for ViTSTR architecture
 
     Args:
-    ----
         vocab: string containing the ordered sequence of supported characters
     """
 
@@ -115,7 +113,7 @@ def _vitstr(
     arch: str,
     model_path: str,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> ViTSTR:
     # Patch the config
@@ -134,7 +132,7 @@ def _vitstr(
 def vitstr_small(
     model_path: str = default_cfgs["vitstr_small"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> ViTSTR:
     """ViTSTR-Small as described in `"Vision Transformer for Fast and Efficient Scene Text Recognition"
@@ -147,14 +145,12 @@ def vitstr_small(
     >>> out = model(input_tensor)
 
     Args:
-    ----
         model_path: path to onnx model file, defaults to url in default_cfgs
         load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False
         engine_cfg: configuration for the inference engine
         **kwargs: keyword arguments of the ViTSTR architecture
 
     Returns:
-    -------
         text recognition architecture
     """
     return _vitstr("vitstr_small", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -163,7 +159,7 @@ def vitstr_small(
 def vitstr_base(
     model_path: str = default_cfgs["vitstr_base"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> ViTSTR:
     """ViTSTR-Base as described in `"Vision Transformer for Fast and Efficient Scene Text Recognition"
@@ -176,14 +172,12 @@ def vitstr_base(
     >>> out = model(input_tensor)
 
     Args:
-    ----
         model_path: path to onnx model file, defaults to url in default_cfgs
         load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False
         engine_cfg: configuration for the inference engine
         **kwargs: keyword arguments of the ViTSTR architecture
 
     Returns:
-    -------
         text recognition architecture
     """
     return _vitstr("vitstr_base", model_path, load_in_8_bit, engine_cfg, **kwargs)

onnxtr/models/recognition/predictor/_utils.py

@@ -3,7 +3,6 @@
 # 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 List, Tuple, Union
 
 import numpy as np
 
@@ -13,16 +12,15 @@ __all__ = ["split_crops", "remap_preds"]
 
 
 def split_crops(
-    crops: List[np.ndarray],
+    crops: list[np.ndarray],
     max_ratio: float,
     target_ratio: int,
     dilation: float,
     channels_last: bool = True,
-) -> Tuple[List[np.ndarray], List[Union[int, Tuple[int, int]]], bool]:
+) -> tuple[list[np.ndarray], list[int | tuple[int, int]], bool]:
     """Chunk crops horizontally to match a given aspect ratio
 
     Args:
-    ----
         crops: list of numpy array of shape (H, W, 3) if channels_last or (3, H, W) otherwise
         max_ratio: the maximum aspect ratio that won't trigger the chunk
         target_ratio: when crops are chunked, they will be chunked to match this aspect ratio
@@ -30,12 +28,11 @@ def split_crops(
         channels_last: whether the numpy array has dimensions in channels last order
 
     Returns:
-    -------
         a tuple with the new crops, their mapping, and a boolean specifying whether any remap is required
     """
     _remap_required = False
-    crop_map: List[Union[int, Tuple[int, int]]] = []
-    new_crops: List[np.ndarray] = []
+    crop_map: list[int | tuple[int, int]] = []
+    new_crops: list[np.ndarray] = []
     for crop in crops:
         h, w = crop.shape[:2] if channels_last else crop.shape[-2:]
         aspect_ratio = w / h
@@ -71,8 +68,8 @@ def split_crops(
 
 
 def remap_preds(
-    preds: List[Tuple[str, float]], crop_map: List[Union[int, Tuple[int, int]]], dilation: float
-) -> List[Tuple[str, float]]:
+    preds: list[tuple[str, float]], crop_map: list[int | tuple[int, int]], dilation: float
+) -> list[tuple[str, float]]:
     remapped_out = []
     for _idx in crop_map:
         # Crop hasn't been split

onnxtr/models/recognition/predictor/base.py

@@ -3,7 +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 typing import Any, List, Sequence, Tuple
+from collections.abc import Sequence
+from typing import Any
 
 import numpy as np
 
@@ -19,7 +20,6 @@ class RecognitionPredictor(NestedObject):
     """Implements an object able to identify character sequences in images
 
     Args:
-    ----
         pre_processor: transform inputs for easier batched model inference
         model: core recognition architecture
         split_wide_crops: wether to use crop splitting for high aspect ratio crops
@@ -43,7 +43,7 @@ class RecognitionPredictor(NestedObject):
         self,
         crops: Sequence[np.ndarray],
         **kwargs: Any,
-    ) -> List[Tuple[str, float]]:
+    ) -> list[tuple[str, float]]:
         if len(crops) == 0:
             return []
         # Dimension check

onnxtr/models/recognition/utils.py

@@ -3,7 +3,6 @@
 # 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 List
 
 from rapidfuzz.distance import Levenshtein
 
@@ -14,14 +13,12 @@ def merge_strings(a: str, b: str, dil_factor: float) -> str:
     """Merges 2 character sequences in the best way to maximize the alignment of their overlapping characters.
 
     Args:
-    ----
         a: first char seq, suffix should be similar to b's prefix.
         b: second char seq, prefix should be similar to a's suffix.
         dil_factor: dilation factor of the boxes to overlap, should be > 1. This parameter is
             only used when the mother sequence is splitted on a character repetition
 
     Returns:
-    -------
         A merged character sequence.
 
     Example::
@@ -61,17 +58,15 @@ def merge_strings(a: str, b: str, dil_factor: float) -> str:
     return a[:-1] + b[index - 1 :]
 
 
-def merge_multi_strings(seq_list: List[str], dil_factor: float) -> str:
+def merge_multi_strings(seq_list: list[str], dil_factor: float) -> str:
     """Recursively merges consecutive string sequences with overlapping characters.
 
     Args:
-    ----
         seq_list: list of sequences to merge. Sequences need to be ordered from left to right.
         dil_factor: dilation factor of the boxes to overlap, should be > 1. This parameter is
             only used when the mother sequence is splitted on a character repetition
 
     Returns:
-    -------
         A merged character sequence
 
     Example::
@@ -80,7 +75,7 @@ def merge_multi_strings(seq_list: List[str], dil_factor: float) -> str:
         'abcdefghijkl'
     """
 
-    def _recursive_merge(a: str, seq_list: List[str], dil_factor: float) -> str:
+    def _recursive_merge(a: str, seq_list: list[str], dil_factor: float) -> str:
         # Recursive version of compute_overlap
         if len(seq_list) == 1:
             return merge_strings(a, seq_list[0], dil_factor)

onnxtr/models/recognition/zoo.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
 
 from .. import recognition
 from ..engine import EngineConfig
@@ -13,7 +13,7 @@ from .predictor import RecognitionPredictor
 __all__ = ["recognition_predictor"]
 
 
-ARCHS: List[str] = [
+ARCHS: list[str] = [
     "crnn_vgg16_bn",
     "crnn_mobilenet_v3_small",
     "crnn_mobilenet_v3_large",
@@ -26,7 +26,7 @@ ARCHS: List[str] = [
 
 
 def _predictor(
-    arch: Any, load_in_8_bit: bool = False, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any
+    arch: Any, load_in_8_bit: bool = False, engine_cfg: EngineConfig | None = None, **kwargs: Any
 ) -> RecognitionPredictor:
     if isinstance(arch, str):
         if arch not in ARCHS:
@@ -50,7 +50,12 @@ def _predictor(
 
 
 def recognition_predictor(
-    arch: Any = "crnn_vgg16_bn", load_in_8_bit: bool = False, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any
+    arch: Any = "crnn_vgg16_bn",
+    symmetric_pad: bool = False,
+    batch_size: int = 128,
+    load_in_8_bit: bool = False,
+    engine_cfg: EngineConfig | None = None,
+    **kwargs: Any,
 ) -> RecognitionPredictor:
     """Text recognition architecture.
 
@@ -62,14 +67,21 @@ def recognition_predictor(
         >>> out = model([input_page])
 
     Args:
-    ----
         arch: name of the architecture or model itself to use (e.g. 'crnn_vgg16_bn')
+        symmetric_pad: if True, pad the image symmetrically instead of padding at the bottom-right
+        batch_size: number of samples the model processes in parallel
         load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False
         engine_cfg: configuration of inference engine
         **kwargs: optional parameters to be passed to the architecture
 
     Returns:
-    -------
         Recognition predictor
     """
-    return _predictor(arch, load_in_8_bit, engine_cfg, **kwargs)
+    return _predictor(
+        arch=arch,
+        symmetric_pad=symmetric_pad,
+        batch_size=batch_size,
+        load_in_8_bit=load_in_8_bit,
+        engine_cfg=engine_cfg,
+        **kwargs,
+    )

onnxtr/models/zoo.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, Optional
+from typing import Any
 
 from .detection.zoo import detection_predictor
 from .engine import EngineConfig
@@ -25,9 +25,9 @@ def _predictor(
     straighten_pages: bool = False,
     detect_language: bool = False,
     load_in_8_bit: bool = False,
-    det_engine_cfg: Optional[EngineConfig] = None,
-    reco_engine_cfg: Optional[EngineConfig] = None,
-    clf_engine_cfg: Optional[EngineConfig] = None,
+    det_engine_cfg: EngineConfig | None = None,
+    reco_engine_cfg: EngineConfig | None = None,
+    clf_engine_cfg: EngineConfig | None = None,
     **kwargs,
 ) -> OCRPredictor:
     # Detection
@@ -74,9 +74,9 @@ def ocr_predictor(
     straighten_pages: bool = False,
     detect_language: bool = False,
     load_in_8_bit: bool = False,
-    det_engine_cfg: Optional[EngineConfig] = None,
-    reco_engine_cfg: Optional[EngineConfig] = None,
-    clf_engine_cfg: Optional[EngineConfig] = None,
+    det_engine_cfg: EngineConfig | None = None,
+    reco_engine_cfg: EngineConfig | None = None,
+    clf_engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> OCRPredictor:
     """End-to-end OCR architecture using one model for localization, and another for text recognition.
@@ -88,7 +88,6 @@ def ocr_predictor(
     >>> out = model([input_page])
 
     Args:
-    ----
         det_arch: name of the detection architecture or the model itself to use
             (e.g. 'db_resnet50', 'db_mobilenet_v3_large')
         reco_arch: name of the recognition architecture or the model itself to use
@@ -115,7 +114,6 @@ def ocr_predictor(
         kwargs: keyword args of `OCRPredictor`
 
     Returns:
-    -------
         OCR predictor
     """
     return _predictor(

onnxtr/transforms/base.py

@@ -3,7 +3,6 @@
 # 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 Tuple, Union
 
 import numpy as np
 from PIL import Image, ImageOps
@@ -12,11 +11,18 @@ __all__ = ["Resize", "Normalize"]
 
 
 class Resize:
-    """Resize the input image to the given size"""
+    """Resize the input image to the given size
+
+    Args:
+        size: the target size of the image
+        interpolation: the interpolation method to use
+        preserve_aspect_ratio: whether to preserve the aspect ratio of the image
+        symmetric_pad: whether to symmetrically pad the image
+    """
 
     def __init__(
         self,
-        size: Union[int, Tuple[int, int]],
+        size: int | tuple[int, int],
         interpolation=Image.Resampling.BILINEAR,
         preserve_aspect_ratio: bool = False,
         symmetric_pad: bool = False,
@@ -72,12 +78,17 @@ class Resize:
 
 
 class Normalize:
-    """Normalize the input image"""
+    """Normalize the input image
+
+    Args:
+        mean: mean values to subtract
+        std: standard deviation values to divide
+    """
 
     def __init__(
         self,
-        mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406),
-        std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225),
+        mean: float | tuple[float, float, float] = (0.485, 0.456, 0.406),
+        std: float | tuple[float, float, float] = (0.229, 0.224, 0.225),
     ) -> None:
         self.mean = mean
         self.std = std

onnxtr/utils/common_types.py

@@ -4,15 +4,14 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 from pathlib import Path
-from typing import List, Tuple, Union
 
 __all__ = ["Point2D", "BoundingBox", "Polygon4P", "Polygon", "Bbox"]
 
 
-Point2D = Tuple[float, float]
-BoundingBox = Tuple[Point2D, Point2D]
-Polygon4P = Tuple[Point2D, Point2D, Point2D, Point2D]
-Polygon = List[Point2D]
-AbstractPath = Union[str, Path]
-AbstractFile = Union[AbstractPath, bytes]
-Bbox = Tuple[float, float, float, float]
+Point2D = tuple[float, float]
+BoundingBox = tuple[Point2D, Point2D]
+Polygon4P = tuple[Point2D, Point2D, Point2D, Point2D]
+Polygon = list[Point2D]
+AbstractPath = str | Path
+AbstractFile = AbstractPath | bytes
+Bbox = tuple[float, float, float, float]

onnxtr/utils/data.py

@@ -13,7 +13,6 @@ import urllib
 import urllib.error
 import urllib.request
 from pathlib import Path
-from typing import Optional, Union
 
 from tqdm.auto import tqdm
 
@@ -25,7 +24,7 @@ HASH_REGEX = re.compile(r"-([a-f0-9]*)\.")
 USER_AGENT = "felixdittrich92/OnnxTR"
 
 
-def _urlretrieve(url: str, filename: Union[Path, str], chunk_size: int = 1024) -> None:
+def _urlretrieve(url: str, filename: Path | str, chunk_size: int = 1024) -> None:
     with open(filename, "wb") as fh:
         with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT})) as response:
             with tqdm(total=response.length) as pbar:
@@ -36,7 +35,7 @@ def _urlretrieve(url: str, filename: Union[Path, str], chunk_size: int = 1024) -
                     fh.write(chunk)
 
 
-def _check_integrity(file_path: Union[str, Path], hash_prefix: str) -> bool:
+def _check_integrity(file_path: str | Path, hash_prefix: str) -> bool:
     with open(file_path, "rb") as f:
         sha_hash = hashlib.sha256(f.read()).hexdigest()
 
@@ -45,10 +44,10 @@ def _check_integrity(file_path: Union[str, Path], hash_prefix: str) -> bool:
 
 def download_from_url(
     url: str,
-    file_name: Optional[str] = None,
-    hash_prefix: Optional[str] = None,
-    cache_dir: Optional[str] = None,
-    cache_subdir: Optional[str] = None,
+    file_name: str | None = None,
+    hash_prefix: str | None = None,
+    cache_dir: str | None = None,
+    cache_subdir: str | None = None,
 ) -> Path:
     """Download a file using its URL
 
@@ -56,7 +55,6 @@ def download_from_url(
     >>> download_from_url("https://yoursource.com/yourcheckpoint-yourhash.zip")
 
     Args:
-    ----
         url: the URL of the file to download
         file_name: optional name of the file once downloaded
         hash_prefix: optional expected SHA256 hash of the file
@@ -64,11 +62,9 @@ def download_from_url(
         cache_subdir: subfolder to use in the cache
 
     Returns:
-    -------
         the location of the downloaded file
 
     Note:
-    ----
         You can change cache directory location by using `ONNXTR_CACHE_DIR` environment variable.
     """
     if not isinstance(file_name, str):
@@ -112,7 +108,7 @@ def download_from_url(
     except (urllib.error.URLError, IOError) as e:  # pragma: no cover
         if url[:5] == "https":
             url = url.replace("https:", "http:")
-            print("Failed download. Trying https -> http instead." f" Downloading {url} to {file_path}")
+            print(f"Failed download. Trying https -> http instead. Downloading {url} to {file_path}")
             _urlretrieve(url, file_path)
         else:
             raise e

onnxtr/utils/fonts.py

@@ -5,25 +5,20 @@
 
 import logging
 import platform
-from typing import Optional, Union
 
 from PIL import ImageFont
 
 __all__ = ["get_font"]
 
 
-def get_font(
-    font_family: Optional[str] = None, font_size: int = 13
-) -> Union[ImageFont.FreeTypeFont, ImageFont.ImageFont]:
+def get_font(font_family: str | None = None, font_size: int = 13) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
     """Resolves a compatible ImageFont for the system
 
     Args:
-    ----
         font_family: the font family to use
         font_size: the size of the font upon rendering
 
     Returns:
-    -------
         the Pillow font
     """
     # Font selection