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

onnxtr/utils/geometry.py

@@ -5,7 +5,6 @@
 
 from copy import deepcopy
 from math import ceil
-from typing import List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
@@ -34,11 +33,9 @@ def bbox_to_polygon(bbox: BoundingBox) -> Polygon4P:
     """Convert a bounding box to a polygon
 
     Args:
-    ----
         bbox: a bounding box
 
     Returns:
-    -------
         a polygon
     """
     return bbox[0], (bbox[1][0], bbox[0][1]), (bbox[0][0], bbox[1][1]), bbox[1]
@@ -48,29 +45,27 @@ def polygon_to_bbox(polygon: Polygon4P) -> BoundingBox:
     """Convert a polygon to a bounding box
 
     Args:
-    ----
         polygon: a polygon
 
     Returns:
-    -------
         a bounding box
     """
     x, y = zip(*polygon)
     return (min(x), min(y)), (max(x), max(y))
 
 
-def detach_scores(boxes: List[np.ndarray]) -> Tuple[List[np.ndarray], List[np.ndarray]]:
+def detach_scores(boxes: list[np.ndarray]) -> tuple[list[np.ndarray], list[np.ndarray]]:
     """Detach the objectness scores from box predictions
+
     Args:
-    ----
         boxes: list of arrays with boxes of shape (N, 5) or (N, 5, 2)
+
     Returns:
-    -------
         a tuple of two lists: the first one contains the boxes without the objectness scores,
         the second one contains the objectness scores
     """
 
-    def _detach(boxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    def _detach(boxes: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
         if boxes.ndim == 2:
             return boxes[:, :-1], boxes[:, -1]
         return boxes[:, :-1], boxes[:, -1, -1]
@@ -83,12 +78,10 @@ def shape_translate(data: np.ndarray, format: str) -> np.ndarray:
     """Translate the shape of the input data to the desired format
 
     Args:
-    ----
         data: input data in shape (B, C, H, W) or (B, H, W, C) or (C, H, W) or (H, W, C)
         format: target format ('BCHW', 'BHWC', 'CHW', or 'HWC')
 
     Returns:
-    -------
         the reshaped data
     """
     # Get the current shape
@@ -120,11 +113,10 @@ def shape_translate(data: np.ndarray, format: str) -> np.ndarray:
         return data
 
 
-def resolve_enclosing_bbox(bboxes: Union[List[BoundingBox], np.ndarray]) -> Union[BoundingBox, np.ndarray]:
+def resolve_enclosing_bbox(bboxes: list[BoundingBox] | np.ndarray) -> BoundingBox | np.ndarray:
     """Compute enclosing bbox either from:
 
     Args:
-    ----
         bboxes: boxes in one of the following formats:
 
             - an array of boxes: (*, 4), where boxes have this shape:
@@ -133,7 +125,6 @@ def resolve_enclosing_bbox(bboxes: Union[List[BoundingBox], np.ndarray]) -> Unio
             - a list of BoundingBox
 
     Returns:
-    -------
         a (1, 4) array (enclosing boxarray), or a BoundingBox
     """
     if isinstance(bboxes, np.ndarray):
@@ -144,11 +135,10 @@ def resolve_enclosing_bbox(bboxes: Union[List[BoundingBox], np.ndarray]) -> Unio
         return (min(x), min(y)), (max(x), max(y))
 
 
-def resolve_enclosing_rbbox(rbboxes: List[np.ndarray], intermed_size: int = 1024) -> np.ndarray:
+def resolve_enclosing_rbbox(rbboxes: list[np.ndarray], intermed_size: int = 1024) -> np.ndarray:
     """Compute enclosing rotated bbox either from:
 
     Args:
-    ----
         rbboxes: boxes in one of the following formats:
 
             - an array of boxes: (*, 4, 2), where boxes have this shape:
@@ -158,26 +148,23 @@ def resolve_enclosing_rbbox(rbboxes: List[np.ndarray], intermed_size: int = 1024
         intermed_size: size of the intermediate image
 
     Returns:
-    -------
         a (4, 2) array (enclosing rotated box)
     """
     cloud: np.ndarray = np.concatenate(rbboxes, axis=0)
     # Convert to absolute for minAreaRect
     cloud *= intermed_size
     rect = cv2.minAreaRect(cloud.astype(np.int32))
-    return cv2.boxPoints(rect) / intermed_size  # type: ignore[return-value]
+    return cv2.boxPoints(rect) / intermed_size
 
 
 def rotate_abs_points(points: np.ndarray, angle: float = 0.0) -> np.ndarray:
     """Rotate points counter-clockwise.
 
     Args:
-    ----
         points: array of size (N, 2)
         angle: angle between -90 and +90 degrees
 
     Returns:
-    -------
         Rotated points
     """
     angle_rad = angle * np.pi / 180.0  # compute radian angle for np functions
@@ -187,16 +174,14 @@ def rotate_abs_points(points: np.ndarray, angle: float = 0.0) -> np.ndarray:
     return np.matmul(points, rotation_mat.T)
 
 
-def compute_expanded_shape(img_shape: Tuple[int, int], angle: float) -> Tuple[int, int]:
+def compute_expanded_shape(img_shape: tuple[int, int], angle: float) -> tuple[int, int]:
     """Compute the shape of an expanded rotated image
 
     Args:
-    ----
         img_shape: the height and width of the image
         angle: angle between -90 and +90 degrees
 
     Returns:
-    -------
         the height and width of the rotated image
     """
     points: np.ndarray = np.array([
@@ -213,21 +198,19 @@ def compute_expanded_shape(img_shape: Tuple[int, int], angle: float) -> Tuple[in
 def rotate_abs_geoms(
     geoms: np.ndarray,
     angle: float,
-    img_shape: Tuple[int, int],
+    img_shape: tuple[int, int],
     expand: bool = True,
 ) -> np.ndarray:
     """Rotate a batch of bounding boxes or polygons by an angle around the
     image center.
 
     Args:
-    ----
         geoms: (N, 4) or (N, 4, 2) array of ABSOLUTE coordinate boxes
         angle: anti-clockwise rotation angle in degrees
         img_shape: the height and width of the image
         expand: whether the image should be padded to avoid information loss
 
     Returns:
-    -------
         A batch of rotated polygons (N, 4, 2)
     """
     # Switch to polygons
@@ -253,19 +236,17 @@ def rotate_abs_geoms(
     return rotated_polys
 
 
-def remap_boxes(loc_preds: np.ndarray, orig_shape: Tuple[int, int], dest_shape: Tuple[int, int]) -> np.ndarray:
+def remap_boxes(loc_preds: np.ndarray, orig_shape: tuple[int, int], dest_shape: tuple[int, int]) -> np.ndarray:
     """Remaps a batch of rotated locpred (N, 4, 2) expressed for an origin_shape to a destination_shape.
     This does not impact the absolute shape of the boxes, but allow to calculate the new relative RotatedBbox
     coordinates after a resizing of the image.
 
     Args:
-    ----
         loc_preds: (N, 4, 2) array of RELATIVE loc_preds
         orig_shape: shape of the origin image
         dest_shape: shape of the destination image
 
     Returns:
-    -------
         A batch of rotated loc_preds (N, 4, 2) expressed in the destination referencial
     """
     if len(dest_shape) != 2:
@@ -284,9 +265,9 @@ def remap_boxes(loc_preds: np.ndarray, orig_shape: Tuple[int, int], dest_shape:
 def rotate_boxes(
     loc_preds: np.ndarray,
     angle: float,
-    orig_shape: Tuple[int, int],
+    orig_shape: tuple[int, int],
     min_angle: float = 1.0,
-    target_shape: Optional[Tuple[int, int]] = None,
+    target_shape: tuple[int, int] | None = None,
 ) -> np.ndarray:
     """Rotate a batch of straight bounding boxes (xmin, ymin, xmax, ymax, c) or rotated bounding boxes
     (4, 2) of an angle, if angle > min_angle, around the center of the page.
@@ -294,7 +275,6 @@ def rotate_boxes(
     is done to remove the padding that is created by rotate_page(expand=True)
 
     Args:
-    ----
         loc_preds: (N, 4) or (N, 4, 2) array of RELATIVE boxes
         angle: angle between -90 and +90 degrees
         orig_shape: shape of the origin image
@@ -302,7 +282,6 @@ def rotate_boxes(
         target_shape: shape of the destination image
 
     Returns:
-    -------
         A batch of rotated boxes (N, 4, 2): or a batch of straight bounding boxes
     """
     # Change format of the boxes to rotated boxes
@@ -349,20 +328,18 @@ def rotate_image(
     """Rotate an image counterclockwise by an given angle.
 
     Args:
-    ----
         image: numpy tensor to rotate
         angle: rotation angle in degrees, between -90 and +90
         expand: whether the image should be padded before the rotation
         preserve_origin_shape: if expand is set to True, resizes the final output to the original image size
 
     Returns:
-    -------
         Rotated array, padded by 0 by default.
     """
     # Compute the expanded padding
     exp_img: np.ndarray
     if expand:
-        exp_shape = compute_expanded_shape(image.shape[:2], angle)  # type: ignore[arg-type]
+        exp_shape = compute_expanded_shape(image.shape[:2], angle)
         h_pad, w_pad = (
             int(max(0, ceil(exp_shape[0] - image.shape[0]))),
             int(max(0, ceil(exp_shape[1] - image.shape[1]))),
@@ -383,7 +360,7 @@ def rotate_image(
             # Pad height
             else:
                 h_pad, w_pad = int(rot_img.shape[1] * image.shape[0] / image.shape[1] - rot_img.shape[0]), 0
-            rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))  # type: ignore[assignment]
+            rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
         if preserve_origin_shape:
             # rescale
             rot_img = cv2.resize(rot_img, image.shape[:-1][::-1], interpolation=cv2.INTER_LINEAR)
@@ -395,11 +372,9 @@ def remove_image_padding(image: np.ndarray) -> np.ndarray:
     """Remove black border padding from an image
 
     Args:
-    ----
         image: numpy tensor to remove padding from
 
     Returns:
-    -------
         Image with padding removed
     """
     # Find the bounding box of the non-black region
@@ -429,16 +404,14 @@ def estimate_page_angle(polys: np.ndarray) -> float:
             return 0.0
 
 
-def convert_to_relative_coords(geoms: np.ndarray, img_shape: Tuple[int, int]) -> np.ndarray:
+def convert_to_relative_coords(geoms: np.ndarray, img_shape: tuple[int, int]) -> np.ndarray:
     """Convert a geometry to relative coordinates
 
     Args:
-    ----
         geoms: a set of polygons of shape (N, 4, 2) or of straight boxes of shape (N, 4)
         img_shape: the height and width of the image
 
     Returns:
-    -------
         the updated geometry
     """
     # Polygon
@@ -456,18 +429,16 @@ def convert_to_relative_coords(geoms: np.ndarray, img_shape: Tuple[int, int]) ->
     raise ValueError(f"invalid format for arg `geoms`: {geoms.shape}")
 
 
-def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True) -> List[np.ndarray]:
+def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True) -> list[np.ndarray]:
     """Created cropped images from list of bounding boxes
 
     Args:
-    ----
         img: input image
         boxes: bounding boxes of shape (N, 4) where N is the number of boxes, and the relative
             coordinates (xmin, ymin, xmax, ymax)
         channels_last: whether the channel dimensions is the last one instead of the last one
 
     Returns:
-    -------
         list of cropped images
     """
     if boxes.shape[0] == 0:
@@ -492,11 +463,10 @@ def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True
 
 def extract_rcrops(
     img: np.ndarray, polys: np.ndarray, dtype=np.float32, channels_last: bool = True, assume_horizontal: bool = False
-) -> List[np.ndarray]:
+) -> list[np.ndarray]:
     """Created cropped images from list of rotated bounding boxes
 
     Args:
-    ----
         img: input image
         polys: bounding boxes of shape (N, 4, 2)
         dtype: target data type of bounding boxes
@@ -504,7 +474,6 @@ def extract_rcrops(
         assume_horizontal: whether the boxes are assumed to be only horizontally oriented
 
     Returns:
-    -------
         list of cropped images
     """
     if polys.shape[0] == 0:
@@ -603,4 +572,4 @@ def extract_rcrops(
             for idx in range(_boxes.shape[0])
         ]
 
-    return crops  # type: ignore[return-value]
+    return crops

onnxtr/utils/multithreading.py

@@ -6,15 +6,16 @@
 
 import multiprocessing as mp
 import os
+from collections.abc import Callable, Iterable, Iterator
 from multiprocessing.pool import ThreadPool
-from typing import Any, Callable, Iterable, Iterator, Optional
+from typing import Any
 
 from onnxtr.file_utils import ENV_VARS_TRUE_VALUES
 
 __all__ = ["multithread_exec"]
 
 
-def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Optional[int] = None) -> Iterator[Any]:
+def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: int | None = None) -> Iterator[Any]:
     """Execute a given function in parallel for each element of a given sequence
 
     >>> from onnxtr.utils.multithreading import multithread_exec
@@ -22,17 +23,14 @@ def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Op
     >>> results = multithread_exec(lambda x: x ** 2, entries)
 
     Args:
-    ----
         func: function to be executed on each element of the iterable
         seq: iterable
         threads: number of workers to be used for multiprocessing
 
     Returns:
-    -------
         iterator of the function's results using the iterable as inputs
 
     Notes:
-    -----
         This function uses ThreadPool from multiprocessing package, which uses `/dev/shm` directory for shared memory.
         If you do not have write permissions for this directory (if you run `onnxtr` on AWS Lambda for instance),
         you might want to disable multiprocessing. To achieve that, set 'ONNXTR_MULTIPROCESSING_DISABLE' to 'TRUE'.

onnxtr/utils/reconstitution.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.
 import logging
-from typing import Any, Dict, Optional
+from typing import Any
 
 import numpy as np
 from anyascii import anyascii
@@ -18,7 +18,7 @@ __all__ = ["synthesize_page"]
 ROTATION_WARNING = False
 
 
-def _warn_rotation(entry: Dict[str, Any]) -> None:  # pragma: no cover
+def _warn_rotation(entry: dict[str, Any]) -> None:  # pragma: no cover
     global ROTATION_WARNING
     if not ROTATION_WARNING and len(entry["geometry"]) == 4:
         logging.warning("Polygons with larger rotations will lead to inaccurate rendering")
@@ -27,11 +27,11 @@ def _warn_rotation(entry: Dict[str, Any]) -> None:  # pragma: no cover
 
 def _synthesize(
     response: Image.Image,
-    entry: Dict[str, Any],
+    entry: dict[str, Any],
     w: int,
     h: int,
     draw_proba: bool = False,
-    font_family: Optional[str] = None,
+    font_family: str | None = None,
     smoothing_factor: float = 0.75,
     min_font_size: int = 6,
     max_font_size: int = 50,
@@ -111,9 +111,9 @@ def _synthesize(
 
 
 def synthesize_page(
-    page: Dict[str, Any],
+    page: dict[str, Any],
     draw_proba: bool = False,
-    font_family: Optional[str] = None,
+    font_family: str | None = None,
     smoothing_factor: float = 0.95,
     min_font_size: int = 8,
     max_font_size: int = 50,
@@ -121,7 +121,6 @@ def synthesize_page(
     """Draw a the content of the element page (OCR response) on a blank page.
 
     Args:
-    ----
         page: exported Page object to represent
         draw_proba: if True, draw words in colors to represent confidence. Blue: p=1, red: p=0
         font_family: family of the font
@@ -130,7 +129,6 @@ def synthesize_page(
         max_font_size: maximum font size
 
     Returns:
-    -------
         the synthesized page
     """
     # Draw template

onnxtr/utils/repr.py

@@ -5,7 +5,6 @@
 
 # Adapted from https://github.com/pytorch/torch/blob/master/torch/nn/modules/module.py
 
-from typing import List
 
 __all__ = ["NestedObject"]
 
@@ -25,7 +24,7 @@ def _addindent(s_, num_spaces):
 class NestedObject:
     """Base class for all nested objects in onnxtr"""
 
-    _children_names: List[str]
+    _children_names: list[str]
 
     def extra_repr(self) -> str:
         return ""

onnxtr/utils/visualization.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, List, Optional, Tuple, Union
+from typing import Any
 
 import cv2
 import matplotlib.patches as patches
@@ -19,9 +19,9 @@ __all__ = ["visualize_page", "draw_boxes"]
 
 def rect_patch(
     geometry: BoundingBox,
-    page_dimensions: Tuple[int, int],
-    label: Optional[str] = None,
-    color: Tuple[float, float, float] = (0, 0, 0),
+    page_dimensions: tuple[int, int],
+    label: str | None = None,
+    color: tuple[float, float, float] = (0, 0, 0),
     alpha: float = 0.3,
     linewidth: int = 2,
     fill: bool = True,
@@ -30,7 +30,6 @@ def rect_patch(
     """Create a matplotlib rectangular patch for the element
 
     Args:
-    ----
         geometry: bounding box of the element
         page_dimensions: dimensions of the Page in format (height, width)
         label: label to display when hovered
@@ -41,7 +40,6 @@ def rect_patch(
         preserve_aspect_ratio: pass True if you passed True to the predictor
 
     Returns:
-    -------
         a rectangular Patch
     """
     if len(geometry) != 2 or any(not isinstance(elt, tuple) or len(elt) != 2 for elt in geometry):
@@ -70,9 +68,9 @@ def rect_patch(
 
 def polygon_patch(
     geometry: np.ndarray,
-    page_dimensions: Tuple[int, int],
-    label: Optional[str] = None,
-    color: Tuple[float, float, float] = (0, 0, 0),
+    page_dimensions: tuple[int, int],
+    label: str | None = None,
+    color: tuple[float, float, float] = (0, 0, 0),
     alpha: float = 0.3,
     linewidth: int = 2,
     fill: bool = True,
@@ -81,7 +79,6 @@ def polygon_patch(
     """Create a matplotlib polygon patch for the element
 
     Args:
-    ----
         geometry: bounding box of the element
         page_dimensions: dimensions of the Page in format (height, width)
         label: label to display when hovered
@@ -92,7 +89,6 @@ def polygon_patch(
         preserve_aspect_ratio: pass True if you passed True to the predictor
 
     Returns:
-    -------
         a polygon Patch
     """
     if not geometry.shape == (4, 2):
@@ -114,20 +110,18 @@ def polygon_patch(
 
 
 def create_obj_patch(
-    geometry: Union[BoundingBox, Polygon4P, np.ndarray],
-    page_dimensions: Tuple[int, int],
+    geometry: BoundingBox | Polygon4P | np.ndarray,
+    page_dimensions: tuple[int, int],
     **kwargs: Any,
 ) -> patches.Patch:
     """Create a matplotlib patch for the element
 
     Args:
-    ----
         geometry: bounding box (straight or rotated) of the element
         page_dimensions: dimensions of the page in format (height, width)
         **kwargs: keyword arguments for the patch
 
     Returns:
-    -------
         a matplotlib Patch
     """
     if isinstance(geometry, tuple):
@@ -141,7 +135,7 @@ def create_obj_patch(
 
 
 def visualize_page(
-    page: Dict[str, Any],
+    page: dict[str, Any],
     image: np.ndarray,
     words_only: bool = True,
     display_artefacts: bool = True,
@@ -163,7 +157,6 @@ def visualize_page(
     >>> plt.show()
 
     Args:
-    ----
         page: the exported Page of a Document
         image: np array of the page, needs to have the same shape than page['dimensions']
         words_only: whether only words should be displayed
@@ -174,7 +167,6 @@ def visualize_page(
         **kwargs: keyword arguments for the polygon patch
 
     Returns:
-    -------
         the matplotlib figure
     """
     # Get proper scale and aspect ratio
@@ -187,7 +179,7 @@ def visualize_page(
     ax.axis("off")
 
     if interactive:
-        artists: List[patches.Patch] = []  # instantiate an empty list of patches (to be drawn on the page)
+        artists: list[patches.Patch] = []  # instantiate an empty list of patches (to be drawn on the page)
 
     for block in page["blocks"]:
         if not words_only:
@@ -266,11 +258,10 @@ def visualize_page(
     return fig
 
 
-def draw_boxes(boxes: np.ndarray, image: np.ndarray, color: Optional[Tuple[int, int, int]] = None, **kwargs) -> None:
+def draw_boxes(boxes: np.ndarray, image: np.ndarray, color: tuple[int, int, int] | None = None, **kwargs) -> None:
     """Draw an array of relative straight boxes on an image
 
     Args:
-    ----
         boxes: array of relative boxes, of shape (*, 4)
         image: np array, float32 or uint8
         color: color to use for bounding box edges

onnxtr/utils/vocabs.py

@@ -4,12 +4,11 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 import string
-from typing import Dict
 
 __all__ = ["VOCABS"]
 
 
-VOCABS: Dict[str, str] = {
+VOCABS: dict[str, str] = {
     "digits": string.digits,
     "ascii_letters": string.ascii_letters,
     "punctuation": string.punctuation,

onnxtr/version.py

@@ -1 +1 @@
-__version__ = 'v0.5.1'
+__version__ = 'v0.6.0'