python-doctr 0.9.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

doctr/models/utils/tensorflow.py

@@ -1,12 +1,11 @@
-# Copyright (C) 2021-2024, Mindee.
+# Copyright (C) 2021-2025, 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.
 
 import logging
-import os
-from typing import Any, Callable, List, Optional, Tuple, Union
-from zipfile import ZipFile
+from collections.abc import Callable
+from typing import Any
 
 import tensorflow as tf
 import tf2onnx
@@ -19,6 +18,7 @@ logging.getLogger("tensorflow").setLevel(logging.DEBUG)
 
 __all__ = [
     "load_pretrained_params",
+    "_build_model",
     "conv_sequence",
     "IntermediateLayerGetter",
     "export_model_to_onnx",
@@ -36,51 +36,50 @@ def _bf16_to_float32(x: tf.Tensor) -> tf.Tensor:
     return tf.cast(x, tf.float32) if x.dtype == tf.bfloat16 else x
 
 
+def _build_model(model: Model):
+    """Build a model by calling it once with dummy input
+
+    Args:
+        model: the model to be built
+    """
+    model(tf.zeros((1, *model.cfg["input_shape"])), training=False)
+
+
 def load_pretrained_params(
     model: Model,
-    url: Optional[str] = None,
-    hash_prefix: Optional[str] = None,
-    overwrite: bool = False,
-    internal_name: str = "weights",
+    url: str | None = None,
+    hash_prefix: str | None = None,
+    skip_mismatch: bool = False,
     **kwargs: Any,
 ) -> None:
     """Load a set of parameters onto a model
 
     >>> from doctr.models import load_pretrained_params
-    >>> load_pretrained_params(model, "https://yoursource.com/yourcheckpoint-yourhash.zip")
+    >>> load_pretrained_params(model, "https://yoursource.com/yourcheckpoint-yourhash.weights.h5")
 
     Args:
-    ----
         model: the keras model to be loaded
         url: URL of the zipped set of parameters
         hash_prefix: first characters of SHA256 expected hash
-        overwrite: should the zip extraction be enforced if the archive has already been extracted
-        internal_name: name of the ckpt files
+        skip_mismatch: skip loading layers with mismatched shapes
         **kwargs: additional arguments to be passed to `doctr.utils.data.download_from_url`
     """
     if url is None:
         logging.warning("Invalid model URL, using default initialization.")
     else:
         archive_path = download_from_url(url, hash_prefix=hash_prefix, cache_subdir="models", **kwargs)
-
-        # Unzip the archive
-        params_path = archive_path.parent.joinpath(archive_path.stem)
-        if not params_path.is_dir() or overwrite:
-            with ZipFile(archive_path, "r") as f:
-                f.extractall(path=params_path)
-
         # Load weights
-        model.load_weights(f"{params_path}{os.sep}{internal_name}")
+        model.load_weights(archive_path, skip_mismatch=skip_mismatch)
 
 
 def conv_sequence(
     out_channels: int,
-    activation: Optional[Union[str, Callable]] = None,
+    activation: str | Callable | None = None,
     bn: bool = False,
     padding: str = "same",
     kernel_initializer: str = "he_normal",
     **kwargs: Any,
-) -> List[layers.Layer]:
+) -> list[layers.Layer]:
     """Builds a convolutional-based layer sequence
 
     >>> from tensorflow.keras import Sequential
@@ -88,7 +87,6 @@ def conv_sequence(
     >>> module = Sequential(conv_sequence(32, 'relu', True, kernel_size=3, input_shape=[224, 224, 3]))
 
     Args:
-    ----
         out_channels: number of output channels
         activation: activation to be used (default: no activation)
         bn: should a batch normalization layer be added
@@ -97,7 +95,6 @@ def conv_sequence(
         **kwargs: additional arguments to be passed to the convolutional layer
 
     Returns:
-    -------
         list of layers
     """
     # No bias before Batch norm
@@ -125,12 +122,11 @@ class IntermediateLayerGetter(Model):
     >>> feat_extractor = IntermediateLayerGetter(ResNet50(include_top=False, pooling=False), target_layers)
 
     Args:
-    ----
         model: the model to extract feature maps from
         layer_names: the list of layers to retrieve the feature map from
     """
 
-    def __init__(self, model: Model, layer_names: List[str]) -> None:
+    def __init__(self, model: Model, layer_names: list[str]) -> None:
         intermediate_fmaps = [model.get_layer(layer_name).get_output_at(0) for layer_name in layer_names]
         super().__init__(model.input, outputs=intermediate_fmaps)
 
@@ -139,8 +135,8 @@ class IntermediateLayerGetter(Model):
 
 
 def export_model_to_onnx(
-    model: Model, model_name: str, dummy_input: List[tf.TensorSpec], **kwargs: Any
-) -> Tuple[str, List[str]]:
+    model: Model, model_name: str, dummy_input: list[tf.TensorSpec], **kwargs: Any
+) -> tuple[str, list[str]]:
     """Export model to ONNX format.
 
     >>> import tensorflow as tf
@@ -151,16 +147,18 @@ def export_model_to_onnx(
     >>> dummy_input=[tf.TensorSpec([None, 32, 32, 3], tf.float32, name="input")])
 
     Args:
-    ----
         model: the keras model to be exported
         model_name: the name for the exported model
         dummy_input: the dummy input to the model
         kwargs: additional arguments to be passed to tf2onnx
 
     Returns:
-    -------
         the path to the exported model and a list with the output layer names
     """
+    # get the users eager mode
+    eager_mode = tf.executing_eagerly()
+    # set eager mode to true to avoid issues with tf2onnx
+    tf.config.run_functions_eagerly(True)
     large_model = kwargs.get("large_model", False)
     model_proto, _ = tf2onnx.convert.from_keras(
         model,
@@ -171,6 +169,9 @@ def export_model_to_onnx(
     # Get the output layer names
     output = [n.name for n in model_proto.graph.output]
 
+    # reset the eager mode to the users mode
+    tf.config.run_functions_eagerly(eager_mode)
+
     # models which are too large (weights > 2GB while converting to ONNX) needs to be handled
     # about an external tensor storage where the graph and weights are seperatly stored in a archive
     if large_model:

doctr/models/zoo.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2024, Mindee.
+# Copyright (C) 2021-2025, 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.
@@ -83,7 +83,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
@@ -108,7 +107,6 @@ def ocr_predictor(
         kwargs: keyword args of `OCRPredictor`
 
     Returns:
-    -------
         OCR predictor
     """
     return _predictor(
@@ -197,7 +195,6 @@ def kie_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
@@ -222,7 +219,6 @@ def kie_predictor(
         kwargs: keyword args of `OCRPredictor`
 
     Returns:
-    -------
         KIE predictor
     """
     return _kie_predictor(

doctr/transforms/functional/__init__.py

@@ -1,6 +1,6 @@
 from doctr.file_utils import is_tf_available, is_torch_available
 
-if is_tf_available():
-    from .tensorflow import *
-elif is_torch_available():
+if is_torch_available():
     from .pytorch import *
+elif is_tf_available():
+    from .tensorflow import *

doctr/transforms/functional/base.py

@@ -1,9 +1,8 @@
-# Copyright (C) 2021-2024, Mindee.
+# Copyright (C) 2021-2025, 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 Tuple, Union
 
 import cv2
 import numpy as np
@@ -15,17 +14,15 @@ __all__ = ["crop_boxes", "create_shadow_mask"]
 
 def crop_boxes(
     boxes: np.ndarray,
-    crop_box: Union[Tuple[int, int, int, int], Tuple[float, float, float, float]],
+    crop_box: tuple[int, int, int, int] | tuple[float, float, float, float],
 ) -> np.ndarray:
     """Crop localization boxes
 
     Args:
-    ----
         boxes: ndarray of shape (N, 4) in relative or abs coordinates
         crop_box: box (xmin, ymin, xmax, ymax) to crop the image, in the same coord format that the boxes
 
     Returns:
-    -------
         the cropped boxes
     """
     is_box_rel = boxes.max() <= 1
@@ -49,17 +46,15 @@ def crop_boxes(
     return boxes[is_valid]
 
 
-def expand_line(line: np.ndarray, target_shape: Tuple[int, int]) -> Tuple[float, float]:
+def expand_line(line: np.ndarray, target_shape: tuple[int, int]) -> tuple[float, float]:
     """Expands a 2-point line, so that the first is on the edge. In other terms, we extend the line in
     the same direction until we meet one of the edges.
 
     Args:
-    ----
         line: array of shape (2, 2) of the point supposed to be on one edge, and the shadow tip.
         target_shape: the desired mask shape
 
     Returns:
-    -------
         2D coordinates of the first point once we extended the line (on one of the edges)
     """
     if any(coord == 0 or coord == size for coord, size in zip(line[0], target_shape[::-1])):
@@ -112,7 +107,7 @@ def expand_line(line: np.ndarray, target_shape: Tuple[int, int]) -> Tuple[float,
 
 
 def create_shadow_mask(
-    target_shape: Tuple[int, int],
+    target_shape: tuple[int, int],
     min_base_width=0.3,
     max_tip_width=0.5,
     max_tip_height=0.3,
@@ -120,14 +115,12 @@ def create_shadow_mask(
     """Creates a random shadow mask
 
     Args:
-    ----
         target_shape: the target shape (H, W)
         min_base_width: the relative minimum shadow base width
         max_tip_width: the relative maximum shadow tip width
         max_tip_height: the relative maximum shadow tip height
 
     Returns:
-    -------
         a numpy ndarray of shape (H, W, 1) with values in the range [0, 1]
     """
     # Default base is top

doctr/transforms/functional/pytorch.py

@@ -1,13 +1,13 @@
-# Copyright (C) 2021-2024, Mindee.
+# Copyright (C) 2021-2025, 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 copy import deepcopy
-from typing import Tuple
 
 import numpy as np
 import torch
+from scipy.ndimage import gaussian_filter
 from torchvision.transforms import functional as F
 
 from doctr.utils.geometry import rotate_abs_geoms
@@ -21,12 +21,10 @@ def invert_colors(img: torch.Tensor, min_val: float = 0.6) -> torch.Tensor:
     """Invert the colors of an image
 
     Args:
-    ----
         img : torch.Tensor, the image to invert
         min_val : minimum value of the random shift
 
     Returns:
-    -------
         the inverted image
     """
     out = F.rgb_to_grayscale(img, num_output_channels=3)
@@ -35,9 +33,9 @@ def invert_colors(img: torch.Tensor, min_val: float = 0.6) -> torch.Tensor:
     rgb_shift = min_val + (1 - min_val) * torch.rand(shift_shape)
     # Inverse the color
     if out.dtype == torch.uint8:
-        out = (out.to(dtype=rgb_shift.dtype) * rgb_shift).to(dtype=torch.uint8)
+        out = (out.to(dtype=rgb_shift.dtype) * rgb_shift).to(dtype=torch.uint8)  # type: ignore[attr-defined]
     else:
-        out = out * rgb_shift.to(dtype=out.dtype)
+        out = out * rgb_shift.to(dtype=out.dtype)  # type: ignore[attr-defined]
     # Inverse the color
     out = 255 - out if out.dtype == torch.uint8 else 1 - out
     return out
@@ -48,18 +46,16 @@ def rotate_sample(
     geoms: np.ndarray,
     angle: float,
     expand: bool = False,
-) -> Tuple[torch.Tensor, np.ndarray]:
+) -> tuple[torch.Tensor, np.ndarray]:
     """Rotate image around the center, interpolation=NEAREST, pad with 0 (black)
 
     Args:
-    ----
         img: image to rotate
         geoms: array of geometries of shape (N, 4) or (N, 4, 2)
         angle: angle in degrees. +: counter-clockwise, -: clockwise
         expand: whether the image should be padded before the rotation
 
     Returns:
-    -------
         A tuple of rotated img (tensor), rotated geometries of shape (N, 4, 2)
     """
     rotated_img = F.rotate(img, angle=angle, fill=0, expand=expand)  # Interpolation NEAREST by default
@@ -81,7 +77,7 @@ def rotate_sample(
     rotated_geoms: np.ndarray = rotate_abs_geoms(
         _geoms,
         angle,
-        img.shape[1:],  # type: ignore[arg-type]
+        img.shape[1:],
         expand,
     ).astype(np.float32)
 
@@ -89,22 +85,20 @@ def rotate_sample(
     rotated_geoms[..., 0] = rotated_geoms[..., 0] / rotated_img.shape[2]
     rotated_geoms[..., 1] = rotated_geoms[..., 1] / rotated_img.shape[1]
 
-    return rotated_img, np.clip(rotated_geoms, 0, 1)
+    return rotated_img, np.clip(np.around(rotated_geoms, decimals=15), 0, 1)
 
 
 def crop_detection(
-    img: torch.Tensor, boxes: np.ndarray, crop_box: Tuple[float, float, float, float]
-) -> Tuple[torch.Tensor, np.ndarray]:
+    img: torch.Tensor, boxes: np.ndarray, crop_box: tuple[float, float, float, float]
+) -> tuple[torch.Tensor, np.ndarray]:
     """Crop and image and associated bboxes
 
     Args:
-    ----
         img: image to crop
         boxes: array of boxes to clip, absolute (int) or relative (float)
         crop_box: box (xmin, ymin, xmax, ymax) to crop the image. Relative coords.
 
     Returns:
-    -------
         A tuple of cropped image, cropped boxes, where the image is not resized.
     """
     if any(val < 0 or val > 1 for val in crop_box):
@@ -119,27 +113,25 @@ def crop_detection(
     return cropped_img, boxes
 
 
-def random_shadow(img: torch.Tensor, opacity_range: Tuple[float, float], **kwargs) -> torch.Tensor:
-    """Crop and image and associated bboxes
+def random_shadow(img: torch.Tensor, opacity_range: tuple[float, float], **kwargs) -> torch.Tensor:
+    """Apply a random shadow effect to an image using NumPy for blurring.
 
     Args:
-    ----
-        img: image to modify
-        opacity_range: the minimum and maximum desired opacity of the shadow
-        **kwargs: additional arguments to pass to `create_shadow_mask`
+        img: Image to modify (C, H, W) as a PyTorch tensor.
+        opacity_range: The minimum and maximum desired opacity of the shadow.
+        **kwargs: Additional arguments to pass to `create_shadow_mask`.
 
     Returns:
-    -------
-        shaded image
+        Shadowed image as a PyTorch tensor (same shape as input).
     """
-    shadow_mask = create_shadow_mask(img.shape[1:], **kwargs)  # type: ignore[arg-type]
-
+    shadow_mask = create_shadow_mask(img.shape[1:], **kwargs)
     opacity = np.random.uniform(*opacity_range)
-    shadow_tensor = 1 - torch.from_numpy(shadow_mask[None, ...])
 
-    # Add some blur to make it believable
-    k = 7 + 2 * int(4 * np.random.rand(1))
+    # Apply Gaussian blur to the shadow mask
     sigma = np.random.uniform(0.5, 5.0)
-    shadow_tensor = F.gaussian_blur(shadow_tensor, k, sigma=[sigma, sigma])
+    blurred_mask = gaussian_filter(shadow_mask, sigma=sigma)
+
+    shadow_tensor = 1 - torch.from_numpy(blurred_mask).float()
+    shadow_tensor = shadow_tensor.to(img.device).unsqueeze(0)  # Add channel dimension
 
     return opacity * shadow_tensor * img + (1 - opacity) * img

doctr/transforms/functional/tensorflow.py

@@ -1,12 +1,12 @@
-# Copyright (C) 2021-2024, Mindee.
+# Copyright (C) 2021-2025, 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.
 
 import math
 import random
+from collections.abc import Iterable
 from copy import deepcopy
-from typing import Iterable, Optional, Tuple, Union
 
 import numpy as np
 import tensorflow as tf
@@ -22,12 +22,10 @@ def invert_colors(img: tf.Tensor, min_val: float = 0.6) -> tf.Tensor:
     """Invert the colors of an image
 
     Args:
-    ----
         img : tf.Tensor, the image to invert
         min_val : minimum value of the random shift
 
     Returns:
-    -------
         the inverted image
     """
     out = tf.image.rgb_to_grayscale(img)  # Convert to gray
@@ -48,13 +46,11 @@ def rotated_img_tensor(img: tf.Tensor, angle: float, expand: bool = False) -> tf
     """Rotate image around the center, interpolation=NEAREST, pad with 0 (black)
 
     Args:
-    ----
         img: image to rotate
         angle: angle in degrees. +: counter-clockwise, -: clockwise
         expand: whether the image should be padded before the rotation
 
     Returns:
-    -------
         the rotated image (tensor)
     """
     # Compute the expanded padding
@@ -103,18 +99,16 @@ def rotate_sample(
     geoms: np.ndarray,
     angle: float,
     expand: bool = False,
-) -> Tuple[tf.Tensor, np.ndarray]:
+) -> tuple[tf.Tensor, np.ndarray]:
     """Rotate image around the center, interpolation=NEAREST, pad with 0 (black)
 
     Args:
-    ----
         img: image to rotate
         geoms: array of geometries of shape (N, 4) or (N, 4, 2)
         angle: angle in degrees. +: counter-clockwise, -: clockwise
         expand: whether the image should be padded before the rotation
 
     Returns:
-    -------
         A tuple of rotated img (tensor), rotated boxes (np array)
     """
     # Rotated the image
@@ -140,22 +134,20 @@ def rotate_sample(
     rotated_geoms[..., 0] = rotated_geoms[..., 0] / rotated_img.shape[1]
     rotated_geoms[..., 1] = rotated_geoms[..., 1] / rotated_img.shape[0]
 
-    return rotated_img, np.clip(rotated_geoms, 0, 1)
+    return rotated_img, np.clip(np.around(rotated_geoms, decimals=15), 0, 1)
 
 
 def crop_detection(
-    img: tf.Tensor, boxes: np.ndarray, crop_box: Tuple[float, float, float, float]
-) -> Tuple[tf.Tensor, np.ndarray]:
+    img: tf.Tensor, boxes: np.ndarray, crop_box: tuple[float, float, float, float]
+) -> tuple[tf.Tensor, np.ndarray]:
     """Crop and image and associated bboxes
 
     Args:
-    ----
         img: image to crop
         boxes: array of boxes to clip, absolute (int) or relative (float)
         crop_box: box (xmin, ymin, xmax, ymax) to crop the image. Relative coords.
 
     Returns:
-    -------
         A tuple of cropped image, cropped boxes, where the image is not resized.
     """
     if any(val < 0 or val > 1 for val in crop_box):
@@ -172,16 +164,15 @@ def crop_detection(
 
 def _gaussian_filter(
     img: tf.Tensor,
-    kernel_size: Union[int, Iterable[int]],
+    kernel_size: int | Iterable[int],
     sigma: float,
-    mode: Optional[str] = None,
-    pad_value: Optional[int] = 0,
+    mode: str | None = None,
+    pad_value: int = 0,
 ):
     """Apply Gaussian filter to image.
     Adapted from: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/filters.py
 
     Args:
-    ----
         img: image to filter of shape (N, H, W, C)
         kernel_size: kernel size of the filter
         sigma: standard deviation of the Gaussian filter
@@ -189,7 +180,6 @@ def _gaussian_filter(
         pad_value: value to pad the image with
 
     Returns:
-    -------
         A tensor of shape (N, H, W, C)
     """
     ksize = tf.convert_to_tensor(tf.broadcast_to(kernel_size, [2]), dtype=tf.int32)
@@ -235,17 +225,15 @@ def _gaussian_filter(
     return tf.nn.depthwise_conv2d(img, g, [1, 1, 1, 1], padding="VALID", data_format="NHWC")
 
 
-def random_shadow(img: tf.Tensor, opacity_range: Tuple[float, float], **kwargs) -> tf.Tensor:
+def random_shadow(img: tf.Tensor, opacity_range: tuple[float, float], **kwargs) -> tf.Tensor:
     """Apply a random shadow to a given image
 
     Args:
-    ----
         img: image to modify
         opacity_range: the minimum and maximum desired opacity of the shadow
         **kwargs: additional arguments to pass to `create_shadow_mask`
 
     Returns:
-    -------
         shadowed image
     """
     shadow_mask = create_shadow_mask(img.shape[:2], **kwargs)

doctr/transforms/modules/__init__.py

@@ -2,7 +2,7 @@ from doctr.file_utils import is_tf_available, is_torch_available
 
 from .base import *
 
-if is_tf_available():
-    from .tensorflow import *
-elif is_torch_available():
-    from .pytorch import *  # type: ignore[assignment]
+if is_torch_available():
+    from .pytorch import *
+elif is_tf_available():
+    from .tensorflow import *  # type: ignore[assignment]