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

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,7 @@ ARCHS: List[str] = [
     "vit_s",
     "vit_b",
 ]
-ORIENTATION_ARCHS: List[str] = ["mobilenet_v3_small_crop_orientation", "mobilenet_v3_small_page_orientation"]
+ORIENTATION_ARCHS: list[str] = ["mobilenet_v3_small_crop_orientation", "mobilenet_v3_small_page_orientation"]
 
 
 def _orientation_predictor(
@@ -48,7 +48,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 +70,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 +81,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 +104,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:
@@ -96,25 +92,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 +158,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 +179,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 +186,6 @@ class _DBNet:
             eps: epsilon to avoid division by zero
 
         Returns:
-        -------
             The computed distance
 
         """
@@ -214,11 +205,10 @@ class _DBNet:
         polygon: np.ndarray,
         canvas: np.ndarray,
         mask: np.ndarray,
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Draw a polygon treshold 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 +268,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()):

doctr/models/detection/differentiable_binarization/pytorch.py

@@ -1,9 +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, Callable, Dict, List, Optional
+from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 import torch
@@ -22,7 +23,7 @@ from .base import DBPostProcessor, _DBNet
 __all__ = ["DBNet", "db_resnet50", "db_resnet34", "db_mobilenet_v3_large"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "db_resnet50": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
@@ -47,7 +48,7 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 class FeaturePyramidNetwork(nn.Module):
     def __init__(
         self,
-        in_channels: List[int],
+        in_channels: list[int],
         out_channels: int,
         deform_conv: bool = False,
     ) -> None:
@@ -76,12 +77,12 @@ class FeaturePyramidNetwork(nn.Module):
             for idx, chans in enumerate(in_channels)
         ])
 
-    def forward(self, x: List[torch.Tensor]) -> torch.Tensor:
+    def forward(self, x: list[torch.Tensor]) -> torch.Tensor:
         if len(x) != len(self.out_branches):
             raise AssertionError
         # Conv1x1 to get the same number of channels
-        _x: List[torch.Tensor] = [branch(t) for branch, t in zip(self.in_branches, x)]
-        out: List[torch.Tensor] = [_x[-1]]
+        _x: list[torch.Tensor] = [branch(t) for branch, t in zip(self.in_branches, x)]
+        out: list[torch.Tensor] = [_x[-1]]
         for t in _x[:-1][::-1]:
             out.append(self.upsample(out[-1]) + t)
 
@@ -96,7 +97,6 @@ class DBNet(_DBNet, nn.Module):
     <https://arxiv.org/pdf/1911.08947.pdf>`_.
 
     Args:
-    ----
         feature extractor: the backbone serving as feature extractor
         head_chans: the number of channels in the head
         deform_conv: whether to use deformable convolution
@@ -117,8 +117,8 @@ class DBNet(_DBNet, nn.Module):
         box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
         exportable: bool = False,
-        cfg: Optional[Dict[str, Any]] = None,
-        class_names: List[str] = [CLASS_NAME],
+        cfg: dict[str, Any] | None = None,
+        class_names: list[str] = [CLASS_NAME],
     ) -> None:
         super().__init__()
         self.class_names = class_names
@@ -182,10 +182,10 @@ class DBNet(_DBNet, nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        target: Optional[List[np.ndarray]] = None,
+        target: list[np.ndarray] | None = None,
         return_model_output: bool = False,
         return_preds: bool = False,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         # Extract feature maps at different stages
         feats = self.feat_extractor(x)
         feats = [feats[str(idx)] for idx in range(len(feats))]
@@ -193,7 +193,7 @@ class DBNet(_DBNet, nn.Module):
         feat_concat = self.fpn(feats)
         logits = self.prob_head(feat_concat)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
         if self.exportable:
             out["logits"] = logits
             return out
@@ -205,11 +205,16 @@ class DBNet(_DBNet, nn.Module):
             out["out_map"] = prob_map
 
         if target is None or return_preds:
+            # Disable for torch.compile compatibility
+            @torch.compiler.disable  # type: ignore[attr-defined]
+            def _postprocess(prob_map: torch.Tensor) -> list[dict[str, Any]]:
+                return [
+                    dict(zip(self.class_names, preds))
+                    for preds in self.postprocessor(prob_map.detach().cpu().permute((0, 2, 3, 1)).numpy())
+                ]
+
             # Post-process boxes (keep only text predictions)
-            out["preds"] = [
-                dict(zip(self.class_names, preds))
-                for preds in self.postprocessor(prob_map.detach().cpu().permute((0, 2, 3, 1)).numpy())
-            ]
+            out["preds"] = _postprocess(prob_map)
 
         if target is not None:
             thresh_map = self.thresh_head(feat_concat)
@@ -222,7 +227,7 @@ class DBNet(_DBNet, nn.Module):
         self,
         out_map: torch.Tensor,
         thresh_map: torch.Tensor,
-        target: List[np.ndarray],
+        target: list[np.ndarray],
         gamma: float = 2.0,
         alpha: float = 0.5,
         eps: float = 1e-8,
@@ -231,7 +236,6 @@ class DBNet(_DBNet, nn.Module):
         and a list of masks for each image. From there it computes the loss with the model output
 
         Args:
-        ----
             out_map: output feature map of the model of shape (N, C, H, W)
             thresh_map: threshold map of shape (N, C, H, W)
             target: list of dictionary where each dict has a `boxes` and a `flags` entry
@@ -240,7 +244,6 @@ class DBNet(_DBNet, nn.Module):
             eps: epsilon factor in dice loss
 
         Returns:
-        -------
             A loss tensor
         """
         if gamma < 0:
@@ -273,7 +276,7 @@ class DBNet(_DBNet, nn.Module):
                 dice_map = torch.softmax(out_map, dim=1)
             else:
                 # compute binary map instead
-                dice_map = 1 / (1 + torch.exp(-50.0 * (prob_map - thresh_map)))
+                dice_map = 1 / (1 + torch.exp(-50.0 * (prob_map - thresh_map)))  # type: ignore[assignment]
             # Class reduced
             inter = (seg_mask * dice_map * seg_target).sum((0, 2, 3))
             cardinality = (seg_mask * (dice_map + seg_target)).sum((0, 2, 3))
@@ -290,10 +293,10 @@ def _dbnet(
     arch: str,
     pretrained: bool,
     backbone_fn: Callable[[bool], nn.Module],
-    fpn_layers: List[str],
-    backbone_submodule: Optional[str] = None,
+    fpn_layers: list[str],
+    backbone_submodule: str | None = None,
     pretrained_backbone: bool = True,
-    ignore_keys: Optional[List[str]] = None,
+    ignore_keys: list[str] | None = None,
     **kwargs: Any,
 ) -> DBNet:
     pretrained_backbone = pretrained_backbone and not pretrained
@@ -341,12 +344,10 @@ def db_resnet34(pretrained: bool = False, **kwargs: Any) -> DBNet:
     >>> out = model(input_tensor)
 
     Args:
-    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
         **kwargs: keyword arguments of the DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _dbnet(
@@ -376,12 +377,10 @@ def db_resnet50(pretrained: bool = False, **kwargs: Any) -> DBNet:
     >>> out = model(input_tensor)
 
     Args:
-    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
         **kwargs: keyword arguments of the DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _dbnet(
@@ -411,12 +410,10 @@ def db_mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> DBNet:
     >>> out = model(input_tensor)
 
     Args:
-    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
         **kwargs: keyword arguments of the DBNet architecture
 
     Returns:
-    -------
         text detection architecture
     """
     return _dbnet(