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

doctr/models/classification/vit/tensorflow.py

@@ -1,10 +1,10 @@
-# 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 Any, Dict, Optional, Tuple
+from typing import Any
 
 import tensorflow as tf
 from tensorflow.keras import Sequential, layers
@@ -19,7 +19,7 @@ from ...utils import _build_model, load_pretrained_params
 __all__ = ["vit_s", "vit_b"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "vit_s": {
         "mean": (0.694, 0.695, 0.693),
         "std": (0.299, 0.296, 0.301),
@@ -41,7 +41,6 @@ class ClassifierHead(layers.Layer, NestedObject):
     """Classifier head for Vision Transformer
 
     Args:
-    ----
         num_classes: number of output classes
     """
 
@@ -61,7 +60,6 @@ class VisionTransformer(Sequential):
     <https://arxiv.org/pdf/2010.11929.pdf>`_.
 
     Args:
-    ----
         d_model: dimension of the transformer layers
         num_layers: number of transformer layers
         num_heads: number of attention heads
@@ -79,12 +77,12 @@ class VisionTransformer(Sequential):
         num_layers: int,
         num_heads: int,
         ffd_ratio: int,
-        patch_size: Tuple[int, int] = (4, 4),
-        input_shape: Tuple[int, int, int] = (32, 32, 3),
+        patch_size: tuple[int, int] = (4, 4),
+        input_shape: tuple[int, int, int] = (32, 32, 3),
         dropout: float = 0.0,
         num_classes: int = 1000,
         include_top: bool = True,
-        cfg: Optional[Dict[str, Any]] = None,
+        cfg: dict[str, Any] | None = None,
     ) -> None:
         _layers = [
             PatchEmbedding(input_shape, d_model, patch_size),
@@ -103,6 +101,15 @@ class VisionTransformer(Sequential):
         super().__init__(_layers)
         self.cfg = cfg
 
+    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:
+        """Load pretrained parameters onto the model
+
+        Args:
+            path_or_url: the path or URL to the model parameters (checkpoint)
+            **kwargs: additional arguments to be passed to `doctr.models.utils.load_pretrained_params`
+        """
+        load_pretrained_params(self, path_or_url, **kwargs)
+
 
 def _vit(
     arch: str,
@@ -148,12 +155,10 @@ def vit_s(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
     >>> out = model(input_tensor)
 
     Args:
-    ----
         pretrained: boolean, True if model is pretrained
         **kwargs: keyword arguments of the VisionTransformer architecture
 
     Returns:
-    -------
         A feature extractor model
     """
     return _vit(
@@ -179,12 +184,10 @@ def vit_b(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
     >>> out = model(input_tensor)
 
     Args:
-    ----
         pretrained: boolean, True if model is pretrained
         **kwargs: keyword arguments of the VisionTransformer architecture
 
     Returns:
-    -------
         A feature extractor model
     """
     return _vit(

doctr/models/classification/zoo.py

@@ -1,11 +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.
 
-from typing import Any, List
+from typing import Any
 
-from doctr.file_utils import is_tf_available
+from doctr.file_utils import is_tf_available, is_torch_available
 
 from .. import classification
 from ..preprocessor import PreProcessor
@@ -13,7 +13,7 @@ from .predictor import OrientationPredictor
 
 __all__ = ["crop_orientation_predictor", "page_orientation_predictor"]
 
-ARCHS: List[str] = [
+ARCHS: list[str] = [
     "magc_resnet31",
     "mobilenet_v3_small",
     "mobilenet_v3_small_r",
@@ -31,7 +31,11 @@ ARCHS: List[str] = [
     "vit_s",
     "vit_b",
 ]
-ORIENTATION_ARCHS: List[str] = ["mobilenet_v3_small_crop_orientation", "mobilenet_v3_small_page_orientation"]
+
+if is_torch_available():
+    ARCHS.extend(["vip_tiny", "vip_base"])
+
+ORIENTATION_ARCHS: list[str] = ["mobilenet_v3_small_crop_orientation", "mobilenet_v3_small_page_orientation"]
 
 
 def _orientation_predictor(
@@ -48,7 +52,14 @@ def _orientation_predictor(
         # Load directly classifier from backbone
         _model = classification.__dict__[arch](pretrained=pretrained)
     else:
-        if not isinstance(arch, classification.MobileNetV3):
+        allowed_archs = [classification.MobileNetV3]
+        if is_torch_available():
+            # Adding the type for torch compiled models to the allowed architectures
+            from doctr.models.utils import _CompiledModule
+
+            allowed_archs.append(_CompiledModule)
+
+        if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")
         _model = arch
 
@@ -63,7 +74,7 @@ def _orientation_predictor(
 
 
 def crop_orientation_predictor(
-    arch: Any = "mobilenet_v3_small_crop_orientation", pretrained: bool = False, **kwargs: Any
+    arch: Any = "mobilenet_v3_small_crop_orientation", pretrained: bool = False, batch_size: int = 128, **kwargs: Any
 ) -> OrientationPredictor:
     """Crop orientation classification architecture.
 
@@ -74,20 +85,19 @@ def crop_orientation_predictor(
     >>> out = model([input_crop])
 
     Args:
-    ----
         arch: name of the architecture to use (e.g. 'mobilenet_v3_small_crop_orientation')
         pretrained: If True, returns a model pre-trained on our recognition crops dataset
+        batch_size: number of samples the model processes in parallel
         **kwargs: keyword arguments to be passed to the OrientationPredictor
 
     Returns:
-    -------
         OrientationPredictor
     """
-    return _orientation_predictor(arch, pretrained, model_type="crop", **kwargs)
+    return _orientation_predictor(arch=arch, pretrained=pretrained, batch_size=batch_size, model_type="crop", **kwargs)
 
 
 def page_orientation_predictor(
-    arch: Any = "mobilenet_v3_small_page_orientation", pretrained: bool = False, **kwargs: Any
+    arch: Any = "mobilenet_v3_small_page_orientation", pretrained: bool = False, batch_size: int = 4, **kwargs: Any
 ) -> OrientationPredictor:
     """Page orientation classification architecture.
 
@@ -98,13 +108,12 @@ def page_orientation_predictor(
     >>> out = model([input_page])
 
     Args:
-    ----
         arch: name of the architecture to use (e.g. 'mobilenet_v3_small_page_orientation')
         pretrained: If True, returns a model pre-trained on our recognition crops dataset
+        batch_size: number of samples the model processes in parallel
         **kwargs: keyword arguments to be passed to the OrientationPredictor
 
     Returns:
-    -------
         OrientationPredictor
     """
-    return _orientation_predictor(arch, pretrained, model_type="page", **kwargs)
+    return _orientation_predictor(arch=arch, pretrained=pretrained, batch_size=batch_size, model_type="page", **kwargs)

doctr/models/core.py

@@ -1,10 +1,10 @@
-# 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 Any, Dict, Optional
+from typing import Any
 
 from doctr.utils.repr import NestedObject
 
@@ -14,6 +14,6 @@ __all__ = ["BaseModel"]
 class BaseModel(NestedObject):
     """Implements abstract DetectionModel class"""
 
-    def __init__(self, cfg: Optional[Dict[str, Any]] = None) -> None:
+    def __init__(self, cfg: dict[str, Any] | None = None) -> None:
         super().__init__()
         self.cfg = cfg

doctr/models/detection/_utils/__init__.py

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

doctr/models/detection/_utils/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 Dict, List
 
 import numpy as np
 
@@ -11,16 +10,15 @@ __all__ = ["_remove_padding"]
 
 
 def _remove_padding(
-    pages: List[np.ndarray],
-    loc_preds: List[Dict[str, np.ndarray]],
+    pages: list[np.ndarray],
+    loc_preds: list[dict[str, np.ndarray]],
     preserve_aspect_ratio: bool,
     symmetric_pad: bool,
     assume_straight_pages: bool,
-) -> List[Dict[str, np.ndarray]]:
+) -> list[dict[str, np.ndarray]]:
     """Remove padding from the localization predictions
 
     Args:
-    ----
         pages: list of pages
         loc_preds: list of localization predictions
         preserve_aspect_ratio: whether the aspect ratio was preserved during padding
@@ -28,7 +26,6 @@ def _remove_padding(
         assume_straight_pages: whether the pages are assumed to be straight
 
     Returns:
-    -------
         list of unpaded localization predictions
     """
     if preserve_aspect_ratio:

doctr/models/detection/_utils/pytorch.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.
@@ -13,12 +13,10 @@ def erode(x: Tensor, kernel_size: int) -> Tensor:
     """Performs erosion on a given tensor
 
     Args:
-    ----
         x: boolean tensor of shape (N, C, H, W)
         kernel_size: the size of the kernel to use for erosion
 
     Returns:
-    -------
         the eroded tensor
     """
     _pad = (kernel_size - 1) // 2
@@ -30,12 +28,10 @@ def dilate(x: Tensor, kernel_size: int) -> Tensor:
     """Performs dilation on a given tensor
 
     Args:
-    ----
         x: boolean tensor of shape (N, C, H, W)
         kernel_size: the size of the kernel to use for dilation
 
     Returns:
-    -------
         the dilated tensor
     """
     _pad = (kernel_size - 1) // 2

doctr/models/detection/_utils/tensorflow.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.
@@ -12,12 +12,10 @@ def erode(x: tf.Tensor, kernel_size: int) -> tf.Tensor:
     """Performs erosion on a given tensor
 
     Args:
-    ----
         x: boolean tensor of shape (N, H, W, C)
         kernel_size: the size of the kernel to use for erosion
 
     Returns:
-    -------
         the eroded tensor
     """
     return 1 - tf.nn.max_pool2d(1 - x, kernel_size, strides=1, padding="SAME")
@@ -27,12 +25,10 @@ def dilate(x: tf.Tensor, kernel_size: int) -> tf.Tensor:
     """Performs dilation on a given tensor
 
     Args:
-    ----
         x: boolean tensor of shape (N, H, W, C)
         kernel_size: the size of the kernel to use for dilation
 
     Returns:
-    -------
         the dilated tensor
     """
     return tf.nn.max_pool2d(x, kernel_size, strides=1, padding="SAME")

doctr/models/detection/core.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 List
 
 import cv2
 import numpy as np
@@ -17,7 +16,6 @@ class DetectionPostProcessor(NestedObject):
     """Abstract class to postprocess the raw output of the model
 
     Args:
-    ----
         box_thresh (float): minimal objectness score to consider a box
         bin_thresh (float): threshold to apply to segmentation raw heatmap
         assume straight_pages (bool): if True, fit straight boxes only
@@ -37,13 +35,11 @@ class DetectionPostProcessor(NestedObject):
         """Compute the confidence score for a polygon : mean of the p values on the polygon
 
         Args:
-        ----
             pred (np.ndarray): p map returned by the model
             points: coordinates of the polygon
             assume_straight_pages: if True, fit straight boxes only
 
         Returns:
-        -------
             polygon objectness
         """
         h, w = pred.shape[:2]
@@ -71,15 +67,13 @@ class DetectionPostProcessor(NestedObject):
     def __call__(
         self,
         proba_map,
-    ) -> List[List[np.ndarray]]:
+    ) -> list[list[np.ndarray]]:
         """Performs postprocessing for a list of model outputs
 
         Args:
-        ----
             proba_map: probability map of shape (N, H, W, C)
 
         Returns:
-        -------
             list of N class predictions (for each input sample), where each class predictions is a list of C tensors
         of shape (*, 5) or (*, 6)
         """

doctr/models/detection/differentiable_binarization/__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():
-    from .pytorch import *  # type: ignore[assignment]
+if is_torch_available():
+    from .pytorch import *
+elif is_tf_available():
+    from .tensorflow import *  # type: ignore[assignment]

doctr/models/detection/differentiable_binarization/base.py

@@ -1,11 +1,10 @@
-# 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.
 
 # Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
 
-from typing import Dict, List, Tuple, Union
 
 import cv2
 import numpy as np
@@ -22,7 +21,6 @@ class DBPostProcessor(DetectionPostProcessor):
     <https://github.com/xuannianz/DifferentiableBinarization>`_.
 
     Args:
-    ----
         unclip ratio: ratio used to unshrink polygons
         min_size_box: minimal length (pix) to keep a box
         max_candidates: maximum boxes to consider in a single page
@@ -47,11 +45,9 @@ class DBPostProcessor(DetectionPostProcessor):
         """Expand a polygon (points) by a factor unclip_ratio, and returns a polygon
 
         Args:
-        ----
             points: The first parameter.
 
         Returns:
-        -------
             a box in absolute coordinates (xmin, ymin, xmax, ymax) or (4, 2) array (quadrangle)
         """
         if not self.assume_straight_pages:
@@ -62,9 +58,8 @@ class DBPostProcessor(DetectionPostProcessor):
             area = (rect[1][0] + 1) * (1 + rect[1][1])
             length = 2 * (rect[1][0] + rect[1][1]) + 2
         else:
-            poly = Polygon(points)
-            area = poly.area
-            length = poly.length
+            area = cv2.contourArea(points)
+            length = cv2.arcLength(points, closed=True)
         distance = area * self.unclip_ratio / length  # compute distance to expand polygon
         offset = pyclipper.PyclipperOffset()
         offset.AddPath(points, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
@@ -96,25 +91,23 @@ class DBPostProcessor(DetectionPostProcessor):
         """Compute boxes from a bitmap/pred_map: find connected components then filter boxes
 
         Args:
-        ----
             pred: Pred map from differentiable binarization output
             bitmap: Bitmap map computed from pred (binarized)
             angle_tol: Comparison tolerance of the angle with the median angle across the page
             ratio_tol: Under this limit aspect ratio, we cannot resolve the direction of the crop
 
         Returns:
-        -------
             np tensor boxes for the bitmap, each box is a 5-element list
                 containing x, y, w, h, score for the box
         """
         height, width = bitmap.shape[:2]
         min_size_box = 2
-        boxes: List[Union[np.ndarray, List[float]]] = []
+        boxes: list[np.ndarray | list[float]] = []
         # get contours from connected components on the bitmap
         contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         for contour in contours:
             # Check whether smallest enclosing bounding box is not too small
-            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < min_size_box):  # type: ignore[index]
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < min_size_box):
                 continue
             # Compute objectness
             if self.assume_straight_pages:
@@ -164,7 +157,6 @@ class _DBNet:
     <https://arxiv.org/pdf/1911.08947.pdf>`_.
 
     Args:
-    ----
         feature extractor: the backbone serving as feature extractor
         fpn_channels: number of channels each extracted feature maps is mapped to
     """
@@ -186,7 +178,6 @@ class _DBNet:
         """Compute the distance for each point of the map (xs, ys) to the (a, b) segment
 
         Args:
-        ----
             xs : map of x coordinates (height, width)
             ys : map of y coordinates (height, width)
             a: first point defining the [ab] segment
@@ -194,7 +185,6 @@ class _DBNet:
             eps: epsilon to avoid division by zero
 
         Returns:
-        -------
             The computed distance
 
         """
@@ -214,11 +204,10 @@ class _DBNet:
         polygon: np.ndarray,
         canvas: np.ndarray,
         mask: np.ndarray,
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-        """Draw a polygon treshold map on a canvas, as described in the DB paper
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Draw a polygon threshold map on a canvas, as described in the DB paper
 
         Args:
-        ----
             polygon : array of coord., to draw the boundary of the polygon
             canvas : threshold map to fill with polygons
             mask : mask for training on threshold polygons
@@ -278,10 +267,10 @@ class _DBNet:
 
     def build_target(
         self,
-        target: List[Dict[str, np.ndarray]],
-        output_shape: Tuple[int, int, int],
+        target: list[dict[str, np.ndarray]],
+        output_shape: tuple[int, int, int],
         channels_last: bool = True,
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         if any(t.dtype != np.float32 for tgt in target for t in tgt.values()):
             raise AssertionError("the expected dtype of target 'boxes' entry is 'np.float32'.")
         if any(np.any((t[:, :4] > 1) | (t[:, :4] < 0)) for tgt in target for t in tgt.values()):