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

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])
 

doctr/transforms/modules/pytorch.py

@@ -1,21 +1,29 @@
-# 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
-from typing import Optional, Tuple, Union
 
 import numpy as np
 import torch
 from PIL.Image import Image
+from scipy.ndimage import gaussian_filter
 from torch.nn.functional import pad
 from torchvision.transforms import functional as F
 from torchvision.transforms import transforms as T
 
 from ..functional.pytorch import random_shadow
 
-__all__ = ["Resize", "GaussianNoise", "ChannelShuffle", "RandomHorizontalFlip", "RandomShadow", "RandomResize"]
+__all__ = [
+    "Resize",
+    "GaussianNoise",
+    "ChannelShuffle",
+    "RandomHorizontalFlip",
+    "RandomShadow",
+    "RandomResize",
+    "GaussianBlur",
+]
 
 
 class Resize(T.Resize):
@@ -23,7 +31,7 @@ class Resize(T.Resize):
 
     def __init__(
         self,
-        size: Union[int, Tuple[int, int]],
+        size: int | tuple[int, int],
         interpolation=F.InterpolationMode.BILINEAR,
         preserve_aspect_ratio: bool = False,
         symmetric_pad: bool = False,
@@ -38,8 +46,8 @@ class Resize(T.Resize):
     def forward(
         self,
         img: torch.Tensor,
-        target: Optional[np.ndarray] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, np.ndarray]]:
+        target: np.ndarray | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, np.ndarray]:
         if isinstance(self.size, int):
             target_ratio = img.shape[-2] / img.shape[-1]
         else:
@@ -74,16 +82,18 @@ class Resize(T.Resize):
                 if self.symmetric_pad:
                     half_pad = (math.ceil(_pad[1] / 2), math.ceil(_pad[3] / 2))
                     _pad = (half_pad[0], _pad[1] - half_pad[0], half_pad[1], _pad[3] - half_pad[1])
+                # Pad image
                 img = pad(img, _pad)
 
             # In case boxes are provided, resize boxes if needed (for detection task if preserve aspect ratio)
             if target is not None:
+                if self.symmetric_pad:
+                    offset = half_pad[0] / img.shape[-1], half_pad[1] / img.shape[-2]
+
                 if self.preserve_aspect_ratio:
                     # Get absolute coords
                     if target.shape[1:] == (4,):
                         if isinstance(self.size, (tuple, list)) and self.symmetric_pad:
-                            if np.max(target) <= 1:
-                                offset = half_pad[0] / img.shape[-1], half_pad[1] / img.shape[-2]
                             target[:, [0, 2]] = offset[0] + target[:, [0, 2]] * raw_shape[-1] / img.shape[-1]
                             target[:, [1, 3]] = offset[1] + target[:, [1, 3]] * raw_shape[-2] / img.shape[-2]
                         else:
@@ -91,16 +101,15 @@ class Resize(T.Resize):
                             target[:, [1, 3]] *= raw_shape[-2] / img.shape[-2]
                     elif target.shape[1:] == (4, 2):
                         if isinstance(self.size, (tuple, list)) and self.symmetric_pad:
-                            if np.max(target) <= 1:
-                                offset = half_pad[0] / img.shape[-1], half_pad[1] / img.shape[-2]
                             target[..., 0] = offset[0] + target[..., 0] * raw_shape[-1] / img.shape[-1]
                             target[..., 1] = offset[1] + target[..., 1] * raw_shape[-2] / img.shape[-2]
                         else:
                             target[..., 0] *= raw_shape[-1] / img.shape[-1]
                             target[..., 1] *= raw_shape[-2] / img.shape[-2]
                     else:
-                        raise AssertionError
-                return img, target
+                        raise AssertionError("Boxes should be in the format (n_boxes, 4, 2) or (n_boxes, 4)")
+
+                return img, np.clip(target, 0, 1)
 
             return img
 
@@ -121,7 +130,6 @@ class GaussianNoise(torch.nn.Module):
     >>> out = transfo(torch.rand((3, 224, 224)))
 
     Args:
-    ----
         mean : mean of the gaussian distribution
         std : std of the gaussian distribution
     """
@@ -135,14 +143,47 @@ class GaussianNoise(torch.nn.Module):
         # Reshape the distribution
         noise = self.mean + 2 * self.std * torch.rand(x.shape, device=x.device) - self.std
         if x.dtype == torch.uint8:
-            return (x + 255 * noise).round().clamp(0, 255).to(dtype=torch.uint8)
+            return (x + 255 * noise).round().clamp(0, 255).to(dtype=torch.uint8)  # type: ignore[attr-defined]
         else:
-            return (x + noise.to(dtype=x.dtype)).clamp(0, 1)
+            return (x + noise.to(dtype=x.dtype)).clamp(0, 1)  # type: ignore[attr-defined]
 
     def extra_repr(self) -> str:
         return f"mean={self.mean}, std={self.std}"
 
 
+class GaussianBlur(torch.nn.Module):
+    """Apply Gaussian Blur to the input tensor
+
+    >>> import torch
+    >>> from doctr.transforms import GaussianBlur
+    >>> transfo = GaussianBlur(sigma=(0.0, 1.0))
+
+    Args:
+        sigma : standard deviation range for the gaussian kernel
+    """
+
+    def __init__(self, sigma: tuple[float, float]) -> None:
+        super().__init__()
+        self.sigma_range = sigma
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Sample a random sigma value within the specified range
+        sigma = torch.empty(1).uniform_(*self.sigma_range).item()
+
+        # Apply Gaussian blur along spatial dimensions only
+        blurred = torch.tensor(
+            gaussian_filter(
+                x.numpy(),
+                sigma=sigma,
+                mode="reflect",
+                truncate=4.0,
+            ),
+            dtype=x.dtype,
+            device=x.device,
+        )
+        return blurred
+
+
 class ChannelShuffle(torch.nn.Module):
     """Randomly shuffle channel order of a given image"""
 
@@ -158,9 +199,7 @@ class ChannelShuffle(torch.nn.Module):
 class RandomHorizontalFlip(T.RandomHorizontalFlip):
     """Randomly flip the input image horizontally"""
 
-    def forward(
-        self, img: Union[torch.Tensor, Image], target: np.ndarray
-    ) -> Tuple[Union[torch.Tensor, Image], np.ndarray]:
+    def forward(self, img: torch.Tensor | Image, target: np.ndarray) -> tuple[torch.Tensor | Image, np.ndarray]:
         if torch.rand(1) < self.p:
             _img = F.hflip(img)
             _target = target.copy()
@@ -182,11 +221,10 @@ class RandomShadow(torch.nn.Module):
     >>> out = transfo(torch.rand((3, 64, 64)))
 
     Args:
-    ----
         opacity_range : minimum and maximum opacity of the shade
     """
 
-    def __init__(self, opacity_range: Optional[Tuple[float, float]] = None) -> None:
+    def __init__(self, opacity_range: tuple[float, float] | None = None) -> None:
         super().__init__()
         self.opacity_range = opacity_range if isinstance(opacity_range, tuple) else (0.2, 0.8)
 
@@ -195,7 +233,7 @@ class RandomShadow(torch.nn.Module):
         try:
             if x.dtype == torch.uint8:
                 return (
-                    (
+                    (  # type: ignore[attr-defined]
                         255
                         * random_shadow(
                             x.to(dtype=torch.float32) / 255,
@@ -224,20 +262,19 @@ class RandomResize(torch.nn.Module):
     >>> out = transfo(torch.rand((3, 64, 64)))
 
     Args:
-    ----
         scale_range: range of the resizing factor for width and height (independently)
         preserve_aspect_ratio: whether to preserve the aspect ratio of the image,
-            given a float value, the aspect ratio will be preserved with this probability
+        given a float value, the aspect ratio will be preserved with this probability
         symmetric_pad: whether to symmetrically pad the image,
-            given a float value, the symmetric padding will be applied with this probability
+        given a float value, the symmetric padding will be applied with this probability
         p: probability to apply the transformation
     """
 
     def __init__(
         self,
-        scale_range: Tuple[float, float] = (0.3, 0.9),
-        preserve_aspect_ratio: Union[bool, float] = False,
-        symmetric_pad: Union[bool, float] = False,
+        scale_range: tuple[float, float] = (0.3, 0.9),
+        preserve_aspect_ratio: bool | float = False,
+        symmetric_pad: bool | float = False,
         p: float = 0.5,
     ) -> None:
         super().__init__()
@@ -247,7 +284,7 @@ class RandomResize(torch.nn.Module):
         self.p = p
         self._resize = Resize
 
-    def forward(self, img: torch.Tensor, target: np.ndarray) -> Tuple[torch.Tensor, np.ndarray]:
+    def forward(self, img: torch.Tensor, target: np.ndarray) -> tuple[torch.Tensor, np.ndarray]:
         if torch.rand(1) < self.p:
             scale_h = np.random.uniform(*self.scale_range)
             scale_w = np.random.uniform(*self.scale_range)