python-doctr 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

doctr/transforms/functional/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -18,15 +18,26 @@ __all__ = ["invert_colors", "rotate_sample", "crop_detection", "random_shadow"]
 
 
 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)
     # Random RGB shift
     shift_shape = [img.shape[0], 3, 1, 1] if img.ndim == 4 else [3, 1, 1]
     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
@@ -41,12 +52,14 @@ def rotate_sample(
     """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
@@ -68,7 +81,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)
 
@@ -85,11 +98,13 @@ def crop_detection(
     """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):
@@ -108,14 +123,16 @@ def random_shadow(img: torch.Tensor, opacity_range: Tuple[float, float], **kwarg
     """Crop and image and associated bboxes
 
     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:
+    -------
         shaded image
     """
-
-    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, ...])

doctr/transforms/functional/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -15,10 +15,21 @@ from doctr.utils.geometry import compute_expanded_shape, rotate_abs_geoms
 
 from .base import create_shadow_mask, crop_boxes
 
-__all__ = ["invert_colors", "rotate_sample", "crop_detection", "random_shadow"]
+__all__ = ["invert_colors", "rotate_sample", "crop_detection", "random_shadow", "rotated_img_tensor"]
 
 
 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
     # Random RGB shift
     shift_shape = [img.shape[0], 1, 1, 3] if img.ndim == 4 else [1, 1, 3]
@@ -37,11 +48,13 @@ 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
@@ -94,12 +107,14 @@ def rotate_sample(
     """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
@@ -134,11 +149,13 @@ def crop_detection(
     """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):
@@ -164,14 +181,15 @@ def _gaussian_filter(
     Adapted from: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/filters.py
 
     Args:
-
-        input: image to filter of shape (N, H, W, C)
+    ----
+        img: image to filter of shape (N, H, W, C)
         kernel_size: kernel size of the filter
         sigma: standard deviation of the Gaussian filter
         mode: padding mode, one of "CONSTANT", "REFLECT", "SYMMETRIC"
         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)
@@ -221,13 +239,15 @@ def random_shadow(img: tf.Tensor, opacity_range: Tuple[float, float], **kwargs)
     """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)
 
     opacity = np.random.uniform(*opacity_range)

doctr/transforms/modules/base.py

@@ -1,11 +1,11 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
@@ -42,6 +42,7 @@ class SampleCompose(NestedObject):
                 >>> out, out_boxes = transfos(torch.rand(8, 64, 64, 3), np.zeros((2, 4)))
 
     Args:
+    ----
         transforms: list of transformation modules
     """
 
@@ -81,6 +82,7 @@ class ImageTransform(NestedObject):
                 >>> out, _ = transfo(torch.rand(8, 64, 64, 3), None)
 
     Args:
+    ----
         transform: the image transformation module to wrap
     """
 
@@ -119,6 +121,7 @@ class ColorInversion(NestedObject):
                 >>> out = transfo(torch.rand(8, 64, 64, 3))
 
     Args:
+    ----
         min_val: range [min_val, 1] to colorize RGB pixels
     """
 
@@ -156,6 +159,7 @@ class OneOf(NestedObject):
                 >>> out = transfo(torch.rand(1, 64, 64, 3))
 
     Args:
+    ----
         transforms: list of transformations, one only will be picked
     """
 
@@ -195,6 +199,7 @@ class RandomApply(NestedObject):
                 >>> out = transfo(torch.rand(1, 64, 64, 3))
 
     Args:
+    ----
         transform: transformation to apply
         p: probability to apply
     """
@@ -206,10 +211,10 @@ class RandomApply(NestedObject):
     def extra_repr(self) -> str:
         return f"transform={self.transform}, p={self.p}"
 
-    def __call__(self, img: Any) -> Any:
+    def __call__(self, img: Any, target: Optional[np.ndarray] = None) -> Union[Any, Tuple[Any, np.ndarray]]:
         if random.random() < self.p:
-            return self.transform(img)
-        return img
+            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)
 
 
 class RandomRotate(NestedObject):
@@ -219,6 +224,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]
         expand: whether the image should be padded before the rotation
@@ -243,6 +249,7 @@ 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
     """

doctr/transforms/modules/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -19,6 +19,8 @@ __all__ = ["Resize", "GaussianNoise", "ChannelShuffle", "RandomHorizontalFlip",
 
 
 class Resize(T.Resize):
+    """Resize the input image to the given size"""
+
     def __init__(
         self,
         size: Union[int, Tuple[int, int]],
@@ -119,6 +121,7 @@ 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
     """
@@ -132,9 +135,9 @@ 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}"
@@ -153,17 +156,11 @@ class ChannelShuffle(torch.nn.Module):
 
 
 class RandomHorizontalFlip(T.RandomHorizontalFlip):
+    """Randomly flip the input image horizontally"""
+
     def forward(
         self, img: Union[torch.Tensor, Image], target: Dict[str, Any]
     ) -> Tuple[Union[torch.Tensor, Image], Dict[str, Any]]:
-        """
-        Args:
-            img: Image to be flipped.
-            target: Dictionary with boxes (in relative coordinates of shape (N, 4)) and labels as keys
-
-        Returns:
-            Tuple of PIL Image or Tensor and target
-        """
         if torch.rand(1) < self.p:
             _img = F.hflip(img)
             _target = target.copy()
@@ -182,6 +179,7 @@ class RandomShadow(torch.nn.Module):
     >>> out = transfo(torch.rand((3, 64, 64)))
 
     Args:
+    ----
         opacity_range : minimum and maximum opacity of the shade
     """
 
@@ -201,7 +199,7 @@ class RandomShadow(torch.nn.Module):
                             self.opacity_range,
                         )
                     )
-                    .round()
+                    .round()  # type: ignore[attr-defined]
                     .clip(0, 255)
                     .to(dtype=torch.uint8)
                 )

doctr/transforms/modules/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -42,6 +42,7 @@ class Compose(NestedObject):
     >>> out = transfos(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         transforms: list of transformation modules
     """
 
@@ -66,6 +67,7 @@ class Resize(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         output_size: expected output size
         method: interpolation method
         preserve_aspect_ratio: if `True`, preserve aspect ratio and pad the rest with zeros
@@ -83,6 +85,7 @@ class Resize(NestedObject):
         self.method = method
         self.preserve_aspect_ratio = preserve_aspect_ratio
         self.symmetric_pad = symmetric_pad
+        self.antialias = True
 
         if isinstance(self.output_size, int):
             self.wanted_size = (self.output_size, self.output_size)
@@ -104,7 +107,7 @@ class Resize(NestedObject):
     ) -> Union[tf.Tensor, Tuple[tf.Tensor, np.ndarray]]:
         input_dtype = img.dtype
 
-        img = tf.image.resize(img, self.wanted_size, self.method, self.preserve_aspect_ratio)
+        img = tf.image.resize(img, self.wanted_size, self.method, self.preserve_aspect_ratio, self.antialias)
         # It will produce an un-padded resized image, with a side shorter than wanted if we preserve aspect ratio
         raw_shape = img.shape[:2]
         if self.preserve_aspect_ratio:
@@ -156,6 +159,7 @@ class Normalize(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         mean: average value per channel
         std: standard deviation per channel
     """
@@ -182,6 +186,7 @@ class LambdaTransformation(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         fn: the function to be applied to the input tensor
     """
 
@@ -219,6 +224,7 @@ class RandomBrightness(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         max_delta: offset to add to each pixel is randomly picked in [-max_delta, max_delta]
         p: probability to apply transformation
     """
@@ -243,6 +249,7 @@ class RandomContrast(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         delta: multiplicative factor is picked in [1-delta, 1+delta] (reduce contrast if factor<1)
     """
 
@@ -266,6 +273,7 @@ class RandomSaturation(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         delta: multiplicative factor is picked in [1-delta, 1+delta] (reduce saturation if factor<1)
     """
 
@@ -288,6 +296,7 @@ class RandomHue(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         max_delta: offset to add to each pixel is randomly picked in [-max_delta, max_delta]
     """
 
@@ -310,6 +319,7 @@ class RandomGamma(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[8, 64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         min_gamma: non-negative real number, lower bound for gamma param
         max_gamma: non-negative real number, upper bound for gamma
         min_gain: lower bound for constant multiplier
@@ -347,6 +357,7 @@ class RandomJpegQuality(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         min_quality: int between [0, 100]
         max_quality: int between [0, 100]
     """
@@ -371,6 +382,7 @@ class GaussianBlur(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         kernel_shape: size of the blurring kernel
         std: min and max value of the standard deviation
     """
@@ -414,6 +426,7 @@ class GaussianNoise(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         mean : mean of the gaussian distribution
         std : std of the gaussian distribution
     """
@@ -451,6 +464,7 @@ class RandomHorizontalFlip(NestedObject):
     >>> out = transfo(image, target)
 
     Args:
+    ----
         p : probability of Horizontal Flip
     """
 
@@ -459,13 +473,6 @@ class RandomHorizontalFlip(NestedObject):
         self.p = p
 
     def __call__(self, img: Union[tf.Tensor, np.ndarray], target: Dict[str, Any]) -> Tuple[tf.Tensor, Dict[str, Any]]:
-        """
-        Args:
-            img: Image to be flipped.
-            target: Dictionary with boxes (in relative coordinates of shape (N, 4)) and labels as keys
-        Returns:
-            Tuple of numpy nd-array or Tensor and target
-        """
         if np.random.rand(1) <= self.p:
             _img = tf.image.flip_left_right(img)
             _target = target.copy()
@@ -484,6 +491,7 @@ class RandomShadow(NestedObject):
     >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
 
     Args:
+    ----
         opacity_range : minimum and maximum opacity of the shade
     """
 

doctr/utils/common_types.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.

doctr/utils/data.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -56,6 +56,7 @@ def download_from_url(
     >>> download_from_url("https://yoursource.com/yourcheckpoint-yourhash.zip")
 
     Args:
+    ----
         url: the URL of the file to download
         file_name: optional name of the file once downloaded
         hash_prefix: optional expected SHA256 hash of the file
@@ -63,12 +64,13 @@ def download_from_url(
         cache_subdir: subfolder to use in the cache
 
     Returns:
+    -------
         the location of the downloaded file
 
     Note:
+    ----
         You can change cache directory location by using `DOCTR_CACHE_DIR` environment variable.
     """
-
     if not isinstance(file_name, str):
         file_name = url.rpartition("/")[-1].split("&")[0]
 

doctr/utils/fonts.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -16,13 +16,14 @@ def get_font(font_family: Optional[str] = None, font_size: int = 13) -> ImageFon
     """Resolves a compatible ImageFont for the system
 
     Args:
+    ----
         font_family: the font family to use
         font_size: the size of the font upon rendering
 
     Returns:
+    -------
         the Pillow font
     """
-
     # Font selection
     if font_family is None:
         try: