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

doctr/models/classification/resnet/tensorflow.py

@@ -1,418 +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 types
-from collections.abc import Callable
-from copy import deepcopy
-from typing import Any
-
-import tensorflow as tf
-from tensorflow.keras import layers
-from tensorflow.keras.applications import ResNet50
-from tensorflow.keras.models import Sequential
-
-from doctr.datasets import VOCABS
-
-from ...utils import _build_model, conv_sequence, load_pretrained_params
-
-__all__ = ["ResNet", "resnet18", "resnet31", "resnet34", "resnet50", "resnet34_wide"]
-
-
-default_cfgs: dict[str, dict[str, Any]] = {
-    "resnet18": {
-        "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/resnet18-f42d3854.weights.h5&src=0",
-    },
-    "resnet31": {
-        "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/resnet31-ab75f78c.weights.h5&src=0",
-    },
-    "resnet34": {
-        "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/resnet34-03967df9.weights.h5&src=0",
-    },
-    "resnet50": {
-        "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/resnet50-82358f34.weights.h5&src=0",
-    },
-    "resnet34_wide": {
-        "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/resnet34_wide-b18fdf79.weights.h5&src=0",
-    },
-}
-
-
-class ResnetBlock(layers.Layer):
-    """Implements a resnet31 block with shortcut
-
-    Args:
-        conv_shortcut: Use of shortcut
-        output_channels: number of channels to use in Conv2D
-        kernel_size: size of square kernels
-        strides: strides to use in the first convolution of the block
-    """
-
-    def __init__(self, output_channels: int, conv_shortcut: bool, strides: int = 1, **kwargs) -> None:
-        super().__init__(**kwargs)
-        if conv_shortcut:
-            self.shortcut = Sequential([
-                layers.Conv2D(
-                    filters=output_channels,
-                    strides=strides,
-                    padding="same",
-                    kernel_size=1,
-                    use_bias=False,
-                    kernel_initializer="he_normal",
-                ),
-                layers.BatchNormalization(),
-            ])
-        else:
-            self.shortcut = layers.Lambda(lambda x: x)
-        self.conv_block = Sequential(self.conv_resnetblock(output_channels, 3, strides))
-        self.act = layers.Activation("relu")
-
-    @staticmethod
-    def conv_resnetblock(
-        output_channels: int,
-        kernel_size: int,
-        strides: int = 1,
-    ) -> list[layers.Layer]:
-        return [
-            *conv_sequence(output_channels, "relu", bn=True, strides=strides, kernel_size=kernel_size),
-            *conv_sequence(output_channels, None, bn=True, kernel_size=kernel_size),
-        ]
-
-    def call(self, inputs: tf.Tensor) -> tf.Tensor:
-        clone = self.shortcut(inputs)
-        conv_out = self.conv_block(inputs)
-        out = self.act(clone + conv_out)
-
-        return out
-
-
-def resnet_stage(
-    num_blocks: int, out_channels: int, shortcut: bool = False, downsample: bool = False
-) -> list[layers.Layer]:
-    _layers: list[layers.Layer] = [ResnetBlock(out_channels, conv_shortcut=shortcut, strides=2 if downsample else 1)]
-
-    for _ in range(1, num_blocks):
-        _layers.append(ResnetBlock(out_channels, conv_shortcut=False))
-
-    return _layers
-
-
-class ResNet(Sequential):
-    """Implements a ResNet architecture
-
-    Args:
-        num_blocks: number of resnet block in each stage
-        output_channels: number of channels in each stage
-        stage_downsample: whether the first residual block of a stage should downsample
-        stage_conv: whether to add a conv_sequence after each stage
-        stage_pooling: pooling to add after each stage (if None, no pooling)
-        origin_stem: whether to use the orginal ResNet stem or ResNet-31's
-        stem_channels: number of output channels of the stem convolutions
-        attn_module: attention module to use in each stage
-        include_top: whether the classifier head should be instantiated
-        num_classes: number of output classes
-        input_shape: shape of inputs
-    """
-
-    def __init__(
-        self,
-        num_blocks: list[int],
-        output_channels: list[int],
-        stage_downsample: list[bool],
-        stage_conv: list[bool],
-        stage_pooling: list[tuple[int, int] | None],
-        origin_stem: bool = True,
-        stem_channels: int = 64,
-        attn_module: Callable[[int], layers.Layer] | None = None,
-        include_top: bool = True,
-        num_classes: int = 1000,
-        cfg: dict[str, Any] | None = None,
-        input_shape: tuple[int, int, int] | None = None,
-    ) -> None:
-        inplanes = stem_channels
-        if origin_stem:
-            _layers = [
-                *conv_sequence(inplanes, "relu", True, kernel_size=7, strides=2, input_shape=input_shape),
-                layers.MaxPool2D(pool_size=(3, 3), strides=2, padding="same"),
-            ]
-        else:
-            _layers = [
-                *conv_sequence(inplanes // 2, "relu", True, kernel_size=3, input_shape=input_shape),
-                *conv_sequence(inplanes, "relu", True, kernel_size=3),
-                layers.MaxPool2D(pool_size=2, strides=2, padding="valid"),
-            ]
-
-        for n_blocks, out_chan, down, conv, pool in zip(
-            num_blocks, output_channels, stage_downsample, stage_conv, stage_pooling
-        ):
-            _layers.extend(resnet_stage(n_blocks, out_chan, out_chan != inplanes, down))
-            if attn_module is not None:
-                _layers.append(attn_module(out_chan))
-            if conv:
-                _layers.extend(conv_sequence(out_chan, activation="relu", bn=True, kernel_size=3))
-            if pool:
-                _layers.append(layers.MaxPool2D(pool_size=pool, strides=pool, padding="valid"))
-            inplanes = out_chan
-
-        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 _resnet(
-    arch: str,
-    pretrained: bool,
-    num_blocks: list[int],
-    output_channels: list[int],
-    stage_downsample: list[bool],
-    stage_conv: list[bool],
-    stage_pooling: list[tuple[int, int] | None],
-    origin_stem: bool = True,
-    **kwargs: Any,
-) -> ResNet:
-    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 = ResNet(
-        num_blocks, output_channels, stage_downsample, stage_conv, stage_pooling, origin_stem, 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 resnet18(pretrained: bool = False, **kwargs: Any) -> ResNet:
-    """Resnet-18 architecture as described in `"Deep Residual Learning for Image Recognition",
-    <https://arxiv.org/pdf/1512.03385.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import resnet18
-    >>> model = resnet18(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the ResNet architecture
-
-    Returns:
-        A classification model
-    """
-    return _resnet(
-        "resnet18",
-        pretrained,
-        [2, 2, 2, 2],
-        [64, 128, 256, 512],
-        [False, True, True, True],
-        [False] * 4,
-        [None] * 4,
-        True,
-        **kwargs,
-    )
-
-
-def resnet31(pretrained: bool = False, **kwargs: Any) -> ResNet:
-    """Resnet31 architecture with rectangular pooling windows as described in
-    `"Show, Attend and Read:A Simple and Strong Baseline for Irregular Text Recognition",
-    <https://arxiv.org/pdf/1811.00751.pdf>`_. Downsizing: (H, W) --> (H/8, W/4)
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import resnet31
-    >>> model = resnet31(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the ResNet architecture
-
-    Returns:
-        A classification model
-    """
-    return _resnet(
-        "resnet31",
-        pretrained,
-        [1, 2, 5, 3],
-        [256, 256, 512, 512],
-        [False] * 4,
-        [True] * 4,
-        [(2, 2), (2, 1), None, None],
-        False,
-        stem_channels=128,
-        **kwargs,
-    )
-
-
-def resnet34(pretrained: bool = False, **kwargs: Any) -> ResNet:
-    """Resnet-34 architecture as described in `"Deep Residual Learning for Image Recognition",
-    <https://arxiv.org/pdf/1512.03385.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import resnet34
-    >>> model = resnet34(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the ResNet architecture
-
-    Returns:
-        A classification model
-    """
-    return _resnet(
-        "resnet34",
-        pretrained,
-        [3, 4, 6, 3],
-        [64, 128, 256, 512],
-        [False, True, True, True],
-        [False] * 4,
-        [None] * 4,
-        True,
-        **kwargs,
-    )
-
-
-def resnet50(pretrained: bool = False, **kwargs: Any) -> ResNet:
-    """Resnet-50 architecture as described in `"Deep Residual Learning for Image Recognition",
-    <https://arxiv.org/pdf/1512.03385.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import resnet50
-    >>> model = resnet50(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the ResNet architecture
-
-    Returns:
-        A classification model
-    """
-    kwargs["num_classes"] = kwargs.get("num_classes", len(default_cfgs["resnet50"]["classes"]))
-    kwargs["input_shape"] = kwargs.get("input_shape", default_cfgs["resnet50"]["input_shape"])
-    kwargs["classes"] = kwargs.get("classes", default_cfgs["resnet50"]["classes"])
-
-    _cfg = deepcopy(default_cfgs["resnet50"])
-    _cfg["num_classes"] = kwargs["num_classes"]
-    _cfg["classes"] = kwargs["classes"]
-    _cfg["input_shape"] = kwargs["input_shape"]
-    kwargs.pop("classes")
-
-    model = ResNet50(
-        weights=None,
-        include_top=True,
-        pooling=True,
-        input_shape=kwargs["input_shape"],
-        classes=kwargs["num_classes"],
-        classifier_activation=None,
-    )
-
-    # monkeypatch the model to allow for loading pretrained parameters
-    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:  # noqa: D417
-        """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)
-
-    model.from_pretrained = types.MethodType(from_pretrained, model)  # Bind method to the instance
-
-    model.cfg = _cfg
-    _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["resnet50"]["url"],
-            skip_mismatch=kwargs["num_classes"] != len(default_cfgs["resnet50"]["classes"]),
-        )
-
-    return model
-
-
-def resnet34_wide(pretrained: bool = False, **kwargs: Any) -> ResNet:
-    """Resnet-34 architecture as described in `"Deep Residual Learning for Image Recognition",
-    <https://arxiv.org/pdf/1512.03385.pdf>`_ with twice as many output channels for each stage.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import resnet34_wide
-    >>> model = resnet34_wide(pretrained=False)
-    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained: boolean, True if model is pretrained
-        **kwargs: keyword arguments of the ResNet architecture
-
-    Returns:
-        A classification model
-    """
-    return _resnet(
-        "resnet34_wide",
-        pretrained,
-        [3, 4, 6, 3],
-        [128, 256, 512, 1024],
-        [False, True, True, True],
-        [False] * 4,
-        [None] * 4,
-        True,
-        stem_channels=128,
-        **kwargs,
-    )

doctr/models/classification/textnet/tensorflow.py

@@ -1,275 +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 Sequential, layers
-
-from doctr.datasets import VOCABS
-
-from ...modules.layers.tensorflow import FASTConvLayer
-from ...utils import _build_model, conv_sequence, load_pretrained_params
-
-__all__ = ["textnet_tiny", "textnet_small", "textnet_base"]
-
-default_cfgs: dict[str, dict[str, Any]] = {
-    "textnet_tiny": {
-        "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/textnet_tiny-a29eeb4a.weights.h5&src=0",
-    },
-    "textnet_small": {
-        "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/textnet_small-1c2df0e3.weights.h5&src=0",
-    },
-    "textnet_base": {
-        "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/textnet_base-8b4b89bc.weights.h5&src=0",
-    },
-}
-
-
-class TextNet(Sequential):
-    """Implements TextNet architecture from `"FAST: Faster Arbitrarily-Shaped Text Detector with
-    Minimalist Kernel Representation" <https://arxiv.org/abs/2111.02394>`_.
-    Implementation based on the official Pytorch implementation: <https://github.com/czczup/FAST>`_.
-
-    Args:
-        stages (list[dict[str, list[int]]]): list of dictionaries containing the parameters of each stage.
-        include_top (bool, optional): Whether to include the classifier head. Defaults to True.
-        num_classes (int, optional): Number of output classes. Defaults to 1000.
-        cfg (dict[str, Any], optional): Additional configuration. Defaults to None.
-    """
-
-    def __init__(
-        self,
-        stages: list[dict[str, list[int]]],
-        input_shape: tuple[int, int, int] = (32, 32, 3),
-        num_classes: int = 1000,
-        include_top: bool = True,
-        cfg: dict[str, Any] | None = None,
-    ) -> None:
-        _layers = [
-            *conv_sequence(
-                out_channels=64, activation="relu", bn=True, kernel_size=3, strides=2, input_shape=input_shape
-            ),
-            *[
-                Sequential(
-                    [
-                        FASTConvLayer(**params)  # type: ignore[arg-type]
-                        for params in [{key: stage[key][i] for key in stage} for i in range(len(stage["in_channels"]))]
-                    ],
-                    name=f"stage_{i}",
-                )
-                for i, stage in enumerate(stages)
-            ],
-        ]
-
-        if include_top:
-            _layers.append(
-                Sequential(
-                    [
-                        layers.AveragePooling2D(1),
-                        layers.Flatten(),
-                        layers.Dense(num_classes),
-                    ],
-                    name="classifier",
-                )
-            )
-
-        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 _textnet(
-    arch: str,
-    pretrained: bool,
-    **kwargs: Any,
-) -> TextNet:
-    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 = TextNet(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 textnet_tiny(pretrained: bool = False, **kwargs: Any) -> TextNet:
-    """Implements TextNet architecture from `"FAST: Faster Arbitrarily-Shaped Text Detector with
-    Minimalist Kernel Representation" <https://arxiv.org/abs/2111.02394>`_.
-    Implementation based on the official Pytorch implementation: <https://github.com/czczup/FAST>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import textnet_tiny
-    >>> model = textnet_tiny(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 TextNet architecture
-
-    Returns:
-        A textnet tiny model
-    """
-    return _textnet(
-        "textnet_tiny",
-        pretrained,
-        stages=[
-            {"in_channels": [64] * 3, "out_channels": [64] * 3, "kernel_size": [(3, 3)] * 3, "stride": [1, 2, 1]},
-            {
-                "in_channels": [64, 128, 128, 128],
-                "out_channels": [128] * 4,
-                "kernel_size": [(3, 3), (1, 3), (3, 3), (3, 1)],
-                "stride": [2, 1, 1, 1],
-            },
-            {
-                "in_channels": [128, 256, 256, 256],
-                "out_channels": [256] * 4,
-                "kernel_size": [(3, 3), (3, 3), (3, 1), (1, 3)],
-                "stride": [2, 1, 1, 1],
-            },
-            {
-                "in_channels": [256, 512, 512, 512],
-                "out_channels": [512] * 4,
-                "kernel_size": [(3, 3), (3, 1), (1, 3), (3, 3)],
-                "stride": [2, 1, 1, 1],
-            },
-        ],
-        **kwargs,
-    )
-
-
-def textnet_small(pretrained: bool = False, **kwargs: Any) -> TextNet:
-    """Implements TextNet architecture from `"FAST: Faster Arbitrarily-Shaped Text Detector with
-    Minimalist Kernel Representation" <https://arxiv.org/abs/2111.02394>`_.
-    Implementation based on the official Pytorch implementation: <https://github.com/czczup/FAST>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import textnet_small
-    >>> model = textnet_small(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 TextNet architecture
-
-    Returns:
-        A TextNet small model
-    """
-    return _textnet(
-        "textnet_small",
-        pretrained,
-        stages=[
-            {"in_channels": [64] * 2, "out_channels": [64] * 2, "kernel_size": [(3, 3)] * 2, "stride": [1, 2]},
-            {
-                "in_channels": [64, 128, 128, 128, 128, 128, 128, 128],
-                "out_channels": [128] * 8,
-                "kernel_size": [(3, 3), (1, 3), (3, 3), (3, 1), (3, 3), (3, 1), (1, 3), (3, 3)],
-                "stride": [2, 1, 1, 1, 1, 1, 1, 1],
-            },
-            {
-                "in_channels": [128, 256, 256, 256, 256, 256, 256, 256],
-                "out_channels": [256] * 8,
-                "kernel_size": [(3, 3), (3, 3), (1, 3), (3, 1), (3, 3), (1, 3), (3, 1), (3, 3)],
-                "stride": [2, 1, 1, 1, 1, 1, 1, 1],
-            },
-            {
-                "in_channels": [256, 512, 512, 512, 512],
-                "out_channels": [512] * 5,
-                "kernel_size": [(3, 3), (3, 1), (1, 3), (1, 3), (3, 1)],
-                "stride": [2, 1, 1, 1, 1],
-            },
-        ],
-        **kwargs,
-    )
-
-
-def textnet_base(pretrained: bool = False, **kwargs: Any) -> TextNet:
-    """Implements TextNet architecture from `"FAST: Faster Arbitrarily-Shaped Text Detector with
-    Minimalist Kernel Representation" <https://arxiv.org/abs/2111.02394>`_.
-    Implementation based on the official Pytorch implementation: <https://github.com/czczup/FAST>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import textnet_base
-    >>> model = textnet_base(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 TextNet architecture
-
-    Returns:
-        A TextNet base model
-    """
-    return _textnet(
-        "textnet_base",
-        pretrained,
-        stages=[
-            {
-                "in_channels": [64] * 10,
-                "out_channels": [64] * 10,
-                "kernel_size": [(3, 3), (3, 3), (3, 1), (3, 3), (3, 1), (3, 3), (3, 3), (1, 3), (3, 3), (3, 3)],
-                "stride": [1, 2, 1, 1, 1, 1, 1, 1, 1, 1],
-            },
-            {
-                "in_channels": [64, 128, 128, 128, 128, 128, 128, 128, 128, 128],
-                "out_channels": [128] * 10,
-                "kernel_size": [(3, 3), (1, 3), (3, 3), (3, 1), (3, 3), (3, 3), (3, 1), (3, 1), (3, 3), (3, 3)],
-                "stride": [2, 1, 1, 1, 1, 1, 1, 1, 1, 1],
-            },
-            {
-                "in_channels": [128, 256, 256, 256, 256, 256, 256, 256],
-                "out_channels": [256] * 8,
-                "kernel_size": [(3, 3), (3, 3), (3, 3), (1, 3), (3, 3), (3, 1), (3, 3), (3, 1)],
-                "stride": [2, 1, 1, 1, 1, 1, 1, 1],
-            },
-            {
-                "in_channels": [256, 512, 512, 512, 512],
-                "out_channels": [512] * 5,
-                "kernel_size": [(3, 3), (1, 3), (3, 1), (3, 1), (1, 3)],
-                "stride": [2, 1, 1, 1, 1],
-            },
-        ],
-        **kwargs,
-    )