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

doctr/models/utils/pytorch.py

@@ -7,6 +7,7 @@ import logging
 from typing import Any
 
 import torch
+import validators
 from torch import nn
 
 from doctr.utils.data import download_from_url
@@ -36,7 +37,7 @@ def _bf16_to_float32(x: torch.Tensor) -> torch.Tensor:
 
 def load_pretrained_params(
     model: nn.Module,
-    url: str | None = None,
+    path_or_url: str | None = None,
     hash_prefix: str | None = None,
     ignore_keys: list[str] | None = None,
     **kwargs: Any,
@@ -44,33 +45,42 @@ def load_pretrained_params(
     """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.pt")
 
     Args:
         model: the PyTorch model to be loaded
-        url: URL of the zipped set of parameters
+        path_or_url: the path or URL to the model parameters (checkpoint)
         hash_prefix: first characters of SHA256 expected hash
         ignore_keys: list of weights to be ignored from the state_dict
         **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)
+    if path_or_url is None:
+        logging.warning("No model URL or Path provided, using default initialization.")
+        return
+
+    archive_path = (
+        download_from_url(path_or_url, hash_prefix=hash_prefix, cache_subdir="models", **kwargs)
+        if validators.url(path_or_url)
+        else path_or_url
+    )
 
-        # Read state_dict
-        state_dict = torch.load(archive_path, map_location="cpu")
+    # Read state_dict
+    state_dict = torch.load(archive_path, map_location="cpu")
 
-        # Remove weights from the state_dict
-        if ignore_keys is not None and len(ignore_keys) > 0:
-            for key in ignore_keys:
+    # Remove weights from the state_dict
+    if ignore_keys is not None and len(ignore_keys) > 0:
+        for key in ignore_keys:
+            if key in state_dict:
                 state_dict.pop(key)
-            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
-            if set(missing_keys) != set(ignore_keys) or len(unexpected_keys) > 0:
-                raise ValueError("unable to load state_dict, due to non-matching keys.")
-        else:
-            # Load weights
-            model.load_state_dict(state_dict)
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+        if any(k not in ignore_keys for k in missing_keys + unexpected_keys):
+            raise ValueError(
+                "Unable to load state_dict, due to non-matching keys.\n"
+                + f"Unexpected keys: {unexpected_keys}\nMissing keys: {missing_keys}"
+            )
+    else:
+        # Load weights
+        model.load_state_dict(state_dict)
 
 
 def conv_sequence_pt(
@@ -154,7 +164,7 @@ def export_model_to_onnx(model: nn.Module, model_name: str, dummy_input: torch.T
     """
     torch.onnx.export(
         model,
-        dummy_input,
+        dummy_input,  # type: ignore[arg-type]
         f"{model_name}.onnx",
         input_names=["input"],
         output_names=["logits"],

doctr/transforms/functional/__init__.py

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

doctr/transforms/functional/pytorch.py

@@ -33,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)  # type: ignore[attr-defined]
+        out = (out.to(dtype=rgb_shift.dtype) * rgb_shift).to(dtype=torch.uint8)
     else:
-        out = out * rgb_shift.to(dtype=out.dtype)  # type: ignore[attr-defined]
+        out = out * rgb_shift.to(dtype=out.dtype)
     # Inverse the color
     out = 255 - out if out.dtype == torch.uint8 else 1 - out
     return out
@@ -77,7 +77,7 @@ def rotate_sample(
     rotated_geoms: np.ndarray = rotate_abs_geoms(
         _geoms,
         angle,
-        img.shape[1:],
+        img.shape[1:],  # type: ignore[arg-type]
         expand,
     ).astype(np.float32)
 
@@ -124,7 +124,7 @@ def random_shadow(img: torch.Tensor, opacity_range: tuple[float, float], **kwarg
     Returns:
         Shadowed image as a PyTorch tensor (same shape as input).
     """
-    shadow_mask = create_shadow_mask(img.shape[1:], **kwargs)
+    shadow_mask = create_shadow_mask(img.shape[1:], **kwargs)  # type: ignore[arg-type]
     opacity = np.random.uniform(*opacity_range)
 
     # Apply Gaussian blur to the shadow mask

doctr/transforms/modules/__init__.py

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

doctr/transforms/modules/base.py

@@ -20,27 +20,13 @@ __all__ = ["SampleCompose", "ImageTransform", "ColorInversion", "OneOf", "Random
 class SampleCompose(NestedObject):
     """Implements a wrapper that will apply transformations sequentially on both image and target
 
-    .. tabs::
+    .. code:: python
 
-        .. tab:: PyTorch
-
-            .. code:: python
-
-                >>> import numpy as np
-                >>> import torch
-                >>> 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)))
-
-        .. tab:: TensorFlow
-
-            .. code:: python
-
-                >>> import numpy as np
-                >>> import tensorflow as tf
-                >>> 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)))
+        >>> import numpy as np
+        >>> import torch
+        >>> 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)))
 
     Args:
         transforms: list of transformation modules
@@ -61,25 +47,12 @@ class SampleCompose(NestedObject):
 class ImageTransform(NestedObject):
     """Implements a transform wrapper to turn an image-only transformation into an image+target transform
 
-    .. tabs::
-
-        .. tab:: PyTorch
-
-            .. code:: python
+    .. code:: python
 
-                >>> import torch
-                >>> from doctr.transforms import ImageTransform, ColorInversion
-                >>> transfo = ImageTransform(ColorInversion((32, 32)))
-                >>> out, _ = transfo(torch.rand(8, 64, 64, 3), None)
-
-        .. tab:: TensorFlow
-
-            .. code:: python
-
-                >>> import tensorflow as tf
-                >>> 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)
+        >>> import torch
+        >>> from doctr.transforms import ImageTransform, ColorInversion
+        >>> transfo = ImageTransform(ColorInversion((32, 32)))
+        >>> out, _ = transfo(torch.rand(8, 64, 64, 3), None)
 
     Args:
         transform: the image transformation module to wrap
@@ -99,25 +72,12 @@ class ColorInversion(NestedObject):
     """Applies the following tranformation to a tensor (image or batch of images):
     convert to grayscale, colorize (shift 0-values randomly), and then invert colors
 
-    .. tabs::
-
-        .. tab:: PyTorch
-
-            .. code:: python
-
-                >>> import torch
-                >>> from doctr.transforms import ColorInversion
-                >>> transfo = ColorInversion(min_val=0.6)
-                >>> out = transfo(torch.rand(8, 64, 64, 3))
+    .. code:: python
 
-        .. tab:: TensorFlow
-
-            .. code:: python
-
-                >>> import tensorflow as tf
-                >>> 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))
+        >>> import torch
+        >>> from doctr.transforms import ColorInversion
+        >>> transfo = ColorInversion(min_val=0.6)
+        >>> out = transfo(torch.rand(8, 64, 64, 3))
 
     Args:
         min_val: range [min_val, 1] to colorize RGB pixels
@@ -136,25 +96,12 @@ class ColorInversion(NestedObject):
 class OneOf(NestedObject):
     """Randomly apply one of the input transformations
 
-    .. tabs::
-
-        .. tab:: PyTorch
-
-            .. code:: python
-
-                >>> import torch
-                >>> from doctr.transforms import OneOf
-                >>> transfo = OneOf([JpegQuality(), Gamma()])
-                >>> out = transfo(torch.rand(1, 64, 64, 3))
-
-        .. tab:: TensorFlow
+    .. code:: python
 
-            .. code:: python
-
-                >>> import tensorflow as tf
-                >>> from doctr.transforms import OneOf
-                >>> transfo = OneOf([JpegQuality(), Gamma()])
-                >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
+        >>> import torch
+        >>> from doctr.transforms import OneOf
+        >>> transfo = OneOf([JpegQuality(), Gamma()])
+        >>> out = transfo(torch.rand(1, 64, 64, 3))
 
     Args:
         transforms: list of transformations, one only will be picked
@@ -175,25 +122,12 @@ class OneOf(NestedObject):
 class RandomApply(NestedObject):
     """Apply with a probability p the input transformation
 
-    .. tabs::
-
-        .. tab:: PyTorch
-
-            .. code:: python
-
-                >>> import torch
-                >>> from doctr.transforms import RandomApply
-                >>> transfo = RandomApply(Gamma(), p=.5)
-                >>> out = transfo(torch.rand(1, 64, 64, 3))
-
-        .. tab:: TensorFlow
-
-            .. code:: python
+    .. code:: python
 
-                >>> import tensorflow as tf
-                >>> from doctr.transforms import RandomApply
-                >>> transfo = RandomApply(Gamma(), p=.5)
-                >>> out = transfo(tf.random.uniform(shape=[64, 64, 3], minval=0, maxval=1))
+        >>> import torch
+        >>> from doctr.transforms import RandomApply
+        >>> transfo = RandomApply(Gamma(), p=.5)
+        >>> out = transfo(torch.rand(1, 64, 64, 3))
 
     Args:
         transform: transformation to apply

doctr/transforms/modules/pytorch.py

@@ -13,7 +13,7 @@ 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
+from ..functional import random_shadow
 
 __all__ = [
     "Resize",
@@ -27,7 +27,21 @@ __all__ = [
 
 
 class Resize(T.Resize):
-    """Resize the input image to the given size"""
+    """Resize the input image to the given size
+
+    >>> import torch
+    >>> from doctr.transforms import Resize
+    >>> transfo = Resize((64, 64), preserve_aspect_ratio=True, symmetric_pad=True)
+    >>> out = transfo(torch.rand((3, 64, 64)))
+
+    Args:
+        size: output size in pixels, either a tuple (height, width) or a single integer for square images
+        interpolation: interpolation mode to use for resizing, default is bilinear
+        preserve_aspect_ratio: whether to preserve the aspect ratio of the image,
+            if True, the image will be resized to fit within the target size while maintaining its aspect ratio
+        symmetric_pad: whether to symmetrically pad the image to the target size,
+            if True, the image will be padded equally on both sides to fit the target size
+    """
 
     def __init__(
         self,
@@ -36,25 +50,19 @@ class Resize(T.Resize):
         preserve_aspect_ratio: bool = False,
         symmetric_pad: bool = False,
     ) -> None:
-        super().__init__(size, interpolation, antialias=True)
+        super().__init__(size if isinstance(size, (list, tuple)) else (size, size), interpolation, antialias=True)
         self.preserve_aspect_ratio = preserve_aspect_ratio
         self.symmetric_pad = symmetric_pad
 
-        if not isinstance(self.size, (int, tuple, list)):
-            raise AssertionError("size should be either a tuple, a list or an int")
-
     def forward(
         self,
         img: torch.Tensor,
         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:
-            target_ratio = self.size[0] / self.size[1]
+        target_ratio = self.size[0] / self.size[1]
         actual_ratio = img.shape[-2] / img.shape[-1]
 
-        if not self.preserve_aspect_ratio or (target_ratio == actual_ratio and (isinstance(self.size, (tuple, list)))):
+        if not self.preserve_aspect_ratio or (target_ratio == actual_ratio):
             # If we don't preserve the aspect ratio or the wanted aspect ratio is the same than the original one
             # We can use with the regular resize
             if target is not None:
@@ -62,16 +70,10 @@ class Resize(T.Resize):
             return super().forward(img)
         else:
             # Resize
-            if isinstance(self.size, (tuple, list)):
-                if actual_ratio > target_ratio:
-                    tmp_size = (self.size[0], max(int(self.size[0] / actual_ratio), 1))
-                else:
-                    tmp_size = (max(int(self.size[1] * actual_ratio), 1), self.size[1])
-            elif isinstance(self.size, int):  # self.size is the longest side, infer the other
-                if img.shape[-2] <= img.shape[-1]:
-                    tmp_size = (max(int(self.size * actual_ratio), 1), self.size)
-                else:
-                    tmp_size = (self.size, max(int(self.size / actual_ratio), 1))
+            if actual_ratio > target_ratio:
+                tmp_size = (self.size[0], max(int(self.size[0] / actual_ratio), 1))
+            else:
+                tmp_size = (max(int(self.size[1] * actual_ratio), 1), self.size[1])
 
             # Scale image
             img = F.resize(img, tmp_size, self.interpolation, antialias=True)
@@ -93,14 +95,14 @@ class Resize(T.Resize):
                 if self.preserve_aspect_ratio:
                     # Get absolute coords
                     if target.shape[1:] == (4,):
-                        if isinstance(self.size, (tuple, list)) and self.symmetric_pad:
+                        if self.symmetric_pad:
                             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:
                             target[:, [0, 2]] *= raw_shape[-1] / img.shape[-1]
                             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 self.symmetric_pad:
                             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:
@@ -143,9 +145,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)  # type: ignore[attr-defined]
+            return (x + 255 * noise).round().clamp(0, 255).to(dtype=torch.uint8)
         else:
-            return (x + noise.to(dtype=x.dtype)).clamp(0, 1)  # type: ignore[attr-defined]
+            return (x + noise.to(dtype=x.dtype)).clamp(0, 1)
 
     def extra_repr(self) -> str:
         return f"mean={self.mean}, std={self.std}"
@@ -233,7 +235,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,

doctr/utils/data.py

@@ -92,7 +92,7 @@ def download_from_url(
         # Create folder hierarchy
         folder_path.mkdir(parents=True, exist_ok=True)
     except OSError:
-        error_message = f"Failed creating cache direcotry at {folder_path}"
+        error_message = f"Failed creating cache directory at {folder_path}"
         if os.environ.get("DOCTR_CACHE_DIR", ""):
             error_message += " using path from 'DOCTR_CACHE_DIR' environment variable."
         else:

doctr/utils/geometry.py

@@ -300,7 +300,7 @@ def rotate_image(
     # 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]))),
@@ -390,14 +390,13 @@ 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) -> 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
@@ -409,21 +408,19 @@ def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True
 
     # Project relative coordinates
     _boxes = boxes.copy()
-    h, w = img.shape[:2] if channels_last else img.shape[-2:]
+    h, w = img.shape[:2]
     if not np.issubdtype(_boxes.dtype, np.integer):
         _boxes[:, [0, 2]] *= w
         _boxes[:, [1, 3]] *= h
         _boxes = _boxes.round().astype(int)
         # Add last index
         _boxes[2:] += 1
-    if channels_last:
-        return deepcopy([img[box[1] : box[3], box[0] : box[2]] for box in _boxes])
 
-    return deepcopy([img[:, box[1] : box[3], box[0] : box[2]] for box in _boxes])
+    return deepcopy([img[box[1] : box[3], box[0] : box[2]] for box in _boxes])
 
 
 def extract_rcrops(
-    img: np.ndarray, polys: np.ndarray, dtype=np.float32, channels_last: bool = True, assume_horizontal: bool = False
+    img: np.ndarray, polys: np.ndarray, dtype=np.float32, assume_horizontal: bool = False
 ) -> list[np.ndarray]:
     """Created cropped images from list of rotated bounding boxes
 
@@ -431,7 +428,6 @@ def extract_rcrops(
         img: input image
         polys: bounding boxes of shape (N, 4, 2)
         dtype: target data type of bounding boxes
-        channels_last: whether the channel dimensions is the last one instead of the last one
         assume_horizontal: whether the boxes are assumed to be only horizontally oriented
 
     Returns:
@@ -444,12 +440,12 @@ def extract_rcrops(
 
     # Project relative coordinates
     _boxes = polys.copy()
-    height, width = img.shape[:2] if channels_last else img.shape[-2:]
+    height, width = img.shape[:2]
     if not np.issubdtype(_boxes.dtype, np.integer):
         _boxes[:, :, 0] *= width
         _boxes[:, :, 1] *= height
 
-    src_img = img if channels_last else img.transpose(1, 2, 0)
+    src_img = img
 
     # Handle only horizontal oriented boxes
     if assume_horizontal:

doctr/utils/visualization.py

@@ -148,7 +148,7 @@ def get_colors(num_colors: int) -> list[tuple[float, float, float]]:
         hue = i / 360.0
         lightness = (50 + np.random.rand() * 10) / 100.0
         saturation = (90 + np.random.rand() * 10) / 100.0
-        colors.append(colorsys.hls_to_rgb(hue, lightness, saturation))
+        colors.append(colorsys.hls_to_rgb(hue, lightness, saturation))  # type: ignore[arg-type]
     return colors
 
 

doctr/version.py

@@ -1 +1 @@
-__version__ = 'v0.11.0'
+__version__ = 'v1.0.0'