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

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)
 
@@ -93,18 +89,16 @@ def rotate_sample(
 
 
 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]

doctr/transforms/modules/base.py

@@ -1,11 +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 typing import Any, Callable, List, Optional, Tuple, Union
+from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 
@@ -21,37 +22,36 @@ class SampleCompose(NestedObject):
 
     .. tabs::
 
-        .. tab:: TensorFlow
+        .. tab:: PyTorch
 
             .. code:: python
 
                 >>> import numpy as np
-                >>> import tensorflow as tf
+                >>> import torch
                 >>> from doctr.transforms import SampleCompose, ImageTransform, ColorInversion, RandomRotate
-                >>> transfo = SampleCompose([ImageTransform(ColorInversion((32, 32))), RandomRotate(30)])
-                >>> out, out_boxes = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1), np.zeros((2, 4)))
+                >>> transfos = SampleCompose([ImageTransform(ColorInversion((32, 32))), RandomRotate(30)])
+                >>> out, out_boxes = transfos(torch.rand(8, 64, 64, 3), np.zeros((2, 4)))
 
-        .. tab:: PyTorch
+        .. tab:: TensorFlow
 
             .. code:: python
 
                 >>> import numpy as np
-                >>> import torch
+                >>> import tensorflow as tf
                 >>> from doctr.transforms import SampleCompose, ImageTransform, ColorInversion, RandomRotate
-                >>> transfos = SampleCompose([ImageTransform(ColorInversion((32, 32))), RandomRotate(30)])
-                >>> out, out_boxes = transfos(torch.rand(8, 64, 64, 3), np.zeros((2, 4)))
+                >>> transfo = SampleCompose([ImageTransform(ColorInversion((32, 32))), RandomRotate(30)])
+                >>> out, out_boxes = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1), np.zeros((2, 4)))
 
     Args:
-    ----
         transforms: list of transformation modules
     """
 
-    _children_names: List[str] = ["sample_transforms"]
+    _children_names: list[str] = ["sample_transforms"]
 
-    def __init__(self, transforms: List[Callable[[Any, Any], Tuple[Any, Any]]]) -> None:
+    def __init__(self, transforms: list[Callable[[Any, Any], tuple[Any, Any]]]) -> None:
         self.sample_transforms = transforms
 
-    def __call__(self, x: Any, target: Any) -> Tuple[Any, Any]:
+    def __call__(self, x: Any, target: Any) -> tuple[Any, Any]:
         for t in self.sample_transforms:
             x, target = t(x, target)
 
@@ -63,35 +63,34 @@ class ImageTransform(NestedObject):
 
     .. tabs::
 
-        .. tab:: TensorFlow
+        .. tab:: PyTorch
 
             .. code:: python
 
-                >>> import tensorflow as tf
+                >>> import torch
                 >>> from doctr.transforms import ImageTransform, ColorInversion
                 >>> transfo = ImageTransform(ColorInversion((32, 32)))
-                >>> out, _ = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1), None)
+                >>> out, _ = transfo(torch.rand(8, 64, 64, 3), None)
 
-        .. tab:: PyTorch
+        .. tab:: TensorFlow
 
             .. code:: python
 
-                >>> import torch
+                >>> import tensorflow as tf
                 >>> from doctr.transforms import ImageTransform, ColorInversion
                 >>> transfo = ImageTransform(ColorInversion((32, 32)))
-                >>> out, _ = transfo(torch.rand(8, 64, 64, 3), None)
+                >>> out, _ = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1), None)
 
     Args:
-    ----
         transform: the image transformation module to wrap
     """
 
-    _children_names: List[str] = ["img_transform"]
+    _children_names: list[str] = ["img_transform"]
 
     def __init__(self, transform: Callable[[Any], Any]) -> None:
         self.img_transform = transform
 
-    def __call__(self, img: Any, target: Any) -> Tuple[Any, Any]:
+    def __call__(self, img: Any, target: Any) -> tuple[Any, Any]:
         img = self.img_transform(img)
         return img, target
 
@@ -102,26 +101,25 @@ class ColorInversion(NestedObject):
 
     .. tabs::
 
-        .. tab:: TensorFlow
+        .. tab:: PyTorch
 
             .. code:: python
 
-                >>> import tensorflow as tf
+                >>> import torch
                 >>> from doctr.transforms import ColorInversion
                 >>> transfo = ColorInversion(min_val=0.6)
-                >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
+                >>> out = transfo(torch.rand(8, 64, 64, 3))
 
-        .. tab:: PyTorch
+        .. tab:: TensorFlow
 
             .. code:: python
 
-                >>> import torch
+                >>> import tensorflow as tf
                 >>> from doctr.transforms import ColorInversion
                 >>> transfo = ColorInversion(min_val=0.6)
-                >>> out = transfo(torch.rand(8, 64, 64, 3))
+                >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
-    ----
         min_val: range [min_val, 1] to colorize RGB pixels
     """
 
@@ -140,35 +138,34 @@ class OneOf(NestedObject):
 
     .. tabs::
 
-        .. tab:: TensorFlow
+        .. tab:: PyTorch
 
             .. code:: python
 
-                >>> import tensorflow as tf
+                >>> import torch
                 >>> from doctr.transforms import OneOf
                 >>> transfo = OneOf([JpegQuality(), Gamma()])
-                >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
+                >>> out = transfo(torch.rand(1, 64, 64, 3))
 
-        .. tab:: PyTorch
+        .. tab:: TensorFlow
 
             .. code:: python
 
-                >>> import torch
+                >>> import tensorflow as tf
                 >>> from doctr.transforms import OneOf
                 >>> transfo = OneOf([JpegQuality(), Gamma()])
-                >>> out = transfo(torch.rand(1, 64, 64, 3))
+                >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
-    ----
         transforms: list of transformations, one only will be picked
     """
 
-    _children_names: List[str] = ["transforms"]
+    _children_names: list[str] = ["transforms"]
 
-    def __init__(self, transforms: List[Callable[[Any], Any]]) -> None:
+    def __init__(self, transforms: list[Callable[[Any], Any]]) -> None:
         self.transforms = transforms
 
-    def __call__(self, img: Any, target: Optional[np.ndarray] = None) -> Union[Any, Tuple[Any, np.ndarray]]:
+    def __call__(self, img: Any, target: np.ndarray | None = None) -> Any | tuple[Any, np.ndarray]:
         # Pick transformation
         transfo = self.transforms[int(random.random() * len(self.transforms))]
         # Apply
@@ -180,26 +177,25 @@ class RandomApply(NestedObject):
 
     .. tabs::
 
-        .. tab:: TensorFlow
+        .. tab:: PyTorch
 
             .. code:: python
 
-                >>> import tensorflow as tf
+                >>> import torch
                 >>> from doctr.transforms import RandomApply
                 >>> transfo = RandomApply(Gamma(), p=.5)
-                >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
+                >>> out = transfo(torch.rand(1, 64, 64, 3))
 
-        .. tab:: PyTorch
+        .. tab:: TensorFlow
 
             .. code:: python
 
-                >>> import torch
+                >>> import tensorflow as tf
                 >>> from doctr.transforms import RandomApply
                 >>> transfo = RandomApply(Gamma(), p=.5)
-                >>> out = transfo(torch.rand(1, 64, 64, 3))
+                >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
-    ----
         transform: transformation to apply
         p: probability to apply
     """
@@ -211,7 +207,7 @@ class RandomApply(NestedObject):
     def extra_repr(self) -> str:
         return f"transform={self.transform}, p={self.p}"
 
-    def __call__(self, img: Any, target: Optional[np.ndarray] = None) -> Union[Any, Tuple[Any, np.ndarray]]:
+    def __call__(self, img: Any, target: np.ndarray | None = None) -> Any | tuple[Any, np.ndarray]:
         if random.random() < self.p:
             return self.transform(img) if target is None else self.transform(img, target)  # type: ignore[call-arg]
         return img if target is None else (img, target)
@@ -224,9 +220,7 @@ class RandomRotate(NestedObject):
         :align: center
 
     Args:
-    ----
-        max_angle: maximum angle for rotation, in degrees. Angles will be uniformly picked in
-            [-max_angle, max_angle]
+        max_angle: maximum angle for rotation, in degrees. Angles will be uniformly picked in [-max_angle, max_angle]
         expand: whether the image should be padded before the rotation
     """
 
@@ -237,7 +231,7 @@ class RandomRotate(NestedObject):
     def extra_repr(self) -> str:
         return f"max_angle={self.max_angle}, expand={self.expand}"
 
-    def __call__(self, img: Any, target: np.ndarray) -> Tuple[Any, np.ndarray]:
+    def __call__(self, img: Any, target: np.ndarray) -> tuple[Any, np.ndarray]:
         angle = random.uniform(-self.max_angle, self.max_angle)
         r_img, r_polys = F.rotate_sample(img, target, angle, self.expand)
         # Removes deleted boxes
@@ -249,19 +243,18 @@ class RandomCrop(NestedObject):
     """Randomly crop a tensor image and its boxes
 
     Args:
-    ----
         scale: tuple of floats, relative (min_area, max_area) of the crop
         ratio: tuple of float, relative (min_ratio, max_ratio) where ratio = h/w
     """
 
-    def __init__(self, scale: Tuple[float, float] = (0.08, 1.0), ratio: Tuple[float, float] = (0.75, 1.33)) -> None:
+    def __init__(self, scale: tuple[float, float] = (0.08, 1.0), ratio: tuple[float, float] = (0.75, 1.33)) -> None:
         self.scale = scale
         self.ratio = ratio
 
     def extra_repr(self) -> str:
         return f"scale={self.scale}, ratio={self.ratio}"
 
-    def __call__(self, img: Any, target: np.ndarray) -> Tuple[Any, np.ndarray]:
+    def __call__(self, img: Any, target: np.ndarray) -> tuple[Any, np.ndarray]:
         scale = random.uniform(self.scale[0], self.scale[1])
         ratio = random.uniform(self.ratio[0], self.ratio[1])