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

doctr/models/detection/linknet/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.
@@ -15,60 +15,51 @@ from tensorflow.keras import Model, Sequential, layers
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.classification import resnet18, resnet34, resnet50
-from doctr.models.utils import IntermediateLayerGetter, conv_sequence, load_pretrained_params
+from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, conv_sequence, load_pretrained_params
 from doctr.utils.repr import NestedObject
 
 from .base import LinkNetPostProcessor, _LinkNet
 
-__all__ = ["LinkNet", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50", "linknet_resnet18_rotation"]
+__all__ = ["LinkNet", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
 
 default_cfgs: Dict[str, Dict[str, Any]] = {
     "linknet_resnet18": {
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
         "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.6.0/linknet_resnet18-611b50f2.zip&src=0",
-    },
-    "linknet_resnet18_rotation": {
-        "mean": (0.798, 0.785, 0.772),
-        "std": (0.264, 0.2749, 0.287),
-        "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.5.0/linknet_resnet18-a48e6ed3.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/linknet_resnet18-b9ee56e6.zip&src=0",
     },
     "linknet_resnet34": {
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
         "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.6.0/linknet_resnet34-bf30afb1.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/linknet_resnet34-51909c56.zip&src=0",
     },
     "linknet_resnet50": {
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
         "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.6.0/linknet_resnet50-cd299262.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/linknet_resnet50-ac9f3829.zip&src=0",
     },
 }
 
 
 def decoder_block(in_chan: int, out_chan: int, stride: int, **kwargs: Any) -> Sequential:
     """Creates a LinkNet decoder block"""
-
-    return Sequential(
-        [
-            *conv_sequence(in_chan // 4, "relu", True, kernel_size=1, **kwargs),
-            layers.Conv2DTranspose(
-                filters=in_chan // 4,
-                kernel_size=3,
-                strides=stride,
-                padding="same",
-                use_bias=False,
-                kernel_initializer="he_normal",
-            ),
-            layers.BatchNormalization(),
-            layers.Activation("relu"),
-            *conv_sequence(out_chan, "relu", True, kernel_size=1),
-        ]
-    )
+    return Sequential([
+        *conv_sequence(in_chan // 4, "relu", True, kernel_size=1, **kwargs),
+        layers.Conv2DTranspose(
+            filters=in_chan // 4,
+            kernel_size=3,
+            strides=stride,
+            padding="same",
+            use_bias=False,
+            kernel_initializer="he_normal",
+        ),
+        layers.BatchNormalization(),
+        layers.Activation("relu"),
+        *conv_sequence(out_chan, "relu", True, kernel_size=1),
+    ])
 
 
 class LinkNetFPN(Model, NestedObject):
@@ -104,8 +95,11 @@ class LinkNet(_LinkNet, keras.Model):
     <https://arxiv.org/pdf/1707.03718.pdf>`_.
 
     Args:
+    ----
         feature extractor: the backbone serving as feature extractor
         fpn_channels: number of channels each extracted feature maps is mapped to
+        bin_thresh: threshold for binarization of the output feature map
+        box_thresh: minimal objectness score to consider a box
         assume_straight_pages: if True, fit straight bounding boxes only
         exportable: onnx exportable returns only logits
         cfg: the configuration dict of the model
@@ -119,6 +113,7 @@ class LinkNet(_LinkNet, keras.Model):
         feat_extractor: IntermediateLayerGetter,
         fpn_channels: int = 64,
         bin_thresh: float = 0.1,
+        box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
         exportable: bool = False,
         cfg: Optional[Dict[str, Any]] = None,
@@ -137,32 +132,32 @@ class LinkNet(_LinkNet, keras.Model):
         self.fpn = LinkNetFPN(fpn_channels, [_shape[1:] for _shape in self.feat_extractor.output_shape])
         self.fpn.build(self.feat_extractor.output_shape)
 
-        self.classifier = Sequential(
-            [
-                layers.Conv2DTranspose(
-                    filters=32,
-                    kernel_size=3,
-                    strides=2,
-                    padding="same",
-                    use_bias=False,
-                    kernel_initializer="he_normal",
-                    input_shape=self.fpn.decoders[-1].output_shape[1:],
-                ),
-                layers.BatchNormalization(),
-                layers.Activation("relu"),
-                *conv_sequence(32, "relu", True, kernel_size=3, strides=1),
-                layers.Conv2DTranspose(
-                    filters=num_classes,
-                    kernel_size=2,
-                    strides=2,
-                    padding="same",
-                    use_bias=True,
-                    kernel_initializer="he_normal",
-                ),
-            ]
-        )
+        self.classifier = Sequential([
+            layers.Conv2DTranspose(
+                filters=32,
+                kernel_size=3,
+                strides=2,
+                padding="same",
+                use_bias=False,
+                kernel_initializer="he_normal",
+                input_shape=self.fpn.decoders[-1].output_shape[1:],
+            ),
+            layers.BatchNormalization(),
+            layers.Activation("relu"),
+            *conv_sequence(32, "relu", True, kernel_size=3, strides=1),
+            layers.Conv2DTranspose(
+                filters=num_classes,
+                kernel_size=2,
+                strides=2,
+                padding="same",
+                use_bias=True,
+                kernel_initializer="he_normal",
+            ),
+        ])
 
-        self.postprocessor = LinkNetPostProcessor(assume_straight_pages=assume_straight_pages, bin_thresh=bin_thresh)
+        self.postprocessor = LinkNetPostProcessor(
+            assume_straight_pages=assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh
+        )
 
     def compute_loss(
         self,
@@ -176,6 +171,7 @@ class LinkNet(_LinkNet, keras.Model):
         <https://github.com/tensorflow/addons/>`_.
 
         Args:
+        ----
             out_map: output feature map of the model of shape N x H x W x 1
             target: list of dictionary where each dict has a `boxes` and a `flags` entry
             gamma: modulating factor in the focal loss formula
@@ -183,6 +179,7 @@ class LinkNet(_LinkNet, keras.Model):
             eps: epsilon factor in dice loss
 
         Returns:
+        -------
             A loss tensor
         """
         seg_target, seg_mask = self.build_target(target, out_map.shape[1:], True)
@@ -204,10 +201,12 @@ class LinkNet(_LinkNet, keras.Model):
         # Class reduced
         focal_loss = tf.reduce_sum(seg_mask * focal_loss, (0, 1, 2, 3)) / tf.reduce_sum(seg_mask, (0, 1, 2, 3))
 
-        # Dice loss
-        inter = tf.math.reduce_sum(seg_mask * proba_map * seg_target, (0, 1, 2, 3))
-        cardinality = tf.math.reduce_sum((proba_map + seg_target), (0, 1, 2, 3))
-        dice_loss = 1 - 2 * (inter + eps) / (cardinality + eps)
+        # Compute dice loss for each class
+        dice_map = tf.nn.softmax(out_map, axis=-1) if len(self.class_names) > 1 else proba_map
+        # Class-reduced dice loss
+        inter = tf.reduce_sum(seg_mask * dice_map * seg_target, axis=[0, 1, 2])
+        cardinality = tf.reduce_sum(seg_mask * (dice_map + seg_target), axis=[0, 1, 2])
+        dice_loss = tf.reduce_mean(1 - 2 * inter / (cardinality + eps))
 
         return focal_loss + dice_loss
 
@@ -229,7 +228,8 @@ class LinkNet(_LinkNet, keras.Model):
             return out
 
         if return_model_output or target is None or return_preds:
-            prob_map = tf.math.sigmoid(logits)
+            prob_map = _bf16_to_float32(tf.math.sigmoid(logits))
+
         if return_model_output:
             out["out_map"] = prob_map
 
@@ -293,12 +293,14 @@ def linknet_resnet18(pretrained: bool = False, **kwargs: Any) -> LinkNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
+        **kwargs: keyword arguments of the LinkNet architecture
 
     Returns:
+    -------
         text detection architecture
     """
-
     return _linknet(
         "linknet_resnet18",
         pretrained,
@@ -308,32 +310,6 @@ def linknet_resnet18(pretrained: bool = False, **kwargs: Any) -> LinkNet:
     )
 
 
-def linknet_resnet18_rotation(pretrained: bool = False, **kwargs: Any) -> LinkNet:
-    """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
-    <https://arxiv.org/pdf/1707.03718.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import linknet_resnet18_rotation
-    >>> model = linknet_resnet18_rotation(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on our text detection dataset
-
-    Returns:
-        text detection architecture
-    """
-
-    return _linknet(
-        "linknet_resnet18_rotation",
-        pretrained,
-        resnet18,
-        ["resnet_block_1", "resnet_block_3", "resnet_block_5", "resnet_block_7"],
-        **kwargs,
-    )
-
-
 def linknet_resnet34(pretrained: bool = False, **kwargs: Any) -> LinkNet:
     """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
     <https://arxiv.org/pdf/1707.03718.pdf>`_.
@@ -345,12 +321,14 @@ def linknet_resnet34(pretrained: bool = False, **kwargs: Any) -> LinkNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
+        **kwargs: keyword arguments of the LinkNet architecture
 
     Returns:
+    -------
         text detection architecture
     """
-
     return _linknet(
         "linknet_resnet34",
         pretrained,
@@ -371,12 +349,14 @@ def linknet_resnet50(pretrained: bool = False, **kwargs: Any) -> LinkNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
+        **kwargs: keyword arguments of the LinkNet architecture
 
     Returns:
+    -------
         text detection architecture
     """
-
     return _linknet(
         "linknet_resnet50",
         pretrained,

doctr/models/detection/predictor/pytorch.py

@@ -1,9 +1,9 @@
-# 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.
 
-from typing import Any, List, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import numpy as np
 import torch
@@ -19,6 +19,7 @@ class DetectionPredictor(nn.Module):
     """Implements an object able to localize text elements in a document
 
     Args:
+    ----
         pre_processor: transform inputs for easier batched model inference
         model: core detection architecture
     """
@@ -32,12 +33,13 @@ class DetectionPredictor(nn.Module):
         self.pre_processor = pre_processor
         self.model = model.eval()
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         pages: List[Union[np.ndarray, torch.Tensor]],
+        return_maps: bool = False,
         **kwargs: Any,
-    ) -> List[np.ndarray]:
+    ) -> Union[List[Dict[str, np.ndarray]], Tuple[List[Dict[str, np.ndarray]], List[np.ndarray]]]:
         # Dimension check
         if any(page.ndim != 3 for page in pages):
             raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
@@ -47,5 +49,13 @@ class DetectionPredictor(nn.Module):
         self.model, processed_batches = set_device_and_dtype(
             self.model, processed_batches, _params.device, _params.dtype
         )
-        predicted_batches = [self.model(batch, return_preds=True, **kwargs)["preds"] for batch in processed_batches]
-        return [pred for batch in predicted_batches for pred in batch]
+        predicted_batches = [
+            self.model(batch, return_preds=True, return_model_output=True, **kwargs) for batch in processed_batches
+        ]
+        preds = [pred for batch in predicted_batches for pred in batch["preds"]]
+        if return_maps:
+            seg_maps = [
+                pred.permute(1, 2, 0).detach().cpu().numpy() for batch in predicted_batches for pred in batch["out_map"]
+            ]
+            return preds, seg_maps
+        return preds

doctr/models/detection/predictor/tensorflow.py

@@ -1,9 +1,9 @@
-# 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.
 
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import numpy as np
 import tensorflow as tf
@@ -19,6 +19,7 @@ class DetectionPredictor(NestedObject):
     """Implements an object able to localize text elements in a document
 
     Args:
+    ----
         pre_processor: transform inputs for easier batched model inference
         model: core detection architecture
     """
@@ -36,14 +37,21 @@ class DetectionPredictor(NestedObject):
     def __call__(
         self,
         pages: List[Union[np.ndarray, tf.Tensor]],
+        return_maps: bool = False,
         **kwargs: Any,
-    ) -> List[Dict[str, np.ndarray]]:
+    ) -> Union[List[Dict[str, np.ndarray]], Tuple[List[Dict[str, np.ndarray]], List[np.ndarray]]]:
         # Dimension check
         if any(page.ndim != 3 for page in pages):
             raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
 
         processed_batches = self.pre_processor(pages)
         predicted_batches = [
-            self.model(batch, return_preds=True, training=False, **kwargs)["preds"] for batch in processed_batches
+            self.model(batch, return_preds=True, return_model_output=True, training=False, **kwargs)
+            for batch in processed_batches
         ]
-        return [pred for batch in predicted_batches for pred in batch]
+
+        preds = [pred for batch in predicted_batches for pred in batch["preds"]]
+        if return_maps:
+            seg_maps = [pred.numpy() for batch in predicted_batches for pred in batch["out_map"]]
+            return preds, seg_maps
+        return preds

doctr/models/detection/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.
@@ -14,12 +14,19 @@ from .predictor import DetectionPredictor
 __all__ = ["detection_predictor"]
 
 ARCHS: List[str]
-ROT_ARCHS: List[str]
 
 
 if is_tf_available():
-    ARCHS = ["db_resnet50", "db_mobilenet_v3_large", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
-    ROT_ARCHS = ["linknet_resnet18_rotation"]
+    ARCHS = [
+        "db_resnet50",
+        "db_mobilenet_v3_large",
+        "linknet_resnet18",
+        "linknet_resnet34",
+        "linknet_resnet50",
+        "fast_tiny",
+        "fast_small",
+        "fast_base",
+    ]
 elif is_torch_available():
     ARCHS = [
         "db_resnet34",
@@ -28,30 +35,24 @@ elif is_torch_available():
         "linknet_resnet18",
         "linknet_resnet34",
         "linknet_resnet50",
+        "fast_tiny",
+        "fast_small",
+        "fast_base",
     ]
-    ROT_ARCHS = ["db_resnet50_rotation"]
 
 
 def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True, **kwargs: Any) -> DetectionPredictor:
     if isinstance(arch, str):
-        if arch not in ARCHS + ROT_ARCHS:
+        if arch not in ARCHS:
             raise ValueError(f"unknown architecture '{arch}'")
 
-        if arch not in ROT_ARCHS and not assume_straight_pages:
-            raise AssertionError(
-                "You are trying to use a model trained on straight pages while not assuming"
-                " your pages are straight. If you have only straight documents, don't pass"
-                " assume_straight_pages=False, otherwise you should use one of these archs:"
-                f"{ROT_ARCHS}"
-            )
-
         _model = detection.__dict__[arch](
             pretrained=pretrained,
             pretrained_backbone=kwargs.get("pretrained_backbone", True),
             assume_straight_pages=assume_straight_pages,
         )
     else:
-        if not isinstance(arch, (detection.DBNet, detection.LinkNet)):
+        if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):
             raise ValueError(f"unknown architecture: {type(arch)}")
 
         _model = arch
@@ -84,12 +85,14 @@ def detection_predictor(
     >>> out = model([input_page])
 
     Args:
+    ----
         arch: name of the architecture or model itself to use (e.g. 'db_resnet50')
         pretrained: If True, returns a model pre-trained on our text detection dataset
         assume_straight_pages: If True, fit straight boxes to the page
+        **kwargs: optional keyword arguments passed to the architecture
 
     Returns:
+    -------
         Detection predictor
     """
-
     return _predictor(arch, pretrained, assume_straight_pages, **kwargs)

doctr/models/factory/hub.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,7 +13,15 @@ import textwrap
 from pathlib import Path
 from typing import Any
 
-from huggingface_hub import HfApi, HfFolder, Repository, hf_hub_download, snapshot_download
+from huggingface_hub import (
+    HfApi,
+    Repository,
+    get_token,
+    get_token_permission,
+    hf_hub_download,
+    login,
+    snapshot_download,
+)
 
 from doctr import models
 from doctr.file_utils import is_tf_available, is_torch_available
@@ -26,7 +34,7 @@ __all__ = ["login_to_hub", "push_to_hf_hub", "from_hub", "_save_model_and_config
 
 AVAILABLE_ARCHS = {
     "classification": models.classification.zoo.ARCHS,
-    "detection": models.detection.zoo.ARCHS + models.detection.zoo.ROT_ARCHS,
+    "detection": models.detection.zoo.ARCHS,
     "recognition": models.recognition.zoo.ARCHS,
     "obj_detection": ["fasterrcnn_mobilenet_v3_large_fpn"] if is_torch_available() else None,
 }
@@ -34,13 +42,12 @@ AVAILABLE_ARCHS = {
 
 def login_to_hub() -> None:  # pragma: no cover
     """Login to huggingface hub"""
-    access_token = HfFolder.get_token()
-    if access_token is not None and HfApi()._is_valid_token(access_token):
+    access_token = get_token()
+    if access_token is not None and get_token_permission(access_token):
         logging.info("Huggingface Hub token found and valid")
-        HfApi().set_access_token(access_token)
+        login(token=access_token, write_permission=True)
     else:
-        subprocess.call(["huggingface-cli", "login"])
-        HfApi().set_access_token(HfFolder().get_token())
+        login()
     # check if git lfs is installed
     try:
         subprocess.call(["git", "lfs", "version"])
@@ -56,6 +63,7 @@ def _save_model_and_config_for_hf_hub(model: Any, save_dir: str, arch: str, task
     """Save model and config to disk for pushing to huggingface hub
 
     Args:
+    ----
         model: TF or PyTorch model to be saved
         save_dir: directory to save model and config
         arch: architecture name
@@ -91,6 +99,7 @@ def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:  #
     >>> push_to_hf_hub(model, 'my-model', 'recognition', arch='crnn_mobilenet_v3_small')
 
     Args:
+    ----
         model: TF or PyTorch model to be saved
         model_name: name of the model which is also the repository name
         task: task name
@@ -165,7 +174,7 @@ def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:  #
     commit_message = f"Add {model_name} model"
 
     local_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub", model_name)
-    repo_url = HfApi().create_repo(model_name, token=HfFolder.get_token(), exist_ok=False)
+    repo_url = HfApi().create_repo(model_name, token=get_token(), exist_ok=False)
     repo = Repository(local_dir=local_cache_dir, clone_from=repo_url, use_auth_token=True)
 
     with repo.commit(commit_message):
@@ -183,13 +192,14 @@ def from_hub(repo_id: str, **kwargs: Any):
     >>> model = from_hub("mindee/fasterrcnn_mobilenet_v3_large_fpn")
 
     Args:
+    ----
         repo_id: HuggingFace model hub repo
         kwargs: kwargs of `hf_hub_download` or `snapshot_download`
 
     Returns:
+    -------
         Model loaded with the checkpoint
     """
-
     # Get the config
     with open(hf_hub_download(repo_id, filename="config.json", **kwargs), "rb") as f:
         cfg = json.load(f)

doctr/models/kie_predictor/base.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 _KIEPredictor(_OCRPredictor):
     """Implements an object able to localize and identify text elements in a set of documents
 
     Args:
+    ----
         assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
             without rotated textual elements.
         straighten_pages: if True, estimates the page general orientation based on the median line orientation.