python-doctr 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

doctr/models/recognition/crnn/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -45,10 +45,10 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 
 
 class CTCPostProcessor(RecognitionPostProcessor):
-    """
-    Postprocess raw prediction of the model (logits) to a list of words using CTC decoding
+    """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
     """
 
@@ -62,14 +62,15 @@ class CTCPostProcessor(RecognitionPostProcessor):
         <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
 
@@ -82,14 +83,15 @@ class CTCPostProcessor(RecognitionPostProcessor):
         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
+        """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)
 
         """
@@ -102,6 +104,7 @@ class CRNN(RecognitionModel, nn.Module):
     Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
 
     Args:
+    ----
         feature_extractor: the backbone serving as feature extractor
         vocab: vocabulary used for encoding
         rnn_units: number of units in the LSTM layers
@@ -128,12 +131,9 @@ class CRNN(RecognitionModel, nn.Module):
         self.feat_extractor = feature_extractor
 
         # Resolve the input_size of the LSTM
-        self.feat_extractor.eval()
-        with torch.no_grad():
+        with torch.inference_mode():
             out_shape = self.feat_extractor(torch.zeros((1, *input_shape))).shape
         lstm_in = out_shape[1] * out_shape[2]
-        # Switch back to original mode
-        self.feat_extractor.train()
 
         self.decoder = nn.LSTM(
             input_size=lstm_in,
@@ -168,11 +168,12 @@ class CRNN(RecognitionModel, nn.Module):
         """Compute CTC loss for the model.
 
         Args:
-            gt: the encoded tensor with gt labels
+        ----
             model_output: predicted logits of the model
-            seq_len: lengths of each gt word inside the batch
+            target: list of target strings
 
         Returns:
+        -------
             The loss of the model on the batch
         """
         gt, seq_len = self.build_target(target)
@@ -249,7 +250,7 @@ def _crnn(
     _cfg["input_shape"] = kwargs["input_shape"]
 
     # Build the model
-    model = CRNN(feat_extractor, cfg=_cfg, **kwargs)  # type: ignore[arg-type]
+    model = CRNN(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 =>
@@ -271,12 +272,14 @@ def crnn_vgg16_bn(pretrained: bool = False, **kwargs: Any) -> CRNN:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the CRNN architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _crnn("crnn_vgg16_bn", pretrained, vgg16_bn_r, ignore_keys=["linear.weight", "linear.bias"], **kwargs)
 
 
@@ -291,12 +294,14 @@ def crnn_mobilenet_v3_small(pretrained: bool = False, **kwargs: Any) -> CRNN:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the CRNN architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _crnn(
         "crnn_mobilenet_v3_small",
         pretrained,
@@ -317,12 +322,14 @@ def crnn_mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> CRNN:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the CRNN architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _crnn(
         "crnn_mobilenet_v3_large",
         pretrained,

doctr/models/recognition/crnn/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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,7 +13,7 @@ from tensorflow.keras.models import Model, Sequential
 from doctr.datasets import VOCABS
 
 from ...classification import mobilenet_v3_large_r, mobilenet_v3_small_r, vgg16_bn_r
-from ...utils.tensorflow import load_pretrained_params
+from ...utils.tensorflow import _bf16_to_float32, load_pretrained_params
 from ..core import RecognitionModel, RecognitionPostProcessor
 
 __all__ = ["CRNN", "crnn_vgg16_bn", "crnn_mobilenet_v3_small", "crnn_mobilenet_v3_large"]
@@ -44,10 +44,10 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 
 
 class CTCPostProcessor(RecognitionPostProcessor):
-    """
-    Postprocess raw prediction of the model (logits) to a list of words using CTC decoding
+    """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
         ignore_case: if True, ignore case of letters
         ignore_accents: if True, ignore accents of letters
@@ -59,16 +59,17 @@ class CTCPostProcessor(RecognitionPostProcessor):
         beam_width: int = 1,
         top_paths: int = 1,
     ) -> Union[List[Tuple[str, float]], List[Tuple[List[str], List[float]]]]:
-        """
-        Performs decoding of raw output with CTC and decoding of CTC predictions
+        """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 BATCH_SIZE X SEQ_LEN X NUM_CLASSES + 1
             beam_width: An int scalar >= 0 (beam search beam width).
             top_paths: An int scalar >= 0, <= beam_width (controls output size).
 
         Returns:
+        -------
             A list of decoded words of length BATCH_SIZE
 
 
@@ -113,6 +114,7 @@ class CRNN(RecognitionModel, Model):
     Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
 
     Args:
+    ----
         feature_extractor: the backbone serving as feature extractor
         vocab: vocabulary used for encoding
         rnn_units: number of units in the LSTM layers
@@ -144,13 +146,11 @@ class CRNN(RecognitionModel, Model):
         self.exportable = exportable
         self.feat_extractor = feature_extractor
 
-        self.decoder = Sequential(
-            [
-                layers.Bidirectional(layers.LSTM(units=rnn_units, return_sequences=True)),
-                layers.Bidirectional(layers.LSTM(units=rnn_units, return_sequences=True)),
-                layers.Dense(units=len(vocab) + 1),
-            ]
-        )
+        self.decoder = Sequential([
+            layers.Bidirectional(layers.LSTM(units=rnn_units, return_sequences=True)),
+            layers.Bidirectional(layers.LSTM(units=rnn_units, return_sequences=True)),
+            layers.Dense(units=len(vocab) + 1),
+        ])
         self.decoder.build(input_shape=(None, w, h * c))
 
         self.postprocessor = CTCPostProcessor(vocab=vocab)
@@ -166,10 +166,12 @@ class CRNN(RecognitionModel, Model):
         """Compute CTC loss for the model.
 
         Args:
+        ----
             model_output: predicted logits of the model
             target: lengths of each gt word inside the batch
 
         Returns:
+        -------
             The loss of the model on the batch
         """
         gt, seq_len = self.build_target(target)
@@ -199,7 +201,7 @@ class CRNN(RecognitionModel, Model):
         w, h, c = transposed_feat.get_shape().as_list()[1:]
         # B x W x H x C --> B x W x H * C
         features_seq = tf.reshape(transposed_feat, shape=(-1, w, h * c))
-        logits = self.decoder(features_seq, **kwargs)
+        logits = _bf16_to_float32(self.decoder(features_seq, **kwargs))
 
         out: Dict[str, tf.Tensor] = {}
         if self.exportable:
@@ -261,12 +263,14 @@ def crnn_vgg16_bn(pretrained: bool = False, **kwargs: Any) -> CRNN:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the CRNN architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _crnn("crnn_vgg16_bn", pretrained, vgg16_bn_r, **kwargs)
 
 
@@ -281,12 +285,14 @@ def crnn_mobilenet_v3_small(pretrained: bool = False, **kwargs: Any) -> CRNN:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the CRNN architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _crnn("crnn_mobilenet_v3_small", pretrained, mobilenet_v3_small_r, **kwargs)
 
 
@@ -301,10 +307,12 @@ def crnn_mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> CRNN:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keyword arguments of the CRNN architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _crnn("crnn_mobilenet_v3_large", pretrained, mobilenet_v3_large_r, **kwargs)

doctr/models/recognition/master/base.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -23,9 +23,11 @@ class _MASTER:
         sequence lengths.
 
         Args:
+        ----
             gts: list of ground-truth labels
 
         Returns:
+        -------
             A tuple of 2 tensors: Encoded labels and sequence lengths (for each entry of the batch)
         """
         encoded = encode_sequences(
@@ -44,6 +46,7 @@ class _MASTERPostProcessor(RecognitionPostProcessor):
     """Abstract class to postprocess the raw output of the model
 
     Args:
+    ----
         vocab: string containing the ordered sequence of supported characters
     """
 

doctr/models/recognition/master/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -15,7 +15,7 @@ from doctr.datasets import VOCABS
 from doctr.models.classification import magc_resnet31
 from doctr.models.modules.transformer import Decoder, PositionalEncoding
 
-from ...utils.pytorch import load_pretrained_params
+from ...utils.pytorch import _bf16_to_float32, load_pretrained_params
 from .base import _MASTER, _MASTERPostProcessor
 
 __all__ = ["MASTER", "master"]
@@ -27,7 +27,7 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
         "std": (0.299, 0.296, 0.301),
         "input_shape": (3, 32, 128),
         "vocab": VOCABS["french"],
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/master-fde31e4a.pt&src=0",
     },
 }
 
@@ -37,6 +37,7 @@ class MASTER(_MASTER, nn.Module):
     Implementation based on the official Pytorch implementation: <https://github.com/wenwenyu/MASTER-pytorch>`_.
 
     Args:
+    ----
         feature_extractor: the backbone serving as feature extractor
         vocab: vocabulary, (without EOS, SOS, PAD)
         d_model: d parameter for the transformer decoder
@@ -105,7 +106,8 @@ class MASTER(_MASTER, nn.Module):
         # borrowed and slightly modified from  https://github.com/wenwenyu/MASTER-pytorch
         # NOTE: nn.TransformerDecoder takes the inverse from this implementation
         # [True, True, True, ..., False, False, False] -> False is masked
-        target_pad_mask = (target != self.vocab_size + 2).unsqueeze(1).unsqueeze(1)  # (N, 1, 1, max_length)
+        # (N, 1, 1, max_length)
+        target_pad_mask = (target != self.vocab_size + 2).unsqueeze(1).unsqueeze(1)  # type: ignore[attr-defined]
         target_length = target.size(1)
         # sub mask filled diagonal with True = see and False = masked (max_length, max_length)
         # NOTE: onnxruntime tril/triu works only with float currently (onnxruntime 1.11.1 - opset 14)
@@ -128,17 +130,19 @@ class MASTER(_MASTER, nn.Module):
         Sequences are masked after the EOS character.
 
         Args:
+        ----
             gt: the encoded tensor with gt labels
             model_output: predicted logits of the model
             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 (sos disappear in shift!)
-        seq_len = seq_len + 1
+        seq_len = seq_len + 1  # type: ignore[assignment]
         # Compute loss: don't forget to shift gt! Otherwise the model learns to output the gt[t-1]!
         # The "masked" first gt char is <sos>. Delete last logit of the model output.
         cce = F.cross_entropy(model_output[:, :-1, :].permute(0, 2, 1), gt[:, 1:], reduction="none")
@@ -159,15 +163,16 @@ class MASTER(_MASTER, nn.Module):
         """Call function for training
 
         Args:
+        ----
             x: images
             target: list of str labels
             return_model_output: if True, return logits
             return_preds: if True, decode logits
 
         Returns:
+        -------
             A dictionnary containing eventually loss, logits and predictions.
         """
-
         # Encode
         features = self.feat_extractor(x)["features"]
         b, c, h, w = features.shape
@@ -195,6 +200,8 @@ class MASTER(_MASTER, nn.Module):
         else:
             logits = self.decode(encoded)
 
+        logits = _bf16_to_float32(logits)
+
         if self.exportable:
             out["logits"] = logits
             return out
@@ -214,9 +221,11 @@ class MASTER(_MASTER, nn.Module):
         """Decode function for prediction
 
         Args:
+        ----
             encoded: input tensor
 
-        Return:
+        Returns:
+        -------
             A Tuple of torch.Tensor: predictions, logits
         """
         b = encoded.size(0)
@@ -259,7 +268,7 @@ class MASTERPostProcessor(_MASTERPostProcessor):
             for encoded_seq in out_idxs.cpu().numpy()
         ]
 
-        return list(zip(word_values, probs.numpy().tolist()))
+        return list(zip(word_values, probs.numpy().clip(0, 1).tolist()))
 
 
 def _master(
@@ -307,12 +316,14 @@ def master(pretrained: bool = False, **kwargs: Any) -> MASTER:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keywoard arguments passed to the MASTER architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _master(
         "master",
         pretrained,

doctr/models/recognition/master/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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,7 +13,7 @@ from doctr.datasets import VOCABS
 from doctr.models.classification import magc_resnet31
 from doctr.models.modules.transformer import Decoder, PositionalEncoding
 
-from ...utils.tensorflow import load_pretrained_params
+from ...utils.tensorflow import _bf16_to_float32, load_pretrained_params
 from .base import _MASTER, _MASTERPostProcessor
 
 __all__ = ["MASTER", "master"]
@@ -31,11 +31,11 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 
 
 class MASTER(_MASTER, Model):
-
     """Implements MASTER as described in paper: <https://arxiv.org/pdf/1910.02562.pdf>`_.
     Implementation based on the official TF implementation: <https://github.com/jiangxiluning/MASTER-TF>`_.
 
     Args:
+    ----
         feature_extractor: the backbone serving as feature extractor
         vocab: vocabulary, (without EOS, SOS, PAD)
         d_model: d parameter for the transformer decoder
@@ -115,11 +115,13 @@ class MASTER(_MASTER, Model):
         Sequences are masked after the EOS character.
 
         Args:
+        ----
             gt: the encoded tensor with gt labels
             model_output: predicted logits of the model
             seq_len: lengths of each gt word inside the batch
 
         Returns:
+        -------
             The loss of the model on the batch
         """
         # Input length : number of timesteps
@@ -150,15 +152,17 @@ class MASTER(_MASTER, Model):
         """Call function for training
 
         Args:
+        ----
             x: images
             target: list of str labels
             return_model_output: if True, return logits
             return_preds: if True, decode logits
+            **kwargs: keyword arguments passed to the decoder
 
-        Return:
+        Returns:
+        -------
             A dictionnary containing eventually loss, logits and predictions.
         """
-
         # Encode
         feature = self.feat_extractor(x, **kwargs)
         b, h, w, c = feature.get_shape()
@@ -183,6 +187,8 @@ class MASTER(_MASTER, Model):
         else:
             logits = self.decode(encoded, **kwargs)
 
+        logits = _bf16_to_float32(logits)
+
         if self.exportable:
             out["logits"] = logits
             return out
@@ -203,9 +209,12 @@ class MASTER(_MASTER, Model):
         """Decode function for prediction
 
         Args:
+        ----
             encoded: encoded features
+            **kwargs: keyword arguments passed to the decoder
 
-        Return:
+        Returns:
+        -------
             A Tuple of tf.Tensor: predictions, logits
         """
         b = encoded.shape[0]
@@ -238,6 +247,7 @@ class MASTERPostProcessor(_MASTERPostProcessor):
     """Post processor for MASTER architectures
 
     Args:
+    ----
         vocab: string containing the ordered sequence of supported characters
     """
 
@@ -260,7 +270,7 @@ class MASTERPostProcessor(_MASTERPostProcessor):
         decoded_strings_pred = tf.sparse.to_dense(decoded_strings_pred.to_sparse(), default_value="not valid")[:, 0]
         word_values = [word.decode() for word in decoded_strings_pred.numpy().tolist()]
 
-        return list(zip(word_values, probs.numpy().tolist()))
+        return list(zip(word_values, probs.numpy().clip(0, 1).tolist()))
 
 
 def _master(arch: str, pretrained: bool, backbone_fn, pretrained_backbone: bool = True, **kwargs: Any) -> MASTER:
@@ -297,10 +307,12 @@ def master(pretrained: bool = False, **kwargs: Any) -> MASTER:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text recognition dataset
+        **kwargs: keywoard arguments passed to the MASTER architecture
 
     Returns:
+    -------
         text recognition architecture
     """
-
     return _master("master", pretrained, magc_resnet31, **kwargs)

doctr/models/recognition/parseq/base.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -23,9 +23,11 @@ class _PARSeq:
         sequence lengths.
 
         Args:
+        ----
             gts: list of ground-truth labels
 
         Returns:
+        -------
             A tuple of 2 tensors: Encoded labels and sequence lengths (for each entry of the batch)
         """
         encoded = encode_sequences(
@@ -44,6 +46,7 @@ class _PARSeqPostProcessor(RecognitionPostProcessor):
     """Abstract class to postprocess the raw output of the model
 
     Args:
+    ----
         vocab: string containing the ordered sequence of supported characters
     """