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

doctr/models/classification/textnet/tensorflow.py

@@ -0,0 +1,267 @@
+# 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.
+
+
+from copy import deepcopy
+from typing import Any, Dict, List, Optional, Tuple
+
+from tensorflow.keras import Sequential, layers
+
+from doctr.datasets import VOCABS
+
+from ...modules.layers.tensorflow import FASTConvLayer
+from ...utils import 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.7.0/textnet_tiny-9e605bd8.zip&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.7.0/textnet_small-4784b292.zip&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.7.0/textnet_base-2c3f3265.zip&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 (Optional[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: Optional[Dict[str, Any]] = 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 _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)
+    # Load pretrained parameters
+    if pretrained:
+        load_pretrained_params(model, default_cfgs[arch]["url"])
+
+    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,
+    )

doctr/models/classification/vgg/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.
@@ -77,12 +77,14 @@ def vgg16_bn_r(pretrained: bool = False, **kwargs: Any) -> tv_vgg.VGG:
     >>> 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,

doctr/models/classification/vgg/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.
@@ -32,6 +32,7 @@ class VGG(Sequential):
     <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
@@ -99,12 +100,14 @@ def vgg16_bn_r(pretrained: bool = False, **kwargs: Any) -> VGG:
     >>> 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/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.
@@ -40,6 +40,7 @@ class ClassifierHead(nn.Module):
     """Classifier head for Vision Transformer
 
     Args:
+    ----
         in_channels: number of input channels
         num_classes: number of output classes
     """
@@ -64,6 +65,7 @@ class VisionTransformer(nn.Sequential):
     <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
@@ -141,12 +143,14 @@ def vit_s(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
     >>> 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,
@@ -171,12 +175,14 @@ def vit_b(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
     >>> 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,

doctr/models/classification/vit/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.
@@ -41,6 +41,7 @@ class ClassifierHead(layers.Layer, NestedObject):
     """Classifier head for Vision Transformer
 
     Args:
+    ----
         num_classes: number of output classes
     """
 
@@ -60,6 +61,7 @@ class VisionTransformer(Sequential):
     <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
@@ -140,12 +142,14 @@ def vit_s(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
     >>> 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,
@@ -169,12 +173,14 @@ def vit_b(pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
     >>> 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,

doctr/models/classification/zoo.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.
@@ -24,6 +24,9 @@ ARCHS: List[str] = [
     "resnet34",
     "resnet50",
     "resnet34_wide",
+    "textnet_tiny",
+    "textnet_small",
+    "textnet_base",
     "vgg16_bn_r",
     "vit_s",
     "vit_b",
@@ -59,11 +62,13 @@ def crop_orientation_predictor(
     >>> out = model([input_crop])
 
     Args:
+    ----
         arch: name of the architecture to use (e.g. 'mobilenet_v3_small')
         pretrained: If True, returns a model pre-trained on our recognition crops dataset
+        **kwargs: keyword arguments to be passed to the CropOrientationPredictor
 
     Returns:
+    -------
         CropOrientationPredictor
     """
-
     return _crop_orientation_predictor(arch, pretrained, **kwargs)

doctr/models/core.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.

doctr/models/detection/__init__.py

@@ -1,3 +1,4 @@
 from .differentiable_binarization import *
 from .linknet import *
+from .fast import *
 from .zoo import *

doctr/models/detection/_utils/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.
@@ -13,9 +13,12 @@ def erode(x: Tensor, kernel_size: int) -> Tensor:
     """Performs erosion on a given tensor
 
     Args:
+    ----
         x: boolean tensor of shape (N, C, H, W)
         kernel_size: the size of the kernel to use for erosion
+
     Returns:
+    -------
         the eroded tensor
     """
     _pad = (kernel_size - 1) // 2
@@ -27,9 +30,12 @@ def dilate(x: Tensor, kernel_size: int) -> Tensor:
     """Performs dilation on a given tensor
 
     Args:
+    ----
         x: boolean tensor of shape (N, C, H, W)
         kernel_size: the size of the kernel to use for dilation
+
     Returns:
+    -------
         the dilated tensor
     """
     _pad = (kernel_size - 1) // 2

doctr/models/detection/_utils/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.
@@ -12,12 +12,14 @@ 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")
 
 
@@ -25,10 +27,12 @@ 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")

doctr/models/detection/core.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.
@@ -17,6 +17,7 @@ class DetectionPostProcessor(NestedObject):
     """Abstract class to postprocess the raw output of the model
 
     Args:
+    ----
         box_thresh (float): minimal objectness score to consider a box
         bin_thresh (float): threshold to apply to segmentation raw heatmap
         assume straight_pages (bool): if True, fit straight boxes only
@@ -36,9 +37,13 @@ class DetectionPostProcessor(NestedObject):
         """Compute the confidence score for a polygon : mean of the p values on the polygon
 
         Args:
+        ----
             pred (np.ndarray): p map returned by the model
+            points: coordinates of the polygon
+            assume_straight_pages: if True, fit straight boxes only
 
         Returns:
+        -------
             polygon objectness
         """
         h, w = pred.shape[:2]
@@ -52,7 +57,7 @@ class DetectionPostProcessor(NestedObject):
 
         else:
             mask: np.ndarray = np.zeros((h, w), np.int32)
-            cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)
+            cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)  # type: ignore[call-overload]
             product = pred * mask
             return np.sum(product) / np.count_nonzero(product)
 
@@ -70,13 +75,14 @@ class DetectionPostProcessor(NestedObject):
         """Performs postprocessing for a list of model outputs
 
         Args:
+        ----
             proba_map: probability map of shape (N, H, W, C)
 
         Returns:
+        -------
             list of N class predictions (for each input sample), where each class predictions is a list of C tensors
         of shape (*, 5) or (*, 6)
         """
-
         if proba_map.ndim != 4:
             raise AssertionError(f"arg `proba_map` is expected to be 4-dimensional, got {proba_map.ndim}.")