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

doctr/models/recognition/sar/pytorch.py

@@ -228,6 +228,15 @@ 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,
@@ -364,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
 

doctr/models/recognition/sar/tensorflow.py

@@ -255,6 +255,15 @@ class SAR(Model, RecognitionModel):
 
         self.postprocessor = SARPostProcessor(vocab=vocab)
 
+    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)
+
     @staticmethod
     def compute_loss(
         model_output: tf.Tensor,
@@ -389,9 +398,7 @@ def _sar(
     # Load pretrained parameters
     if pretrained:
         # The given vocab differs from the pretrained model => skip the mismatching layers for fine tuning
-        load_pretrained_params(
-            model, default_cfgs[arch]["url"], skip_mismatch=kwargs["vocab"] != default_cfgs[arch]["vocab"]
-        )
+        model.from_pretrained(default_cfgs[arch]["url"], skip_mismatch=kwargs["vocab"] != default_cfgs[arch]["vocab"])
 
     return model
 

doctr/models/recognition/utils.py

@@ -4,81 +4,90 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 
-from rapidfuzz.distance import Levenshtein
+from rapidfuzz.distance import Hamming
 
 __all__ = ["merge_strings", "merge_multi_strings"]
 
 
-def merge_strings(a: str, b: str, dil_factor: float) -> str:
+def merge_strings(a: str, b: str, overlap_ratio: float) -> str:
     """Merges 2 character sequences in the best way to maximize the alignment of their overlapping characters.
 
     Args:
         a: first char seq, suffix should be similar to b's prefix.
         b: second char seq, prefix should be similar to a's suffix.
-        dil_factor: dilation factor of the boxes to overlap, should be > 1. This parameter is
-            only used when the mother sequence is splitted on a character repetition
+        overlap_ratio: estimated ratio of overlapping characters.
 
     Returns:
         A merged character sequence.
 
     Example::
-        >>> from doctr.models.recognition.utils import merge_sequences
-        >>> merge_sequences('abcd', 'cdefgh', 1.4)
+        >>> from doctr.models.recognition.utils import merge_strings
+        >>> merge_strings('abcd', 'cdefgh', 0.5)
         'abcdefgh'
-        >>> merge_sequences('abcdi', 'cdefgh', 1.4)
+        >>> merge_strings('abcdi', 'cdefgh', 0.5)
         'abcdefgh'
     """
     seq_len = min(len(a), len(b))
-    if seq_len == 0:  # One sequence is empty, return the other
-        return b if len(a) == 0 else a
-
-    # Initialize merging index and corresponding score (mean Levenstein)
-    min_score, index = 1.0, 0  # No overlap, just concatenate
-
-    scores = [Levenshtein.distance(a[-i:], b[:i], processor=None) / i for i in range(1, seq_len + 1)]
-
-    # Edge case (split in the middle of char repetitions): if it starts with 2 or more 0
-    if len(scores) > 1 and (scores[0], scores[1]) == (0, 0):
-        # Compute n_overlap (number of overlapping chars, geometrically determined)
-        n_overlap = round(len(b) * (dil_factor - 1) / dil_factor)
-        # Find the number of consecutive zeros in the scores list
-        # Impossible to have a zero after a non-zero score in that case
-        n_zeros = sum(val == 0 for val in scores)
-        # Index is bounded by the geometrical overlap to avoid collapsing repetitions
-        min_score, index = 0, min(n_zeros, n_overlap)
-
-    else:  # Common case: choose the min score index
-        for i, score in enumerate(scores):
-            if score < min_score:
-                min_score, index = score, i + 1  # Add one because first index is an overlap of 1 char
-
-    # Merge with correct overlap
-    if index == 0:
+    if seq_len <= 1:  # One sequence is empty or will be after cropping in next step, return both to keep data
         return a + b
-    return a[:-1] + b[index - 1 :]
 
+    a_crop, b_crop = a[:-1], b[1:]  # Remove last letter of "a" and first of "b", because they might be cut off
+    max_overlap = min(len(a_crop), len(b_crop))
 
-def merge_multi_strings(seq_list: list[str], dil_factor: float) -> str:
-    """Recursively merges consecutive string sequences with overlapping characters.
+    # Compute Hamming distances for all possible overlaps
+    scores = [Hamming.distance(a_crop[-i:], b_crop[:i], processor=None) for i in range(1, max_overlap + 1)]
+
+    # Find zero-score matches
+    zero_matches = [i for i, score in enumerate(scores) if score == 0]
+
+    expected_overlap = round(len(b) * overlap_ratio) - 3  # adjust for cropping and index
+
+    # Case 1: One perfect match - exactly one zero score - just merge there
+    if len(zero_matches) == 1:
+        i = zero_matches[0]
+        return a_crop + b_crop[i + 1 :]
+
+    # Case 2: Multiple perfect matches - likely due to repeated characters.
+    # Use the estimated overlap length to choose the match closest to the expected alignment.
+    elif len(zero_matches) > 1:
+        best_i = min(zero_matches, key=lambda x: abs(x - expected_overlap))
+        return a_crop + b_crop[best_i + 1 :]
+
+    # Case 3: Absence of zero scores indicates that the same character in the image was recognized differently OR that
+    # the overlap was too small and we just need to merge the crops fully
+    if expected_overlap < -1:
+        return a + b
+    elif expected_overlap < 0:
+        return a_crop + b_crop
+
+    # Find best overlap by minimizing Hamming distance + distance from expected overlap size
+    combined_scores = [score + abs(i - expected_overlap) for i, score in enumerate(scores)]
+    best_i = combined_scores.index(min(combined_scores))
+    return a_crop + b_crop[best_i + 1 :]
+
+
+def merge_multi_strings(seq_list: list[str], overlap_ratio: float, last_overlap_ratio: float) -> str:
+    """
+    Merges consecutive string sequences with overlapping characters.
 
     Args:
         seq_list: list of sequences to merge. Sequences need to be ordered from left to right.
-        dil_factor: dilation factor of the boxes to overlap, should be > 1. This parameter is
-            only used when the mother sequence is splitted on a character repetition
+        overlap_ratio: Estimated ratio of overlapping letters between neighboring strings.
+        last_overlap_ratio: Estimated ratio of overlapping letters for the last element in seq_list.
 
     Returns:
         A merged character sequence
 
     Example::
-        >>> from doctr.models.recognition.utils import merge_multi_sequences
-        >>> merge_multi_sequences(['abc', 'bcdef', 'difghi', 'aijkl'], 1.4)
+        >>> from doctr.models.recognition.utils import merge_multi_strings
+        >>> merge_multi_strings(['abc', 'bcdef', 'difghi', 'aijkl'], 0.5, 0.1)
         'abcdefghijkl'
     """
-
-    def _recursive_merge(a: str, seq_list: list[str], dil_factor: float) -> str:
-        # Recursive version of compute_overlap
-        if len(seq_list) == 1:
-            return merge_strings(a, seq_list[0], dil_factor)
-        return _recursive_merge(merge_strings(a, seq_list[0], dil_factor), seq_list[1:], dil_factor)
-
-    return _recursive_merge("", seq_list, dil_factor)
+    if not seq_list:
+        return ""
+    result = seq_list[0]
+    for i in range(1, len(seq_list)):
+        text_b = seq_list[i]
+        ratio = last_overlap_ratio if i == len(seq_list) - 1 else overlap_ratio
+        result = merge_strings(result, text_b, ratio)
+    return result

doctr/models/recognition/viptr/__init__.py

@@ -0,0 +1,4 @@
+from doctr.file_utils import is_torch_available
+
+if is_torch_available():
+    from .pytorch import *

doctr/models/recognition/viptr/pytorch.py

@@ -0,0 +1,277 @@
+# 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 itertools import groupby
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models._utils import IntermediateLayerGetter
+
+from doctr.datasets import VOCABS, decode_sequence
+
+from ...classification import vip_tiny
+from ...utils.pytorch import _bf16_to_float32, load_pretrained_params
+from ..core import RecognitionModel, RecognitionPostProcessor
+
+__all__ = ["VIPTR", "viptr_tiny"]
+
+
+default_cfgs: dict[str, dict[str, Any]] = {
+    "viptr_tiny": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 32, 128),
+        "vocab": VOCABS["french"],
+        "url": "https://doctr-static.mindee.com/models?id=v0.11.0/viptr_tiny-1cb2515e.pt&src=0",
+    },
+}
+
+
+class VIPTRPostProcessor(RecognitionPostProcessor):
+    """Postprocess raw prediction of the model (logits) to a list of words using CTC decoding
+
+    Args:
+        vocab: string containing the ordered sequence of supported characters
+    """
+
+    @staticmethod
+    def ctc_best_path(
+        logits: torch.Tensor,
+        vocab: str = VOCABS["french"],
+        blank: int = 0,
+    ) -> list[tuple[str, float]]:
+        """Implements best path decoding as shown by Graves (Dissertation, p63), highly inspired from
+        <https://github.com/githubharald/CTCDecoder>`_.
+
+        Args:
+            logits: model output, shape: N x T x C
+            vocab: vocabulary to use
+            blank: index of blank label
+
+        Returns:
+            A list of tuples: (word, confidence)
+        """
+        # Gather the most confident characters, and assign the smallest conf among those to the sequence prob
+        probs = F.softmax(logits, dim=-1).max(dim=-1).values.min(dim=1).values
+
+        # collapse best path (using itertools.groupby), map to chars, join char list to string
+        words = [
+            decode_sequence([k for k, _ in groupby(seq.tolist()) if k != blank], vocab)
+            for seq in torch.argmax(logits, dim=-1)
+        ]
+
+        return list(zip(words, probs.tolist()))
+
+    def __call__(self, logits: torch.Tensor) -> list[tuple[str, float]]:
+        """Performs decoding of raw output with CTC and decoding of CTC predictions
+        with label_to_idx mapping dictionnary
+
+        Args:
+            logits: raw output of the model, shape (N, C + 1, seq_len)
+
+        Returns:
+            A tuple of 2 lists: a list of str (words) and a list of float (probs)
+
+        """
+        # Decode CTC
+        return self.ctc_best_path(logits=logits, vocab=self.vocab, blank=len(self.vocab))
+
+
+class VIPTR(RecognitionModel, nn.Module):
+    """Implements a VIPTR architecture as described in `"A Vision Permutable Extractor for Fast and Efficient
+    Scene Text Recognition" <https://arxiv.org/abs/2401.10110>`_.
+
+    Args:
+        feature_extractor: the backbone serving as feature extractor
+        vocab: vocabulary used for encoding
+        input_shape: input shape of the image
+        exportable: onnx exportable returns only logits
+        cfg: configuration dictionary
+    """
+
+    def __init__(
+        self,
+        feature_extractor: nn.Module,
+        vocab: str,
+        input_shape: tuple[int, int, int] = (3, 32, 128),
+        exportable: bool = False,
+        cfg: dict[str, Any] | None = None,
+    ):
+        super().__init__()
+        self.vocab = vocab
+        self.exportable = exportable
+        self.cfg = cfg
+        self.max_length = 32
+        self.vocab_size = len(vocab)
+
+        self.feat_extractor = feature_extractor
+        with torch.inference_mode():
+            embedding_units = self.feat_extractor(torch.zeros((1, *input_shape)))["features"].shape[-1]
+
+        self.postprocessor = VIPTRPostProcessor(vocab=self.vocab)
+        self.head = nn.Linear(embedding_units, len(self.vocab) + 1)  # +1 for PAD
+
+        for n, m in self.named_modules():
+            # Don't override the initialization of the backbone
+            if n.startswith("feat_extractor."):
+                continue
+            if isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    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: list[str] | None = None,
+        return_model_output: bool = False,
+        return_preds: bool = False,
+    ) -> dict[str, Any]:
+        if target is not None:
+            _gt, _seq_len = self.build_target(target)
+            gt, seq_len = torch.from_numpy(_gt).to(dtype=torch.long), torch.tensor(_seq_len)
+            gt, seq_len = gt.to(x.device), seq_len.to(x.device)
+
+        if self.training and target is None:
+            raise ValueError("Need to provide labels during training")
+
+        features = self.feat_extractor(x)["features"]  # (B, max_len, embed_dim)
+        B, N, E = features.size()
+        logits = self.head(features).view(B, N, len(self.vocab) + 1)
+
+        decoded_features = _bf16_to_float32(logits)
+
+        out: dict[str, Any] = {}
+        if self.exportable:
+            out["logits"] = decoded_features
+            return out
+
+        if return_model_output:
+            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"] = _postprocess(decoded_features)
+
+        if target is not None:
+            out["loss"] = self.compute_loss(decoded_features, gt, seq_len, len(self.vocab))
+
+        return out
+
+    @staticmethod
+    def compute_loss(
+        model_output: torch.Tensor,
+        gt: torch.Tensor,
+        seq_len: torch.Tensor,
+        blank_idx: int = 0,
+    ) -> torch.Tensor:
+        """Compute CTC loss for the model.
+
+        Args:
+            model_output: predicted logits of the model
+            gt: ground truth tensor
+            seq_len: sequence lengths of the ground truth
+            blank_idx: index of the blank label
+
+        Returns:
+            The loss of the model on the batch
+        """
+        batch_len = model_output.shape[0]
+        input_length = model_output.shape[1] * torch.ones(size=(batch_len,), dtype=torch.int32)
+        # N x T x C -> T x N x C
+        logits = model_output.permute(1, 0, 2)
+        probs = F.log_softmax(logits, dim=-1)
+        ctc_loss = F.ctc_loss(
+            probs,
+            gt,
+            input_length,
+            seq_len,
+            blank_idx,
+            zero_infinity=True,
+        )
+
+        return ctc_loss
+
+
+def _viptr(
+    arch: str,
+    pretrained: bool,
+    backbone_fn: Callable[[bool], nn.Module],
+    layer: str,
+    pretrained_backbone: bool = True,
+    ignore_keys: list[str] | None = None,
+    **kwargs: Any,
+) -> VIPTR:
+    pretrained_backbone = pretrained_backbone and not pretrained
+
+    # Patch the config
+    _cfg = deepcopy(default_cfgs[arch])
+    _cfg["vocab"] = kwargs.get("vocab", _cfg["vocab"])
+    _cfg["input_shape"] = kwargs.get("input_shape", _cfg["input_shape"])
+
+    # Feature extractor
+    feat_extractor = IntermediateLayerGetter(
+        backbone_fn(pretrained_backbone, input_shape=_cfg["input_shape"]),  # type: ignore[call-arg]
+        {layer: "features"},
+    )
+
+    kwargs["vocab"] = _cfg["vocab"]
+    kwargs["input_shape"] = _cfg["input_shape"]
+
+    model = VIPTR(feat_extractor, cfg=_cfg, **kwargs)
+
+    # Load pretrained parameters
+    if pretrained:
+        # 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
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+
+    return model
+
+
+def viptr_tiny(pretrained: bool = False, **kwargs: Any) -> VIPTR:
+    """VIPTR-Tiny as described in `"A Vision Permutable Extractor for Fast and Efficient Scene Text Recognition"
+    <https://arxiv.org/abs/2401.10110>`_.
+
+    >>> import torch
+    >>> from doctr.models import viptr_tiny
+    >>> model = viptr_tiny(pretrained=False)
+    >>> input_tensor = torch.rand((1, 3, 32, 128))
+    >>> out = model(input_tensor)
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the VIPTR architecture
+
+    Returns:
+        VIPTR: a VIPTR model instance
+    """
+    return _viptr(
+        "viptr_tiny",
+        pretrained,
+        vip_tiny,
+        "5",
+        ignore_keys=["head.weight", "head.bias"],
+        **kwargs,
+    )

doctr/models/recognition/vitstr/pytorch.py

@@ -74,6 +74,15 @@ class ViTSTR(_ViTSTR, nn.Module):
 
         self.postprocessor = ViTSTRPostProcessor(vocab=self.vocab)
 
+    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,
@@ -214,7 +223,7 @@ def _vitstr(
         # 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
 

doctr/models/recognition/vitstr/tensorflow.py

@@ -74,6 +74,15 @@ class ViTSTR(_ViTSTR, Model):
 
         self.postprocessor = ViTSTRPostProcessor(vocab=self.vocab)
 
+    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)
+
     @staticmethod
     def compute_loss(
         model_output: tf.Tensor,
@@ -217,9 +226,7 @@ def _vitstr(
     # Load pretrained parameters
     if pretrained:
         # The given vocab differs from the pretrained model => skip the mismatching layers for fine tuning
-        load_pretrained_params(
-            model, default_cfgs[arch]["url"], skip_mismatch=kwargs["vocab"] != default_cfgs[arch]["vocab"]
-        )
+        model.from_pretrained(default_cfgs[arch]["url"], skip_mismatch=kwargs["vocab"] != default_cfgs[arch]["vocab"])
 
     return model
 

doctr/models/recognition/zoo.py

@@ -25,6 +25,9 @@ ARCHS: list[str] = [
     "parseq",
 ]
 
+if is_torch_available():
+    ARCHS.extend(["viptr_tiny"])
+
 
 def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredictor:
     if isinstance(arch, str):
@@ -37,6 +40,8 @@ def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredict
     else:
         allowed_archs = [recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq]
         if is_torch_available():
+            # Add VIPTR which is only available in torch at the moment
+            allowed_archs.append(recognition.VIPTR)
             # Adding the type for torch compiled models to the allowed architectures
             from doctr.models.utils import _CompiledModule
 

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(

doctr/models/utils/tensorflow.py

@@ -9,6 +9,7 @@ from typing import Any
 
 import tensorflow as tf
 import tf2onnx
+import validators
 from tensorflow.keras import Model, layers
 
 from doctr.utils.data import download_from_url
@@ -47,7 +48,7 @@ def _build_model(model: Model):
 
 def load_pretrained_params(
     model: Model,
-    url: str | None = None,
+    path_or_url: str | None = None,
     hash_prefix: str | None = None,
     skip_mismatch: bool = False,
     **kwargs: Any,
@@ -59,17 +60,23 @@ def load_pretrained_params(
 
     Args:
         model: the keras 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
         skip_mismatch: skip loading layers with mismatched shapes
         **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)
-        # Load weights
-        model.load_weights(archive_path, skip_mismatch=skip_mismatch)
+    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
+    )
+
+    # Load weights
+    model.load_weights(archive_path, skip_mismatch=skip_mismatch)
 
 
 def conv_sequence(