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

doctr/models/classification/vgg/tensorflow.py

@@ -1,125 +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
-
-from tensorflow.keras import layers
-from tensorflow.keras.models import Sequential
-
-from doctr.datasets import VOCABS
-
-from ...utils import _build_model, conv_sequence, load_pretrained_params
-
-__all__ = ["VGG", "vgg16_bn_r"]
-
-
-default_cfgs: dict[str, dict[str, Any]] = {
-    "vgg16_bn_r": {
-        "mean": (0.5, 0.5, 0.5),
-        "std": (1.0, 1.0, 1.0),
-        "input_shape": (32, 32, 3),
-        "classes": list(VOCABS["french"]),
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/vgg16_bn_r-b4d69212.weights.h5&src=0",
-    },
-}
-
-
-class VGG(Sequential):
-    """Implements the VGG architecture from `"Very Deep Convolutional Networks for Large-Scale Image Recognition"
-    <https://arxiv.org/pdf/1409.1556.pdf>`_.
-
-    Args:
-        num_blocks: number of convolutional block in each stage
-        planes: number of output channels in each stage
-        rect_pools: whether pooling square kernels should be replace with rectangular ones
-        include_top: whether the classifier head should be instantiated
-        num_classes: number of output classes
-        input_shape: shapes of the input tensor
-    """
-
-    def __init__(
-        self,
-        num_blocks: list[int],
-        planes: list[int],
-        rect_pools: list[bool],
-        include_top: bool = False,
-        num_classes: int = 1000,
-        input_shape: tuple[int, int, int] | None = None,
-        cfg: dict[str, Any] | None = None,
-    ) -> None:
-        _layers = []
-        # Specify input_shape only for the first layer
-        kwargs = {"input_shape": input_shape}
-        for nb_blocks, out_chan, rect_pool in zip(num_blocks, planes, rect_pools):
-            for _ in range(nb_blocks):
-                _layers.extend(conv_sequence(out_chan, "relu", True, kernel_size=3, **kwargs))  # type: ignore[arg-type]
-                kwargs = {}
-            _layers.append(layers.MaxPooling2D((2, 1 if rect_pool else 2)))
-
-        if include_top:
-            _layers.extend([layers.GlobalAveragePooling2D(), layers.Dense(num_classes)])
-        super().__init__(_layers)
-        self.cfg = cfg
-
-    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 _vgg(
-    arch: str, pretrained: bool, num_blocks: list[int], planes: list[int], rect_pools: list[bool], **kwargs: Any
-) -> VGG:
-    kwargs["num_classes"] = kwargs.get("num_classes", len(default_cfgs[arch]["classes"]))
-    kwargs["input_shape"] = kwargs.get("input_shape", default_cfgs[arch]["input_shape"])
-    kwargs["classes"] = kwargs.get("classes", default_cfgs[arch]["classes"])
-
-    _cfg = deepcopy(default_cfgs[arch])
-    _cfg["num_classes"] = kwargs["num_classes"]
-    _cfg["classes"] = kwargs["classes"]
-    _cfg["input_shape"] = kwargs["input_shape"]
-    kwargs.pop("classes")
-
-    # Build the model
-    model = VGG(num_blocks, planes, rect_pools, cfg=_cfg, **kwargs)
-    _build_model(model)
-
-    # Load pretrained parameters
-    if pretrained:
-        # The number of classes is not the same as the number of classes in the pretrained model =>
-        # skip the mismatching layers for fine tuning
-        model.from_pretrained(
-            default_cfgs[arch]["url"], skip_mismatch=kwargs["num_classes"] != len(default_cfgs[arch]["classes"])
-        )
-
-    return model
-
-
-def vgg16_bn_r(pretrained: bool = False, **kwargs: Any) -> VGG:
-    """VGG-16 architecture as described in `"Very Deep Convolutional Networks for Large-Scale Image Recognition"
-    <https://arxiv.org/pdf/1409.1556.pdf>`_, modified by adding batch normalization, rectangular pooling and a simpler
-    classification head.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import vgg16_bn_r
-    >>> model = vgg16_bn_r(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-        **kwargs: keyword arguments of the VGG architecture
-
-    Returns:
-        VGG feature extractor
-    """
-    return _vgg(
-        "vgg16_bn_r", pretrained, [2, 2, 3, 3, 3], [64, 128, 256, 512, 512], [False, False, True, True, True], **kwargs
-    )

doctr/models/classification/vit/tensorflow.py

@@ -1,201 +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 Sequential, layers
-
-from doctr.datasets import VOCABS
-from doctr.models.modules.transformer import EncoderBlock
-from doctr.models.modules.vision_transformer.tensorflow import PatchEmbedding
-from doctr.utils.repr import NestedObject
-
-from ...utils import _build_model, load_pretrained_params
-
-__all__ = ["vit_s", "vit_b"]
-
-
-default_cfgs: dict[str, dict[str, Any]] = {
-    "vit_s": {
-        "mean": (0.694, 0.695, 0.693),
-        "std": (0.299, 0.296, 0.301),
-        "input_shape": (3, 32, 32),
-        "classes": list(VOCABS["french"]),
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/vit_s-69bc459e.weights.h5&src=0",
-    },
-    "vit_b": {
-        "mean": (0.694, 0.695, 0.693),
-        "std": (0.299, 0.296, 0.301),
-        "input_shape": (32, 32, 3),
-        "classes": list(VOCABS["french"]),
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/vit_b-c64705bd.weights.h5&src=0",
-    },
-}
-
-
-class ClassifierHead(layers.Layer, NestedObject):
-    """Classifier head for Vision Transformer
-
-    Args:
-        num_classes: number of output classes
-    """
-
-    def __init__(self, num_classes: int) -> None:
-        super().__init__()
-
-        self.head = layers.Dense(num_classes, kernel_initializer="he_normal", name="dense")
-
-    def call(self, x: tf.Tensor) -> tf.Tensor:
-        # (batch_size, num_classes) cls token
-        return self.head(x[:, 0])
-
-
-class VisionTransformer(Sequential):
-    """VisionTransformer architecture as described in
-    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale",
-    <https://arxiv.org/pdf/2010.11929.pdf>`_.
-
-    Args:
-        d_model: dimension of the transformer layers
-        num_layers: number of transformer layers
-        num_heads: number of attention heads
-        ffd_ratio: multiplier for the hidden dimension of the feedforward layer
-        patch_size: size of the patches
-        input_shape: size of the input image
-        dropout: dropout rate
-        num_classes: number of output classes
-        include_top: whether the classifier head should be instantiated
-    """
-
-    def __init__(
-        self,
-        d_model: int,
-        num_layers: int,
-        num_heads: int,
-        ffd_ratio: int,
-        patch_size: tuple[int, int] = (4, 4),
-        input_shape: tuple[int, int, int] = (32, 32, 3),
-        dropout: float = 0.0,
-        num_classes: int = 1000,
-        include_top: bool = True,
-        cfg: dict[str, Any] | None = None,
-    ) -> None:
-        _layers = [
-            PatchEmbedding(input_shape, d_model, patch_size),
-            EncoderBlock(
-                num_layers,
-                num_heads,
-                d_model,
-                d_model * ffd_ratio,
-                dropout,
-                activation_fct=layers.Activation("gelu"),
-            ),
-        ]
-        if include_top:
-            _layers.append(ClassifierHead(num_classes))
-
-        super().__init__(_layers)
-        self.cfg = cfg
-
-    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 _vit(
-    arch: str,
-    pretrained: bool,
-    **kwargs: Any,
-) -> VisionTransformer:
-    kwargs["num_classes"] = kwargs.get("num_classes", len(default_cfgs[arch]["classes"]))
-    kwargs["input_shape"] = kwargs.get("input_shape", default_cfgs[arch]["input_shape"])
-    kwargs["classes"] = kwargs.get("classes", default_cfgs[arch]["classes"])
-
-    _cfg = deepcopy(default_cfgs[arch])
-    _cfg["num_classes"] = kwargs["num_classes"]
-    _cfg["input_shape"] = kwargs["input_shape"]
-    _cfg["classes"] = kwargs["classes"]
-    kwargs.pop("classes")
-
-    # Build the model
-    model = VisionTransformer(cfg=_cfg, **kwargs)
-    _build_model(model)
-
-    # Load pretrained parameters
-    if pretrained:
-        # The number of classes is not the same as the number of classes in the pretrained model =>
-        # skip the mismatching layers for fine tuning
-        load_pretrained_params(
-            model, default_cfgs[arch]["url"], skip_mismatch=kwargs["num_classes"] != len(default_cfgs[arch]["classes"])
-        )
-
-    return model
-
-
-def vit_s(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
-    """VisionTransformer-S architecture
-    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale",
-    <https://arxiv.org/pdf/2010.11929.pdf>`_. Patches: (H, W) -> (H/8, W/8)
-
-    NOTE: unofficial config used in ViTSTR and ParSeq
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import vit_s
-    >>> model = vit_s(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 32, 32, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the VisionTransformer architecture
-
-    Returns:
-        A feature extractor model
-    """
-    return _vit(
-        "vit_s",
-        pretrained,
-        d_model=384,
-        num_layers=12,
-        num_heads=6,
-        ffd_ratio=4,
-        **kwargs,
-    )
-
-
-def vit_b(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
-    """VisionTransformer-B architecture as described in
-    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale",
-    <https://arxiv.org/pdf/2010.11929.pdf>`_. Patches: (H, W) -> (H/8, W/8)
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import vit_b
-    >>> model = vit_b(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 32, 32, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the VisionTransformer architecture
-
-    Returns:
-        A feature extractor model
-    """
-    return _vit(
-        "vit_b",
-        pretrained,
-        d_model=768,
-        num_layers=12,
-        num_heads=12,
-        ffd_ratio=4,
-        **kwargs,
-    )

doctr/models/detection/_utils/tensorflow.py

@@ -1,34 +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 tensorflow as tf
-
-__all__ = ["erode", "dilate"]
-
-
-def erode(x: tf.Tensor, kernel_size: int) -> tf.Tensor:
-    """Performs erosion on a given tensor
-
-    Args:
-        x: boolean tensor of shape (N, H, W, C)
-        kernel_size: the size of the kernel to use for erosion
-
-    Returns:
-        the eroded tensor
-    """
-    return 1 - tf.nn.max_pool2d(1 - x, kernel_size, strides=1, padding="SAME")
-
-
-def dilate(x: tf.Tensor, kernel_size: int) -> tf.Tensor:
-    """Performs dilation on a given tensor
-
-    Args:
-        x: boolean tensor of shape (N, H, W, C)
-        kernel_size: the size of the kernel to use for dilation
-
-    Returns:
-        the dilated tensor
-    """
-    return tf.nn.max_pool2d(x, kernel_size, strides=1, padding="SAME")