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

onnxtr/models/detection/models/differentiable_binarization.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, Optional
+from typing import Any
 
 import numpy as np
 from scipy.special import expit
@@ -14,7 +14,7 @@ from ..postprocessor.base import GeneralDetectionPostProcessor
 __all__ = ["DBNet", "db_resnet50", "db_resnet34", "db_mobilenet_v3_large"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "db_resnet50": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
@@ -43,7 +43,6 @@ class DBNet(Engine):
     """DBNet Onnx loader
 
     Args:
-    ----
         model_path: path or url to onnx model file
         engine_cfg: configuration for the inference engine
         bin_thresh: threshold for binarization of the output feature map
@@ -56,11 +55,11 @@ class DBNet(Engine):
     def __init__(
         self,
         model_path: str,
-        engine_cfg: Optional[EngineConfig] = None,
+        engine_cfg: EngineConfig | None = None,
         bin_thresh: float = 0.3,
         box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
-        cfg: Optional[Dict[str, Any]] = None,
+        cfg: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
@@ -77,10 +76,10 @@ class DBNet(Engine):
         x: np.ndarray,
         return_model_output: bool = False,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         logits = self.run(x)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
 
         prob_map = expit(logits)
         if return_model_output:
@@ -95,7 +94,7 @@ def _dbnet(
     arch: str,
     model_path: str,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> DBNet:
     # Patch the url
@@ -107,7 +106,7 @@ def _dbnet(
 def db_resnet34(
     model_path: str = default_cfgs["db_resnet34"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> DBNet:
     """DBNet as described in `"Real-time Scene Text Detection with Differentiable Binarization"
@@ -120,14 +119,12 @@ def db_resnet34(
     >>> 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 DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _dbnet("db_resnet34", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -136,7 +133,7 @@ def db_resnet34(
 def db_resnet50(
     model_path: str = default_cfgs["db_resnet50"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> DBNet:
     """DBNet as described in `"Real-time Scene Text Detection with Differentiable Binarization"
@@ -149,14 +146,12 @@ def db_resnet50(
     >>> 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 DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _dbnet("db_resnet50", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -165,7 +160,7 @@ def db_resnet50(
 def db_mobilenet_v3_large(
     model_path: str = default_cfgs["db_mobilenet_v3_large"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> DBNet:
     """DBNet as described in `"Real-time Scene Text Detection with Differentiable Binarization"
@@ -178,14 +173,12 @@ def db_mobilenet_v3_large(
     >>> 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 DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _dbnet("db_mobilenet_v3_large", model_path, load_in_8_bit, engine_cfg, **kwargs)

onnxtr/models/detection/models/fast.py

@@ -4,7 +4,7 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 import logging
-from typing import Any, Dict, Optional
+from typing import Any
 
 import numpy as np
 from scipy.special import expit
@@ -15,7 +15,7 @@ from ..postprocessor.base import GeneralDetectionPostProcessor
 __all__ = ["FAST", "fast_tiny", "fast_small", "fast_base"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "fast_tiny": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
@@ -41,7 +41,6 @@ class FAST(Engine):
     """FAST Onnx loader
 
     Args:
-    ----
         model_path: path or url to onnx model file
         engine_cfg: configuration for the inference engine
         bin_thresh: threshold for binarization of the output feature map
@@ -54,11 +53,11 @@ class FAST(Engine):
     def __init__(
         self,
         model_path: str,
-        engine_cfg: Optional[EngineConfig] = None,
+        engine_cfg: EngineConfig | None = None,
         bin_thresh: float = 0.1,
         box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
-        cfg: Optional[Dict[str, Any]] = None,
+        cfg: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
@@ -75,10 +74,10 @@ class FAST(Engine):
         x: np.ndarray,
         return_model_output: bool = False,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         logits = self.run(x)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
 
         prob_map = expit(logits)
         if return_model_output:
@@ -93,7 +92,7 @@ def _fast(
     arch: str,
     model_path: str,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> FAST:
     if load_in_8_bit:
@@ -105,7 +104,7 @@ def _fast(
 def fast_tiny(
     model_path: str = default_cfgs["fast_tiny"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> FAST:
     """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
@@ -118,14 +117,12 @@ def fast_tiny(
     >>> 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 DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _fast("fast_tiny", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -134,7 +131,7 @@ def fast_tiny(
 def fast_small(
     model_path: str = default_cfgs["fast_small"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> FAST:
     """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
@@ -147,14 +144,12 @@ def fast_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 DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _fast("fast_small", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -163,7 +158,7 @@ def fast_small(
 def fast_base(
     model_path: str = default_cfgs["fast_base"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> FAST:
     """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
@@ -176,14 +171,12 @@ def fast_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 DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _fast("fast_base", model_path, load_in_8_bit, engine_cfg, **kwargs)

onnxtr/models/detection/models/linknet.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, Optional
+from typing import Any
 
 import numpy as np
 from scipy.special import expit
@@ -14,7 +14,7 @@ from ..postprocessor.base import GeneralDetectionPostProcessor
 __all__ = ["LinkNet", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "linknet_resnet18": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
@@ -43,7 +43,6 @@ class LinkNet(Engine):
     """LinkNet Onnx loader
 
     Args:
-    ----
         model_path: path or url to onnx model file
         engine_cfg: configuration for the inference engine
         bin_thresh: threshold for binarization of the output feature map
@@ -56,11 +55,11 @@ class LinkNet(Engine):
     def __init__(
         self,
         model_path: str,
-        engine_cfg: Optional[EngineConfig] = None,
+        engine_cfg: EngineConfig | None = None,
         bin_thresh: float = 0.1,
         box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
-        cfg: Optional[Dict[str, Any]] = None,
+        cfg: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
@@ -77,10 +76,10 @@ class LinkNet(Engine):
         x: np.ndarray,
         return_model_output: bool = False,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         logits = self.run(x)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
 
         prob_map = expit(logits)
         if return_model_output:
@@ -95,7 +94,7 @@ def _linknet(
     arch: str,
     model_path: str,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> LinkNet:
     # Patch the url
@@ -107,7 +106,7 @@ def _linknet(
 def linknet_resnet18(
     model_path: str = default_cfgs["linknet_resnet18"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> LinkNet:
     """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
@@ -120,14 +119,12 @@ def linknet_resnet18(
     >>> 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 LinkNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _linknet("linknet_resnet18", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -136,7 +133,7 @@ def linknet_resnet18(
 def linknet_resnet34(
     model_path: str = default_cfgs["linknet_resnet34"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> LinkNet:
     """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
@@ -149,14 +146,12 @@ def linknet_resnet34(
     >>> 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 LinkNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _linknet("linknet_resnet34", model_path, load_in_8_bit, engine_cfg, **kwargs)
@@ -165,7 +160,7 @@ def linknet_resnet34(
 def linknet_resnet50(
     model_path: str = default_cfgs["linknet_resnet50"]["url"],
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> LinkNet:
     """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
@@ -178,14 +173,12 @@ def linknet_resnet50(
     >>> 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 LinkNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _linknet("linknet_resnet50", model_path, load_in_8_bit, engine_cfg, **kwargs)

onnxtr/models/detection/postprocessor/base.py

@@ -5,7 +5,6 @@
 
 # Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
 
-from typing import List, Union
 
 import cv2
 import numpy as np
@@ -21,7 +20,6 @@ class GeneralDetectionPostProcessor(DetectionPostProcessor):
     """Implements a post processor for FAST model.
 
     Args:
-    ----
         bin_thresh: threshold used to binzarized p_map at inference time
         box_thresh: minimal objectness score to consider a box
         assume_straight_pages: whether the inputs were expected to have horizontal text elements
@@ -43,11 +41,9 @@ class GeneralDetectionPostProcessor(DetectionPostProcessor):
         """Expand a polygon (points) by a factor unclip_ratio, and returns a polygon
 
         Args:
-        ----
             points: The first parameter.
 
         Returns:
-        -------
             a box in absolute coordinates (xmin, ymin, xmax, ymax) or (4, 2) array (quadrangle)
         """
         if not self.assume_straight_pages:
@@ -92,24 +88,22 @@ class GeneralDetectionPostProcessor(DetectionPostProcessor):
         """Compute boxes from a bitmap/pred_map: find connected components then filter boxes
 
         Args:
-        ----
             pred: Pred map from differentiable linknet output
             bitmap: Bitmap map computed from pred (binarized)
             angle_tol: Comparison tolerance of the angle with the median angle across the page
             ratio_tol: Under this limit aspect ratio, we cannot resolve the direction of the crop
 
         Returns:
-        -------
             np tensor boxes for the bitmap, each box is a 6-element list
-                containing x, y, w, h, alpha, score for the box
+            containing x, y, w, h, alpha, score for the box
         """
         height, width = bitmap.shape[:2]
-        boxes: List[Union[np.ndarray, List[float]]] = []
+        boxes: list[np.ndarray | list[float]] = []
         # get contours from connected components on the bitmap
         contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         for contour in contours:
             # Check whether smallest enclosing bounding box is not too small
-            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):  # type: ignore[index]
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
                 continue
             # Compute objectness
             if self.assume_straight_pages:

onnxtr/models/detection/predictor/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, Tuple, Union
+from typing import Any
 
 import numpy as np
 
@@ -18,12 +18,11 @@ class DetectionPredictor(NestedObject):
     """Implements an object able to localize text elements in a document
 
     Args:
-    ----
         pre_processor: transform inputs for easier batched model inference
         model: core detection architecture
     """
 
-    _children_names: List[str] = ["pre_processor", "model"]
+    _children_names: list[str] = ["pre_processor", "model"]
 
     def __init__(
         self,
@@ -35,10 +34,10 @@ class DetectionPredictor(NestedObject):
 
     def __call__(
         self,
-        pages: List[np.ndarray],
+        pages: list[np.ndarray],
         return_maps: bool = False,
         **kwargs: Any,
-    ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]]:
+    ) -> list[np.ndarray] | tuple[list[np.ndarray], list[np.ndarray]]:
         # Extract parameters from the preprocessor
         preserve_aspect_ratio = self.pre_processor.resize.preserve_aspect_ratio
         symmetric_pad = self.pre_processor.resize.symmetric_pad

onnxtr/models/detection/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 .. import detection
 from ..engine import EngineConfig
@@ -29,7 +29,7 @@ def _predictor(
     arch: Any,
     assume_straight_pages: bool = True,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> DetectionPredictor:
     if isinstance(arch, str):
@@ -59,8 +59,11 @@ def _predictor(
 def detection_predictor(
     arch: Any = "fast_base",
     assume_straight_pages: bool = True,
+    preserve_aspect_ratio: bool = True,
+    symmetric_pad: bool = True,
+    batch_size: int = 2,
     load_in_8_bit: bool = False,
-    engine_cfg: Optional[EngineConfig] = None,
+    engine_cfg: EngineConfig | None = None,
     **kwargs: Any,
 ) -> DetectionPredictor:
     """Text detection architecture.
@@ -72,15 +75,26 @@ def detection_predictor(
     >>> out = model([input_page])
 
     Args:
-    ----
         arch: name of the architecture or model itself to use (e.g. 'db_resnet50')
         assume_straight_pages: If True, fit straight boxes to the page
+        preserve_aspect_ratio: If True, pad the input document image to preserve the aspect ratio before
+            running the detection model on it
+        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 for the inference engine
         **kwargs: optional keyword arguments passed to the architecture
 
     Returns:
-    -------
         Detection predictor
     """
-    return _predictor(arch, assume_straight_pages, load_in_8_bit, engine_cfg=engine_cfg, **kwargs)
+    return _predictor(
+        arch=arch,
+        assume_straight_pages=assume_straight_pages,
+        preserve_aspect_ratio=preserve_aspect_ratio,
+        symmetric_pad=symmetric_pad,
+        batch_size=batch_size,
+        load_in_8_bit=load_in_8_bit,
+        engine_cfg=engine_cfg,
+        **kwargs,
+    )

onnxtr/models/engine.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, Dict, List, Optional, Tuple, Union
+import logging
+from typing import Any
 
 import numpy as np
 from onnxruntime import (
@@ -25,22 +26,22 @@ class EngineConfig:
     """Implements a configuration class for the engine of a model
 
     Args:
-    ----
         providers: list of providers to use for inference ref.: https://onnxruntime.ai/docs/execution-providers/
         session_options: configuration for the inference session ref.: https://onnxruntime.ai/docs/api/python/api_summary.html#sessionoptions
     """
 
     def __init__(
         self,
-        providers: Optional[Union[List[Tuple[str, Dict[str, Any]]], List[str]]] = None,
-        session_options: Optional[SessionOptions] = None,
+        providers: list[tuple[str, dict[str, Any]]] | list[str] | None = None,
+        session_options: SessionOptions | None = None,
     ):
         self._providers = providers or self._init_providers()
         self._session_options = session_options or self._init_sess_opts()
 
-    def _init_providers(self) -> List[Tuple[str, Dict[str, Any]]]:
+    def _init_providers(self) -> list[tuple[str, dict[str, Any]]]:
         providers: Any = [("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"})]
         available_providers = get_available_providers()
+        logging.info(f"Available providers: {available_providers}")
         if "CUDAExecutionProvider" in available_providers and get_device() == "GPU":  # pragma: no cover
             providers.insert(
                 0,
@@ -66,7 +67,7 @@ class EngineConfig:
         return session_options
 
     @property
-    def providers(self) -> Union[List[Tuple[str, Dict[str, Any]]], List[str]]:
+    def providers(self) -> list[tuple[str, dict[str, Any]]] | list[str]:
         return self._providers
 
     @property
@@ -81,13 +82,12 @@ class Engine:
     """Implements an abstract class for the engine of a model
 
     Args:
-    ----
         url: the url to use to download a model if needed
         engine_cfg: the configuration of the engine
         **kwargs: additional arguments to be passed to `download_from_url`
     """
 
-    def __init__(self, url: str, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any) -> None:
+    def __init__(self, url: str, engine_cfg: EngineConfig | None = None, **kwargs: Any) -> None:
         engine_cfg = engine_cfg if isinstance(engine_cfg, EngineConfig) else EngineConfig()
         archive_path = download_from_url(url, cache_subdir="models", **kwargs) if "http" in url else url
         # Store model path for each model
@@ -97,7 +97,7 @@ class Engine:
         self.runtime = InferenceSession(archive_path, providers=self.providers, sess_options=self.session_options)
         self.runtime_inputs = self.runtime.get_inputs()[0]
         self.tf_exported = int(self.runtime_inputs.shape[-1]) == 3
-        self.fixed_batch_size: Union[int, str] = self.runtime_inputs.shape[
+        self.fixed_batch_size: int | str = self.runtime_inputs.shape[
             0
         ]  # mostly possible with tensorflow exported models
         self.output_name = [output.name for output in self.runtime.get_outputs()]

onnxtr/models/factory/hub.py

@@ -12,7 +12,7 @@ import shutil
 import subprocess
 import textwrap
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from huggingface_hub import (
     HfApi,
@@ -59,7 +59,6 @@ def _save_model_and_config_for_hf_hub(model: Any, save_dir: str, arch: str, task
     """Save model and config to disk for pushing to huggingface hub
 
     Args:
-    ----
         model: Onnx model to be saved
         save_dir: directory to save model and config
         arch: architecture name
@@ -91,7 +90,6 @@ def push_to_hf_hub(
     >>> push_to_hf_hub(model, 'my-model', 'recognition', arch='crnn_mobilenet_v3_small')
 
     Args:
-    ----
         model: Onnx model to be saved
         model_name: name of the model which is also the repository name
         task: task name
@@ -179,20 +177,18 @@ def push_to_hf_hub(
     repo.git_push()
 
 
-def from_hub(repo_id: str, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any):
+def from_hub(repo_id: str, engine_cfg: EngineConfig | None = None, **kwargs: Any):
     """Instantiate & load a pretrained model from HF hub.
 
     >>> from onnxtr.models import from_hub
     >>> model = from_hub("onnxtr/my-model")
 
     Args:
-    ----
         repo_id: HuggingFace model hub repo
         engine_cfg: configuration for the inference engine (optional)
-        kwargs: kwargs of `hf_hub_download`
+        **kwargs: kwargs of `hf_hub_download`
 
     Returns:
-    -------
         Model loaded with the checkpoint
     """
     # Get the config