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

doctr/models/detection/fast/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, Union
+from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 import torch
@@ -21,7 +22,7 @@ from .base import _FAST, FASTPostProcessor
 __all__ = ["FAST", "fast_tiny", "fast_small", "fast_base", "reparameterize"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "fast_tiny": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
@@ -47,7 +48,6 @@ class FastNeck(nn.Module):
     """Neck of the FAST architecture, composed of a series of 3x3 convolutions and upsampling layers.
 
     Args:
-    ----
         in_channels: number of input channels
         out_channels: number of output channels
     """
@@ -77,7 +77,6 @@ class FastHead(nn.Sequential):
     """Head of the FAST architecture
 
     Args:
-    ----
         in_channels: number of input channels
         num_classes: number of output classes
         out_channels: number of output channels
@@ -91,7 +90,7 @@ class FastHead(nn.Sequential):
         out_channels: int = 128,
         dropout: float = 0.1,
     ) -> None:
-        _layers: List[nn.Module] = [
+        _layers: list[nn.Module] = [
             FASTConvLayer(in_channels, out_channels, kernel_size=3),
             nn.Dropout(dropout),
             nn.Conv2d(out_channels, num_classes, kernel_size=1, bias=False),
@@ -104,7 +103,6 @@ class FAST(_FAST, nn.Module):
     <https://arxiv.org/pdf/2111.02394.pdf>`_.
 
     Args:
-    ----
         feat extractor: the backbone serving as feature extractor
         bin_thresh: threshold for binarization
         box_thresh: minimal objectness score to consider a box
@@ -125,8 +123,8 @@ class FAST(_FAST, nn.Module):
         pooling_size: int = 4,  # different from paper performs better on close text-rich images
         assume_straight_pages: bool = True,
         exportable: bool = False,
-        cfg: Optional[Dict[str, Any]] = {},
-        class_names: List[str] = [CLASS_NAME],
+        cfg: dict[str, Any] = {},
+        class_names: list[str] = [CLASS_NAME],
     ) -> None:
         super().__init__()
         self.class_names = class_names
@@ -172,13 +170,22 @@ class FAST(_FAST, nn.Module):
                 m.weight.data.fill_(1.0)
                 m.bias.data.zero_()
 
+    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 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))]
@@ -186,7 +193,7 @@ class FAST(_FAST, nn.Module):
         feat_concat = self.neck(feats)
         logits = F.interpolate(self.prob_head(feat_concat), size=x.shape[-2:], mode="bilinear")
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
         if self.exportable:
             out["logits"] = logits
             return out
@@ -198,11 +205,16 @@ class FAST(_FAST, 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:
             loss = self.compute_loss(logits, target)
@@ -213,19 +225,17 @@ class FAST(_FAST, nn.Module):
     def compute_loss(
         self,
         out_map: torch.Tensor,
-        target: List[np.ndarray],
+        target: list[np.ndarray],
         eps: float = 1e-6,
     ) -> torch.Tensor:
         """Compute fast loss, 2 x Dice loss where the text kernel loss is scaled by 0.5.
 
         Args:
-        ----
             out_map: output feature map of the model of shape (N, num_classes, H, W)
             target: list of dictionary where each dict has a `boxes` and a `flags` entry
             eps: epsilon factor in dice loss
 
         Returns:
-        -------
             A loss tensor
         """
         targets = self.build_target(target, out_map.shape[1:], False)  # type: ignore[arg-type]
@@ -279,15 +289,13 @@ class FAST(_FAST, nn.Module):
         return text_loss + kernel_loss
 
 
-def reparameterize(model: Union[FAST, nn.Module]) -> FAST:
+def reparameterize(model: FAST | nn.Module) -> FAST:
     """Fuse batchnorm and conv layers and reparameterize the model
 
-    args:
-    ----
+    Args:
         model: the FAST model to reparameterize
 
     Returns:
-    -------
         the reparameterized model
     """
     last_conv = None
@@ -303,12 +311,12 @@ def reparameterize(model: Union[FAST, nn.Module]) -> FAST:
             if last_conv is None:
                 continue
             conv_w = last_conv.weight
-            conv_b = last_conv.bias if last_conv.bias is not None else torch.zeros_like(child.running_mean)
+            conv_b = last_conv.bias if last_conv.bias is not None else torch.zeros_like(child.running_mean)  # type: ignore[arg-type]
 
-            factor = child.weight / torch.sqrt(child.running_var + child.eps)
+            factor = child.weight / torch.sqrt(child.running_var + child.eps)  # type: ignore
             last_conv.weight = nn.Parameter(conv_w * factor.reshape([last_conv.out_channels, 1, 1, 1]))
             last_conv.bias = nn.Parameter((conv_b - child.running_mean) * factor + child.bias)
-            model._modules[last_conv_name] = last_conv
+            model._modules[last_conv_name] = last_conv  # type: ignore[index]
             model._modules[name] = nn.Identity()
             last_conv = None
         elif isinstance(child, nn.Conv2d):
@@ -324,9 +332,9 @@ def _fast(
     arch: str,
     pretrained: bool,
     backbone_fn: Callable[[bool], nn.Module],
-    feat_layers: List[str],
+    feat_layers: list[str],
     pretrained_backbone: bool = True,
-    ignore_keys: Optional[List[str]] = None,
+    ignore_keys: list[str] | None = None,
     **kwargs: Any,
 ) -> FAST:
     pretrained_backbone = pretrained_backbone and not pretrained
@@ -350,7 +358,7 @@ def _fast(
         _ignore_keys = (
             ignore_keys if kwargs["class_names"] != default_cfgs[arch].get("class_names", [CLASS_NAME]) else None
         )
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     return model
 
@@ -366,12 +374,10 @@ def fast_tiny(pretrained: bool = False, **kwargs: Any) -> FAST:
     >>> 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 _fast(
@@ -395,12 +401,10 @@ def fast_small(pretrained: bool = False, **kwargs: Any) -> FAST:
     >>> 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 _fast(
@@ -424,12 +428,10 @@ def fast_base(pretrained: bool = False, **kwargs: Any) -> FAST:
     >>> 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 _fast(

doctr/models/detection/fast/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.
@@ -6,7 +6,7 @@
 # Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
 
 from copy import deepcopy
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 import tensorflow as tf
@@ -23,7 +23,7 @@ from .base import _FAST, FASTPostProcessor
 __all__ = ["FAST", "fast_tiny", "fast_small", "fast_base", "reparameterize"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "fast_tiny": {
         "input_shape": (1024, 1024, 3),
         "mean": (0.798, 0.785, 0.772),
@@ -49,7 +49,6 @@ class FastNeck(layers.Layer, NestedObject):
     """Neck of the FAST architecture, composed of a series of 3x3 convolutions and upsampling layer.
 
     Args:
-    ----
         in_channels: number of input channels
         out_channels: number of output channels
     """
@@ -77,7 +76,6 @@ class FastHead(Sequential):
     """Head of the FAST architecture
 
     Args:
-    ----
         in_channels: number of input channels
         num_classes: number of output classes
         out_channels: number of output channels
@@ -104,7 +102,6 @@ class FAST(_FAST, Model, NestedObject):
     <https://arxiv.org/pdf/2111.02394.pdf>`_.
 
     Args:
-    ----
         feature extractor: the backbone serving as feature extractor
         bin_thresh: threshold for binarization
         box_thresh: minimal objectness score to consider a box
@@ -116,7 +113,7 @@ class FAST(_FAST, Model, NestedObject):
         class_names: list of class names
     """
 
-    _children_names: List[str] = ["feat_extractor", "neck", "head", "postprocessor"]
+    _children_names: list[str] = ["feat_extractor", "neck", "head", "postprocessor"]
 
     def __init__(
         self,
@@ -127,8 +124,8 @@ class FAST(_FAST, Model, NestedObject):
         pooling_size: int = 4,  # different from paper performs better on close text-rich images
         assume_straight_pages: bool = True,
         exportable: bool = False,
-        cfg: Optional[Dict[str, Any]] = {},
-        class_names: List[str] = [CLASS_NAME],
+        cfg: dict[str, Any] = {},
+        class_names: list[str] = [CLASS_NAME],
     ) -> None:
         super().__init__()
         self.class_names = class_names
@@ -156,22 +153,29 @@ class FAST(_FAST, Model, NestedObject):
         # Pooling layer as erosion reversal as described in the paper
         self.pooling = layers.MaxPooling2D(pool_size=pooling_size // 2 + 1, strides=1, padding="same")
 
+    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 compute_loss(
         self,
         out_map: tf.Tensor,
-        target: List[Dict[str, np.ndarray]],
+        target: list[dict[str, np.ndarray]],
         eps: float = 1e-6,
     ) -> tf.Tensor:
         """Compute fast loss, 2 x Dice loss where the text kernel loss is scaled by 0.5.
 
         Args:
-        ----
             out_map: output feature map of the model of shape (N, num_classes, H, W)
             target: list of dictionary where each dict has a `boxes` and a `flags` entry
             eps: epsilon factor in dice loss
 
         Returns:
-        -------
             A loss tensor
         """
         targets = self.build_target(target, out_map.shape[1:], True)
@@ -222,18 +226,18 @@ class FAST(_FAST, Model, NestedObject):
     def call(
         self,
         x: tf.Tensor,
-        target: Optional[List[Dict[str, np.ndarray]]] = None,
+        target: list[dict[str, np.ndarray]] | None = None,
         return_model_output: bool = False,
         return_preds: bool = False,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         feat_maps = self.feat_extractor(x, **kwargs)
         # Pass through the Neck & Head & Upsample
         feat_concat = self.neck(feat_maps, **kwargs)
         logits: tf.Tensor = self.head(feat_concat, **kwargs)
         logits = layers.UpSampling2D(size=x.shape[-2] // logits.shape[-2], interpolation="bilinear")(logits, **kwargs)
 
-        out: Dict[str, tf.Tensor] = {}
+        out: dict[str, tf.Tensor] = {}
         if self.exportable:
             out["logits"] = logits
             return out
@@ -255,15 +259,14 @@ class FAST(_FAST, Model, NestedObject):
         return out
 
 
-def reparameterize(model: Union[FAST, layers.Layer]) -> FAST:
+def reparameterize(model: FAST | layers.Layer) -> FAST:
     """Fuse batchnorm and conv layers and reparameterize the model
 
     args:
-    ----
+
         model: the FAST model to reparameterize
 
     Returns:
-    -------
         the reparameterized model
     """
     last_conv = None
@@ -306,9 +309,9 @@ def _fast(
     arch: str,
     pretrained: bool,
     backbone_fn,
-    feat_layers: List[str],
+    feat_layers: list[str],
     pretrained_backbone: bool = True,
-    input_shape: Optional[Tuple[int, int, int]] = None,
+    input_shape: tuple[int, int, int] | None = None,
     **kwargs: Any,
 ) -> FAST:
     pretrained_backbone = pretrained_backbone and not pretrained
@@ -338,8 +341,7 @@ def _fast(
     # Load pretrained parameters
     if pretrained:
         # The given class_names differs from the pretrained model => skip the mismatching layers for fine tuning
-        load_pretrained_params(
-            model,
+        model.from_pretrained(
             _cfg["url"],
             skip_mismatch=kwargs["class_names"] != default_cfgs[arch].get("class_names", [CLASS_NAME]),
         )
@@ -358,12 +360,10 @@ def fast_tiny(pretrained: bool = False, **kwargs: Any) -> FAST:
     >>> 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 _fast(
@@ -386,12 +386,10 @@ def fast_small(pretrained: bool = False, **kwargs: Any) -> FAST:
     >>> 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 _fast(
@@ -414,12 +412,10 @@ def fast_base(pretrained: bool = False, **kwargs: Any) -> FAST:
     >>> 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 _fast(

doctr/models/detection/linknet/__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/linknet/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
@@ -23,7 +22,6 @@ class LinkNetPostProcessor(DetectionPostProcessor):
     """Implements a post processor for LinkNet model.
 
     Args:
-    ----
         bin_thresh: threshold used to binzarized p_map at inference time
         box_thresh: minimal objectness score to consider a box
         assume_straight_pages: whether the inputs were expected to have horizontal text elements
@@ -45,11 +43,9 @@ class LinkNetPostProcessor(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:
@@ -60,9 +56,8 @@ class LinkNetPostProcessor(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)
@@ -94,24 +89,22 @@ class LinkNetPostProcessor(DetectionPostProcessor):
         """Compute boxes from a bitmap/pred_map: find connected components then filter boxes
 
         Args:
-        ----
             pred: Pred map from differentiable linknet 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 6-element list
                 containing x, y, w, h, alpha, score for the box
         """
         height, width = bitmap.shape[: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) < 2):  # type: ignore[index]
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
                 continue
             # Compute objectness
             if self.assume_straight_pages:
@@ -152,7 +145,6 @@ class _LinkNet(BaseModel):
     <https://arxiv.org/pdf/1707.03718.pdf>`_.
 
     Args:
-    ----
         out_chan: number of channels for the output
     """
 
@@ -162,20 +154,18 @@ class _LinkNet(BaseModel):
 
     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]:
+    ) -> tuple[np.ndarray, np.ndarray]:
         """Build the target, and it's mask to be used from loss computation.
 
         Args:
-        ----
             target: target coming from dataset
             output_shape: shape of the output of the model without batch_size
             channels_last: whether channels are last or not
 
         Returns:
-        -------
             the new formatted target and the mask
         """
         if any(t.dtype != np.float32 for tgt in target for t in tgt.values()):