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

doctr/models/classification/resnet/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.
@@ -61,6 +61,7 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 
 
 def resnet_stage(in_channels: int, out_channels: int, num_blocks: int, stride: int) -> List[nn.Module]:
+    """Build a ResNet stage"""
     _layers: List[nn.Module] = []
 
     in_chan = in_channels
@@ -83,6 +84,7 @@ class ResNet(nn.Sequential):
     Text Recognition" <https://arxiv.org/pdf/1811.00751.pdf>`_.
 
     Args:
+    ----
         num_blocks: number of resnet block in each stage
         output_channels: number of channels in each stage
         stage_conv: whether to add a conv_sequence after each stage
@@ -134,13 +136,11 @@ class ResNet(nn.Sequential):
             _layers.append(nn.Sequential(*_stage))
 
         if include_top:
-            _layers.extend(
-                [
-                    nn.AdaptiveAvgPool2d(1),
-                    nn.Flatten(1),
-                    nn.Linear(output_channels[-1], num_classes, bias=True),
-                ]
-            )
+            _layers.extend([
+                nn.AdaptiveAvgPool2d(1),
+                nn.Flatten(1),
+                nn.Linear(output_channels[-1], num_classes, bias=True),
+            ])
 
         super().__init__(*_layers)
         self.cfg = cfg
@@ -224,12 +224,14 @@ def resnet18(pretrained: bool = False, **kwargs: Any) -> TVResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A resnet18 model
     """
-
     return _tv_resnet(
         "resnet18",
         pretrained,
@@ -251,12 +253,14 @@ def resnet31(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A resnet31 model
     """
-
     return _resnet(
         "resnet31",
         pretrained,
@@ -283,12 +287,14 @@ def resnet34(pretrained: bool = False, **kwargs: Any) -> TVResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A resnet34 model
     """
-
     return _tv_resnet(
         "resnet34",
         pretrained,
@@ -309,12 +315,14 @@ def resnet34_wide(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A resnet34_wide model
     """
-
     return _resnet(
         "resnet34_wide",
         pretrained,
@@ -341,12 +349,14 @@ def resnet50(pretrained: bool = False, **kwargs: Any) -> TVResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A resnet50 model
     """
-
     return _tv_resnet(
         "resnet50",
         pretrained,

doctr/models/classification/resnet/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.
@@ -58,10 +58,10 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 
 
 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
@@ -71,19 +71,17 @@ class ResnetBlock(layers.Layer):
     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(),
-                ]
-            )
+            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))
@@ -123,6 +121,7 @@ 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
@@ -177,12 +176,10 @@ class ResNet(Sequential):
             inplanes = out_chan
 
         if include_top:
-            _layers.extend(
-                [
-                    layers.GlobalAveragePooling2D(),
-                    layers.Dense(num_classes),
-                ]
-            )
+            _layers.extend([
+                layers.GlobalAveragePooling2D(),
+                layers.Dense(num_classes),
+            ])
 
         super().__init__(_layers)
         self.cfg = cfg
@@ -231,12 +228,14 @@ def resnet18(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> 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,
@@ -262,12 +261,14 @@ def resnet31(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> 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,
@@ -293,12 +294,14 @@ def resnet34(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> 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,
@@ -323,12 +326,14 @@ def resnet50(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> 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"])
@@ -368,12 +373,14 @@ def resnet34_wide(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> 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,

doctr/models/classification/textnet/__init__.py

@@ -0,0 +1,6 @@
+from doctr.file_utils import is_tf_available, is_torch_available
+
+if is_tf_available():
+    from .tensorflow import *
+elif is_torch_available():
+    from .pytorch import *  # type: ignore[assignment]

doctr/models/classification/textnet/pytorch.py

@@ -0,0 +1,275 @@
+# 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 torch import nn
+
+from doctr.datasets import VOCABS
+
+from ...modules.layers.pytorch import FASTConvLayer
+from ...utils import conv_sequence_pt, 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": (3, 32, 32),
+        "classes": list(VOCABS["french"]),
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/textnet_tiny-c5970fe0.pt&src=0",
+    },
+    "textnet_small": {
+        "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.7.0/textnet_small-6e8ab0ce.pt&src=0",
+    },
+    "textnet_base": {
+        "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.7.0/textnet_base-8295dc85.pt&src=0",
+    },
+}
+
+
+class TextNet(nn.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] = (3, 32, 32),
+        num_classes: int = 1000,
+        include_top: bool = True,
+        cfg: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        _layers: List[nn.Module] = [
+            *conv_sequence_pt(
+                in_channels=3, out_channels=64, relu=True, bn=True, kernel_size=3, stride=2, padding=(1, 1)
+            ),
+            *[
+                nn.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"]))]
+                ])
+                for stage in stages
+            ],
+        ]
+
+        if include_top:
+            _layers.append(
+                nn.Sequential(
+                    nn.AdaptiveAvgPool2d(1),
+                    nn.Flatten(1),
+                    nn.Linear(stages[-1]["out_channels"][-1], num_classes),
+                )
+            )
+
+        super().__init__(*_layers)
+        self.cfg = cfg
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+
+def _textnet(
+    arch: str,
+    pretrained: bool,
+    ignore_keys: Optional[List[str]] = None,
+    **kwargs: Any,
+) -> TextNet:
+    kwargs["num_classes"] = kwargs.get("num_classes", len(default_cfgs[arch]["classes"]))
+    kwargs["classes"] = kwargs.get("classes", default_cfgs[arch]["classes"])
+
+    _cfg = deepcopy(default_cfgs[arch])
+    _cfg["num_classes"] = kwargs["num_classes"]
+    _cfg["classes"] = kwargs["classes"]
+    kwargs.pop("classes")
+
+    # Build the model
+    model = TextNet(**kwargs)
+    # Load pretrained parameters
+    if pretrained:
+        # The number of classes is not the same as the number of classes in the pretrained model =>
+        # remove the last layer weights
+        _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
+        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+
+    model.cfg = _cfg
+
+    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 torch
+    >>> from doctr.models import textnet_tiny
+    >>> model = textnet_tiny(pretrained=False)
+    >>> input_tensor = torch.rand((1, 3, 512, 512), dtype=torch.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],
+            },
+        ],
+        ignore_keys=["7.2.weight", "7.2.bias"],
+        **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 torch
+    >>> from doctr.models import textnet_small
+    >>> model = textnet_small(pretrained=False)
+    >>> input_tensor = torch.rand((1, 3, 512, 512), dtype=torch.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],
+            },
+        ],
+        ignore_keys=["7.2.weight", "7.2.bias"],
+        **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 torch
+    >>> from doctr.models import textnet_base
+    >>> model = textnet_base(pretrained=False)
+    >>> input_tensor = torch.rand((1, 3, 512, 512), dtype=torch.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],
+            },
+        ],
+        ignore_keys=["7.2.weight", "7.2.bias"],
+        **kwargs,
+    )