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

doctr/models/recognition/predictor/_utils.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 List, Tuple, Union
+
+import math
 
 import numpy as np
 
@@ -13,74 +14,132 @@ __all__ = ["split_crops", "remap_preds"]
 
 
 def split_crops(
-    crops: List[np.ndarray],
+    crops: list[np.ndarray],
     max_ratio: float,
     target_ratio: int,
-    dilation: float,
+    split_overlap_ratio: float,
     channels_last: bool = True,
-) -> Tuple[List[np.ndarray], List[Union[int, Tuple[int, int]]], bool]:
-    """Chunk crops horizontally to match a given aspect ratio
+) -> tuple[list[np.ndarray], list[int | tuple[int, int, float]], bool]:
+    """
+    Split crops horizontally if they exceed a given aspect ratio.
 
     Args:
-    ----
-        crops: list of numpy array of shape (H, W, 3) if channels_last or (3, H, W) otherwise
-        max_ratio: the maximum aspect ratio that won't trigger the chunk
-        target_ratio: when crops are chunked, they will be chunked to match this aspect ratio
-        dilation: the width dilation of final chunks (to provide some overlaps)
-        channels_last: whether the numpy array has dimensions in channels last order
+        crops: List of image crops (H, W, C) if channels_last else (C, H, W).
+        max_ratio: Aspect ratio threshold above which crops are split.
+        target_ratio: Target aspect ratio after splitting (e.g., 4 for 128x32).
+        split_overlap_ratio: Desired overlap between splits (as a fraction of split width).
+        channels_last: Whether the crops are in channels-last format.
 
     Returns:
-    -------
-        a tuple with the new crops, their mapping, and a boolean specifying whether any remap is required
+        A tuple containing:
+            - The new list of crops (possibly with splits),
+            - A mapping indicating how to reassemble predictions,
+            - A boolean indicating whether remapping is required.
     """
-    _remap_required = False
-    crop_map: List[Union[int, Tuple[int, int]]] = []
-    new_crops: List[np.ndarray] = []
+    if split_overlap_ratio <= 0.0 or split_overlap_ratio >= 1.0:
+        raise ValueError(f"Valid range for split_overlap_ratio is (0.0, 1.0), but is: {split_overlap_ratio}")
+
+    remap_required = False
+    new_crops: list[np.ndarray] = []
+    crop_map: list[int | tuple[int, int, float]] = []
+
     for crop in crops:
         h, w = crop.shape[:2] if channels_last else crop.shape[-2:]
         aspect_ratio = w / h
+
         if aspect_ratio > max_ratio:
-            # Determine the number of crops, reference aspect ratio = 4 = 128 / 32
-            num_subcrops = int(aspect_ratio // target_ratio)
-            # Find the new widths, additional dilation factor to overlap crops
-            width = dilation * w / num_subcrops
-            centers = [(w / num_subcrops) * (1 / 2 + idx) for idx in range(num_subcrops)]
-            # Get the crops
-            if channels_last:
-                _crops = [
-                    crop[:, max(0, int(round(center - width / 2))) : min(w - 1, int(round(center + width / 2))), :]
-                    for center in centers
-                ]
+            split_width = max(1, math.ceil(h * target_ratio))
+            overlap_width = max(0, math.floor(split_width * split_overlap_ratio))
+
+            splits, last_overlap = _split_horizontally(crop, split_width, overlap_width, channels_last)
+
+            # Remove any empty splits
+            splits = [s for s in splits if all(dim > 0 for dim in s.shape)]
+            if splits:
+                crop_map.append((len(new_crops), len(new_crops) + len(splits), last_overlap))
+                new_crops.extend(splits)
+                remap_required = True
             else:
-                _crops = [
-                    crop[:, :, max(0, int(round(center - width / 2))) : min(w - 1, int(round(center + width / 2)))]
-                    for center in centers
-                ]
-            # Avoid sending zero-sized crops
-            _crops = [crop for crop in _crops if all(s > 0 for s in crop.shape)]
-            # Record the slice of crops
-            crop_map.append((len(new_crops), len(new_crops) + len(_crops)))
-            new_crops.extend(_crops)
-            # At least one crop will require merging
-            _remap_required = True
+                # Fallback: treat it as a single crop
+                crop_map.append(len(new_crops))
+                new_crops.append(crop)
         else:
             crop_map.append(len(new_crops))
             new_crops.append(crop)
 
-    return new_crops, crop_map, _remap_required
+    return new_crops, crop_map, remap_required
+
+
+def _split_horizontally(
+    image: np.ndarray, split_width: int, overlap_width: int, channels_last: bool
+) -> tuple[list[np.ndarray], float]:
+    """
+    Horizontally split a single image with overlapping regions.
+
+    Args:
+        image: The image to split (H, W, C) if channels_last else (C, H, W).
+        split_width: Width of each split.
+        overlap_width: Width of the overlapping region.
+        channels_last: Whether the image is in channels-last format.
+
+    Returns:
+        - A list of horizontal image slices.
+        - The actual overlap ratio of the last split.
+    """
+    image_width = image.shape[1] if channels_last else image.shape[-1]
+    if image_width <= split_width:
+        return [image], 0.0
+
+    # Compute start columns for each split
+    step = split_width - overlap_width
+    starts = list(range(0, image_width - split_width + 1, step))
+
+    # Ensure the last patch reaches the end of the image
+    if starts[-1] + split_width < image_width:
+        starts.append(image_width - split_width)
+
+    splits = []
+    for start_col in starts:
+        end_col = start_col + split_width
+        if channels_last:
+            split = image[:, start_col:end_col, :]
+        else:
+            split = image[:, :, start_col:end_col]
+        splits.append(split)
+
+    # Calculate the last overlap ratio, if only one split no overlap
+    last_overlap = 0
+    if len(starts) > 1:
+        last_overlap = (starts[-2] + split_width) - starts[-1]
+    last_overlap_ratio = last_overlap / split_width if split_width else 0.0
+
+    return splits, last_overlap_ratio
 
 
 def remap_preds(
-    preds: List[Tuple[str, float]], crop_map: List[Union[int, Tuple[int, int]]], dilation: float
-) -> List[Tuple[str, float]]:
-    remapped_out = []
-    for _idx in crop_map:
-        # Crop hasn't been split
-        if isinstance(_idx, int):
-            remapped_out.append(preds[_idx])
+    preds: list[tuple[str, float]],
+    crop_map: list[int | tuple[int, int, float]],
+    overlap_ratio: float,
+) -> list[tuple[str, float]]:
+    """
+    Reconstruct predictions from possibly split crops.
+
+    Args:
+        preds: List of (text, confidence) tuples from each crop.
+        crop_map: Map returned by `split_crops`.
+        overlap_ratio: Overlap ratio used during splitting.
+
+    Returns:
+        List of merged (text, confidence) tuples corresponding to original crops.
+    """
+    remapped = []
+    for item in crop_map:
+        if isinstance(item, int):
+            remapped.append(preds[item])
         else:
-            # unzip
-            vals, probs = zip(*preds[_idx[0] : _idx[1]])
-            # Merge the string values
-            remapped_out.append((merge_multi_strings(vals, dilation), min(probs)))  # type: ignore[arg-type]
-    return remapped_out
+            start_idx, end_idx, last_overlap = item
+            text_parts, confidences = zip(*preds[start_idx:end_idx])
+            merged_text = merge_multi_strings(list(text_parts), overlap_ratio, last_overlap)
+            merged_conf = sum(confidences) / len(confidences)  # average confidence
+            remapped.append((merged_text, merged_conf))
+    return remapped

doctr/models/recognition/predictor/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, List, Sequence, Tuple, Union
+from collections.abc import Sequence
+from typing import Any
 
 import numpy as np
 import torch
@@ -21,7 +22,6 @@ class RecognitionPredictor(nn.Module):
     """Implements an object able to identify character sequences in images
 
     Args:
-    ----
         pre_processor: transform inputs for easier batched model inference
         model: core detection architecture
         split_wide_crops: wether to use crop splitting for high aspect ratio crops
@@ -38,15 +38,15 @@ class RecognitionPredictor(nn.Module):
         self.model = model.eval()
         self.split_wide_crops = split_wide_crops
         self.critical_ar = 8  # Critical aspect ratio
-        self.dil_factor = 1.4  # Dilation factor to overlap the crops
+        self.overlap_ratio = 0.5  # Ratio of overlap between neighboring crops
         self.target_ar = 6  # Target aspect ratio
 
     @torch.inference_mode()
     def forward(
         self,
-        crops: Sequence[Union[np.ndarray, torch.Tensor]],
+        crops: Sequence[np.ndarray | torch.Tensor],
         **kwargs: Any,
-    ) -> List[Tuple[str, float]]:
+    ) -> list[tuple[str, float]]:
         if len(crops) == 0:
             return []
         # Dimension check
@@ -60,14 +60,14 @@ class RecognitionPredictor(nn.Module):
                 crops,  # type: ignore[arg-type]
                 self.critical_ar,
                 self.target_ar,
-                self.dil_factor,
+                self.overlap_ratio,
                 isinstance(crops[0], np.ndarray),
             )
             if remapped:
                 crops = new_crops
 
         # Resize & batch them
-        processed_batches = self.pre_processor(crops)
+        processed_batches = self.pre_processor(crops)  # type: ignore[arg-type]
 
         # Forward it
         _params = next(self.model.parameters())
@@ -81,6 +81,6 @@ class RecognitionPredictor(nn.Module):
 
         # Remap crops
         if self.split_wide_crops and remapped:
-            out = remap_preds(out, crop_map, self.dil_factor)
+            out = remap_preds(out, crop_map, self.overlap_ratio)
 
         return out

doctr/models/recognition/predictor/tensorflow.py

@@ -1,9 +1,9 @@
-# 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, Tuple, Union
+from typing import Any
 
 import numpy as np
 import tensorflow as tf
@@ -21,13 +21,12 @@ class RecognitionPredictor(NestedObject):
     """Implements an object able to identify character sequences in images
 
     Args:
-    ----
         pre_processor: transform inputs for easier batched model inference
         model: core detection architecture
         split_wide_crops: wether to use crop splitting for high aspect ratio crops
     """
 
-    _children_names: List[str] = ["pre_processor", "model"]
+    _children_names: list[str] = ["pre_processor", "model"]
 
     def __init__(
         self,
@@ -40,14 +39,14 @@ class RecognitionPredictor(NestedObject):
         self.model = model
         self.split_wide_crops = split_wide_crops
         self.critical_ar = 8  # Critical aspect ratio
-        self.dil_factor = 1.4  # Dilation factor to overlap the crops
+        self.overlap_ratio = 0.5  # Ratio of overlap between neighboring crops
         self.target_ar = 6  # Target aspect ratio
 
     def __call__(
         self,
-        crops: List[Union[np.ndarray, tf.Tensor]],
+        crops: list[np.ndarray | tf.Tensor],
         **kwargs: Any,
-    ) -> List[Tuple[str, float]]:
+    ) -> list[tuple[str, float]]:
         if len(crops) == 0:
             return []
         # Dimension check
@@ -57,7 +56,7 @@ class RecognitionPredictor(NestedObject):
         # Split crops that are too wide
         remapped = False
         if self.split_wide_crops:
-            new_crops, crop_map, remapped = split_crops(crops, self.critical_ar, self.target_ar, self.dil_factor)
+            new_crops, crop_map, remapped = split_crops(crops, self.critical_ar, self.target_ar, self.overlap_ratio)
             if remapped:
                 crops = new_crops
 
@@ -75,6 +74,6 @@ class RecognitionPredictor(NestedObject):
 
         # Remap crops
         if self.split_wide_crops and remapped:
-            out = remap_preds(out, crop_map, self.dil_factor)
+            out = remap_preds(out, crop_map, self.overlap_ratio)
 
         return out

doctr/models/recognition/sar/__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/recognition/sar/pytorch.py

@@ -1,10 +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 collections.abc import Callable
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any
 
 import torch
 from torch import nn
@@ -19,7 +20,7 @@ from ..core import RecognitionModel, RecognitionPostProcessor
 
 __all__ = ["SAR", "sar_resnet31"]
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "sar_resnet31": {
         "mean": (0.694, 0.695, 0.693),
         "std": (0.299, 0.296, 0.301),
@@ -80,7 +81,6 @@ class SARDecoder(nn.Module):
     """Implements decoder module of the SAR model
 
     Args:
-    ----
         rnn_units: number of hidden units in recurrent cells
         max_length: maximum length of a sequence
         vocab_size: number of classes in the model alphabet
@@ -114,12 +114,12 @@ class SARDecoder(nn.Module):
         self,
         features: torch.Tensor,  # (N, C, H, W)
         holistic: torch.Tensor,  # (N, C)
-        gt: Optional[torch.Tensor] = None,  # (N, L)
+        gt: torch.Tensor | None = None,  # (N, L)
     ) -> torch.Tensor:
         if gt is not None:
             gt_embedding = self.embed_tgt(gt)
 
-        logits_list: List[torch.Tensor] = []
+        logits_list: list[torch.Tensor] = []
 
         for t in range(self.max_length + 1):  # 32
             if t == 0:
@@ -166,7 +166,6 @@ class SAR(nn.Module, RecognitionModel):
     Irregular Text Recognition" <https://arxiv.org/pdf/1811.00751.pdf>`_.
 
     Args:
-    ----
         feature_extractor: the backbone serving as feature extractor
         vocab: vocabulary used for encoding
         rnn_units: number of hidden units in both encoder and decoder LSTM
@@ -187,9 +186,9 @@ class SAR(nn.Module, RecognitionModel):
         attention_units: int = 512,
         max_length: int = 30,
         dropout_prob: float = 0.0,
-        input_shape: Tuple[int, int, int] = (3, 32, 128),
+        input_shape: tuple[int, int, int] = (3, 32, 128),
         exportable: bool = False,
-        cfg: Optional[Dict[str, Any]] = None,
+        cfg: dict[str, Any] | None = None,
     ) -> None:
         super().__init__()
         self.vocab = vocab
@@ -229,13 +228,22 @@ class SAR(nn.Module, RecognitionModel):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
+    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[str]] = None,
+        target: list[str] | None = None,
         return_model_output: bool = False,
         return_preds: bool = False,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         features = self.feat_extractor(x)["features"]
         # NOTE: use max instead of functional max_pool2d which leads to ONNX incompatibility (kernel_size)
         # Vertical max pooling (N, C, H, W) --> (N, C, W)
@@ -254,7 +262,7 @@ class SAR(nn.Module, RecognitionModel):
 
         decoded_features = _bf16_to_float32(self.decoder(features, encoded, gt=None if target is None else gt))
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
         if self.exportable:
             out["logits"] = decoded_features
             return out
@@ -263,8 +271,13 @@ class SAR(nn.Module, RecognitionModel):
             out["out_map"] = decoded_features
 
         if target is None or return_preds:
+            # Disable for torch.compile compatibility
+            @torch.compiler.disable  # type: ignore[attr-defined]
+            def _postprocess(decoded_features: torch.Tensor) -> list[tuple[str, float]]:
+                return self.postprocessor(decoded_features)
+
             # Post-process boxes
-            out["preds"] = self.postprocessor(decoded_features)
+            out["preds"] = _postprocess(decoded_features)
 
         if target is not None:
             out["loss"] = self.compute_loss(decoded_features, gt, seq_len)
@@ -281,19 +294,17 @@ class SAR(nn.Module, RecognitionModel):
         Sequences are masked after the EOS character.
 
         Args:
-        ----
             model_output: predicted logits of the model
             gt: the encoded tensor with gt labels
             seq_len: lengths of each gt word inside the batch
 
         Returns:
-        -------
             The loss of the model on the batch
         """
         # Input length : number of timesteps
         input_len = model_output.shape[1]
         # Add one for additional <eos> token
-        seq_len = seq_len + 1
+        seq_len = seq_len + 1  # type: ignore[assignment]
         # Compute loss
         # (N, L, vocab_size + 1)
         cce = F.cross_entropy(model_output.permute(0, 2, 1), gt, reduction="none")
@@ -308,14 +319,13 @@ class SARPostProcessor(RecognitionPostProcessor):
     """Post processor for SAR architectures
 
     Args:
-    ----
         vocab: string containing the ordered sequence of supported characters
     """
 
     def __call__(
         self,
         logits: torch.Tensor,
-    ) -> List[Tuple[str, float]]:
+    ) -> list[tuple[str, float]]:
         # compute pred with argmax for attention models
         out_idxs = logits.argmax(-1)
         # N x L
@@ -338,7 +348,7 @@ def _sar(
     backbone_fn: Callable[[bool], nn.Module],
     layer: str,
     pretrained_backbone: bool = True,
-    ignore_keys: Optional[List[str]] = None,
+    ignore_keys: list[str] | None = None,
     **kwargs: Any,
 ) -> SAR:
     pretrained_backbone = pretrained_backbone and not pretrained
@@ -363,7 +373,7 @@ def _sar(
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if _cfg["vocab"] != default_cfgs[arch]["vocab"] 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
 
@@ -379,12 +389,10 @@ def sar_resnet31(pretrained: bool = False, **kwargs: Any) -> SAR:
     >>> out = model(input_tensor)
 
     Args:
-    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
         **kwargs: keyword arguments of the SAR architecture
 
     Returns:
-    -------
         text recognition architecture
     """
     return _sar(