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

doctr/models/preprocessor/tensorflow.py

@@ -1,122 +0,0 @@
-# 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.
-
-import math
-from typing import Any
-
-import numpy as np
-import tensorflow as tf
-
-from doctr.transforms import Normalize, Resize
-from doctr.utils.multithreading import multithread_exec
-from doctr.utils.repr import NestedObject
-
-__all__ = ["PreProcessor"]
-
-
-class PreProcessor(NestedObject):
-    """Implements an abstract preprocessor object which performs casting, resizing, batching and normalization.
-
-    Args:
-        output_size: expected size of each page in format (H, W)
-        batch_size: the size of page batches
-        mean: mean value of the training distribution by channel
-        std: standard deviation of the training distribution by channel
-        **kwargs: additional arguments for the resizing operation
-    """
-
-    _children_names: list[str] = ["resize", "normalize"]
-
-    def __init__(
-        self,
-        output_size: tuple[int, int],
-        batch_size: int,
-        mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
-        std: tuple[float, float, float] = (1.0, 1.0, 1.0),
-        **kwargs: Any,
-    ) -> None:
-        self.batch_size = batch_size
-        self.resize = Resize(output_size, **kwargs)
-        # Perform the division by 255 at the same time
-        self.normalize = Normalize(mean, std)
-        self._runs_on_cuda = tf.config.list_physical_devices("GPU") != []
-
-    def batch_inputs(self, samples: list[tf.Tensor]) -> list[tf.Tensor]:
-        """Gather samples into batches for inference purposes
-
-        Args:
-            samples: list of samples (tf.Tensor)
-
-        Returns:
-            list of batched samples
-        """
-        num_batches = int(math.ceil(len(samples) / self.batch_size))
-        batches = [
-            tf.stack(samples[idx * self.batch_size : min((idx + 1) * self.batch_size, len(samples))], axis=0)
-            for idx in range(int(num_batches))
-        ]
-
-        return batches
-
-    def sample_transforms(self, x: np.ndarray | tf.Tensor) -> tf.Tensor:
-        if x.ndim != 3:
-            raise AssertionError("expected list of 3D Tensors")
-        if isinstance(x, np.ndarray):
-            if x.dtype not in (np.uint8, np.float32):
-                raise TypeError("unsupported data type for numpy.ndarray")
-            x = tf.convert_to_tensor(x)
-        elif x.dtype not in (tf.uint8, tf.float16, tf.float32):
-            raise TypeError("unsupported data type for torch.Tensor")
-        # Data type & 255 division
-        if x.dtype == tf.uint8:
-            x = tf.image.convert_image_dtype(x, dtype=tf.float32)
-        # Resizing
-        x = self.resize(x)
-
-        return x
-
-    def __call__(self, x: tf.Tensor | np.ndarray | list[tf.Tensor | np.ndarray]) -> list[tf.Tensor]:
-        """Prepare document data for model forwarding
-
-        Args:
-            x: list of images (np.array) or tensors (already resized and batched)
-
-        Returns:
-            list of page batches
-        """
-        # Input type check
-        if isinstance(x, (np.ndarray, tf.Tensor)):
-            if x.ndim != 4:
-                raise AssertionError("expected 4D Tensor")
-            if isinstance(x, np.ndarray):
-                if x.dtype not in (np.uint8, np.float32):
-                    raise TypeError("unsupported data type for numpy.ndarray")
-                x = tf.convert_to_tensor(x)
-            elif x.dtype not in (tf.uint8, tf.float16, tf.float32):
-                raise TypeError("unsupported data type for torch.Tensor")
-
-            # Data type & 255 division
-            if x.dtype == tf.uint8:
-                x = tf.image.convert_image_dtype(x, dtype=tf.float32)
-            # Resizing
-            if (x.shape[1], x.shape[2]) != self.resize.output_size:
-                x = tf.image.resize(
-                    x, self.resize.output_size, method=self.resize.method, antialias=self.resize.antialias
-                )
-
-            batches = [x]
-
-        elif isinstance(x, list) and all(isinstance(sample, (np.ndarray, tf.Tensor)) for sample in x):
-            # Sample transform (to tensor, resize)
-            samples = list(multithread_exec(self.sample_transforms, x, threads=1 if self._runs_on_cuda else None))
-            # Batching
-            batches = self.batch_inputs(samples)
-        else:
-            raise TypeError(f"invalid input type: {type(x)}")
-
-        # Batch transforms (normalize)
-        batches = list(multithread_exec(self.normalize, batches, threads=1 if self._runs_on_cuda else None))
-
-        return batches

doctr/models/recognition/crnn/tensorflow.py

@@ -1,317 +0,0 @@
-# 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 copy import deepcopy
-from typing import Any
-
-import tensorflow as tf
-from tensorflow.keras import layers
-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 _bf16_to_float32, _build_model, load_pretrained_params
-from ..core import RecognitionModel, RecognitionPostProcessor
-
-__all__ = ["CRNN", "crnn_vgg16_bn", "crnn_mobilenet_v3_small", "crnn_mobilenet_v3_large"]
-
-default_cfgs: dict[str, dict[str, Any]] = {
-    "crnn_vgg16_bn": {
-        "mean": (0.694, 0.695, 0.693),
-        "std": (0.299, 0.296, 0.301),
-        "input_shape": (32, 128, 3),
-        "vocab": VOCABS["legacy_french"],
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/crnn_vgg16_bn-9c188f45.weights.h5&src=0",
-    },
-    "crnn_mobilenet_v3_small": {
-        "mean": (0.694, 0.695, 0.693),
-        "std": (0.299, 0.296, 0.301),
-        "input_shape": (32, 128, 3),
-        "vocab": VOCABS["french"],
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/crnn_mobilenet_v3_small-54850265.weights.h5&src=0",
-    },
-    "crnn_mobilenet_v3_large": {
-        "mean": (0.694, 0.695, 0.693),
-        "std": (0.299, 0.296, 0.301),
-        "input_shape": (32, 128, 3),
-        "vocab": VOCABS["french"],
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/crnn_mobilenet_v3_large-c64045e5.weights.h5&src=0",
-    },
-}
-
-
-class CTCPostProcessor(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
-        ignore_case: if True, ignore case of letters
-        ignore_accents: if True, ignore accents of letters
-    """
-
-    def __call__(
-        self,
-        logits: tf.Tensor,
-        beam_width: int = 1,
-        top_paths: int = 1,
-    ) -> list[tuple[str, float]] | list[tuple[list[str] | list[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 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
-
-
-        """
-        # Decode CTC
-        _decoded, _log_prob = tf.nn.ctc_beam_search_decoder(
-            tf.transpose(logits, perm=[1, 0, 2]),
-            tf.fill(tf.shape(logits)[:1], tf.shape(logits)[1]),
-            beam_width=beam_width,
-            top_paths=top_paths,
-        )
-
-        _decoded = tf.sparse.concat(
-            1,
-            [tf.sparse.expand_dims(dec, axis=1) for dec in _decoded],
-            expand_nonconcat_dims=True,
-        )  # dim : batchsize x beamwidth x actual_max_len_predictions
-        out_idxs = tf.sparse.to_dense(_decoded, default_value=len(self.vocab))
-
-        # Map it to characters
-        _decoded_strings_pred = tf.strings.reduce_join(
-            inputs=tf.nn.embedding_lookup(tf.constant(self._embedding, dtype=tf.string), out_idxs),
-            axis=-1,
-        )
-        _decoded_strings_pred = tf.strings.split(_decoded_strings_pred, "<eos>")
-        decoded_strings_pred = tf.sparse.to_dense(_decoded_strings_pred.to_sparse(), default_value="not valid")[
-            :, :, 0
-        ]  # dim : batch_size x beam_width
-
-        if top_paths == 1:
-            probs = tf.math.exp(tf.squeeze(_log_prob, axis=1))  # dim : batchsize
-            decoded_strings_pred = tf.squeeze(decoded_strings_pred, axis=1)
-            word_values = [word.decode() for word in decoded_strings_pred.numpy().tolist()]
-        else:
-            probs = tf.math.exp(_log_prob)  # dim : batchsize x beamwidth
-            word_values = [[word.decode() for word in words] for words in decoded_strings_pred.numpy().tolist()]
-        return list(zip(word_values, probs.numpy().tolist()))
-
-
-class CRNN(RecognitionModel, Model):
-    """Implements a CRNN architecture as described in `"An End-to-End Trainable Neural Network for Image-based
-    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
-        exportable: onnx exportable returns only logits
-        beam_width: beam width for beam search decoding
-        top_paths: number of top paths for beam search decoding
-        cfg: configuration dictionary
-    """
-
-    _children_names: list[str] = ["feat_extractor", "decoder", "postprocessor"]
-
-    def __init__(
-        self,
-        feature_extractor: Model,
-        vocab: str,
-        rnn_units: int = 128,
-        exportable: bool = False,
-        beam_width: int = 1,
-        top_paths: int = 1,
-        cfg: dict[str, Any] | None = None,
-    ) -> None:
-        # Initialize kernels
-        h, w, c = feature_extractor.output_shape[1:]
-
-        super().__init__()
-        self.vocab = vocab
-        self.max_length = w
-        self.cfg = cfg
-        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.build(input_shape=(None, w, h * c))
-
-        self.postprocessor = CTCPostProcessor(vocab=vocab)
-
-        self.beam_width = beam_width
-        self.top_paths = top_paths
-
-    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,
-        model_output: tf.Tensor,
-        target: list[str],
-    ) -> tf.Tensor:
-        """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)
-        batch_len = model_output.shape[0]
-        input_length = tf.fill((batch_len,), model_output.shape[1])
-        ctc_loss = tf.nn.ctc_loss(
-            gt, model_output, seq_len, input_length, logits_time_major=False, blank_index=len(self.vocab)
-        )
-        return ctc_loss
-
-    def call(
-        self,
-        x: tf.Tensor,
-        target: list[str] | None = None,
-        return_model_output: bool = False,
-        return_preds: bool = False,
-        beam_width: int = 1,
-        top_paths: int = 1,
-        **kwargs: Any,
-    ) -> dict[str, Any]:
-        if kwargs.get("training", False) and target is None:
-            raise ValueError("Need to provide labels during training")
-
-        features = self.feat_extractor(x, **kwargs)
-        # B x H x W x C --> B x W x H x C
-        transposed_feat = tf.transpose(features, perm=[0, 2, 1, 3])
-        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 = _bf16_to_float32(self.decoder(features_seq, **kwargs))
-
-        out: dict[str, tf.Tensor] = {}
-        if self.exportable:
-            out["logits"] = logits
-            return out
-
-        if return_model_output:
-            out["out_map"] = logits
-
-        if target is None or return_preds:
-            # Post-process boxes
-            out["preds"] = self.postprocessor(logits, beam_width=beam_width, top_paths=top_paths)
-
-        if target is not None:
-            out["loss"] = self.compute_loss(logits, target)
-
-        return out
-
-
-def _crnn(
-    arch: str,
-    pretrained: bool,
-    backbone_fn,
-    pretrained_backbone: bool = True,
-    input_shape: tuple[int, int, int] | None = None,
-    **kwargs: Any,
-) -> CRNN:
-    pretrained_backbone = pretrained_backbone and not pretrained
-
-    kwargs["vocab"] = kwargs.get("vocab", default_cfgs[arch]["vocab"])
-
-    _cfg = deepcopy(default_cfgs[arch])
-    _cfg["vocab"] = kwargs["vocab"]
-    _cfg["input_shape"] = input_shape or default_cfgs[arch]["input_shape"]
-
-    feat_extractor = backbone_fn(
-        input_shape=_cfg["input_shape"],
-        include_top=False,
-        pretrained=pretrained_backbone,
-    )
-
-    # Build the model
-    model = CRNN(feat_extractor, cfg=_cfg, **kwargs)
-    _build_model(model)
-    # Load pretrained parameters
-    if pretrained:
-        # The given vocab differs from the pretrained model => skip the mismatching layers for fine tuning
-        model.from_pretrained(_cfg["url"], skip_mismatch=kwargs["vocab"] != default_cfgs[arch]["vocab"])
-
-    return model
-
-
-def crnn_vgg16_bn(pretrained: bool = False, **kwargs: Any) -> CRNN:
-    """CRNN with a VGG-16 backbone as described in `"An End-to-End Trainable Neural Network for Image-based
-    Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import crnn_vgg16_bn
-    >>> model = crnn_vgg16_bn(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 32, 128, 3], maxval=1, dtype=tf.float32)
-    >>> 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)
-
-
-def crnn_mobilenet_v3_small(pretrained: bool = False, **kwargs: Any) -> CRNN:
-    """CRNN with a MobileNet V3 Small backbone as described in `"An End-to-End Trainable Neural Network for Image-based
-    Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import crnn_mobilenet_v3_small
-    >>> model = crnn_mobilenet_v3_small(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 32, 128, 3], maxval=1, dtype=tf.float32)
-    >>> 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)
-
-
-def crnn_mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> CRNN:
-    """CRNN with a MobileNet V3 Large backbone as described in `"An End-to-End Trainable Neural Network for Image-based
-    Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import crnn_mobilenet_v3_large
-    >>> model = crnn_mobilenet_v3_large(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 32, 128, 3], maxval=1, dtype=tf.float32)
-    >>> 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)