python-doctr 0.9.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

doctr/models/detection/linknet/pytorch.py

@@ -1,9 +1,10 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 typing import Any, Callable, Dict, List, Optional, Tuple
+from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 import torch
@@ -20,7 +21,7 @@ from .base import LinkNetPostProcessor, _LinkNet
 __all__ = ["LinkNet", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
 
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "linknet_resnet18": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
@@ -43,7 +44,7 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
 
 
 class LinkNetFPN(nn.Module):
-    def __init__(self, layer_shapes: List[Tuple[int, int, int]]) -> None:
+    def __init__(self, layer_shapes: list[tuple[int, int, int]]) -> None:
         super().__init__()
         strides = [
             1 if (in_shape[-1] == out_shape[-1]) else 2
@@ -74,7 +75,7 @@ class LinkNetFPN(nn.Module):
             nn.ReLU(inplace=True),
         )
 
-    def forward(self, feats: List[torch.Tensor]) -> torch.Tensor:
+    def forward(self, feats: list[torch.Tensor]) -> torch.Tensor:
         out = feats[-1]
         for decoder, fmap in zip(self.decoders[::-1], feats[:-1][::-1]):
             out = decoder(out) + fmap
@@ -89,7 +90,6 @@ class LinkNet(nn.Module, _LinkNet):
     <https://arxiv.org/pdf/1707.03718.pdf>`_.
 
     Args:
-    ----
         feature extractor: the backbone serving as feature extractor
         bin_thresh: threshold for binarization of the output feature map
         box_thresh: minimal objectness score to consider a box
@@ -108,8 +108,8 @@ class LinkNet(nn.Module, _LinkNet):
         head_chans: int = 32,
         assume_straight_pages: bool = True,
         exportable: bool = False,
-        cfg: Optional[Dict[str, Any]] = None,
-        class_names: List[str] = [CLASS_NAME],
+        cfg: dict[str, Any] | None = None,
+        class_names: list[str] = [CLASS_NAME],
     ) -> None:
         super().__init__()
         self.class_names = class_names
@@ -163,16 +163,16 @@ class LinkNet(nn.Module, _LinkNet):
     def forward(
         self,
         x: torch.Tensor,
-        target: Optional[List[np.ndarray]] = None,
+        target: list[np.ndarray] | None = None,
         return_model_output: bool = False,
         return_preds: bool = False,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         feats = self.feat_extractor(x)
         logits = self.fpn([feats[str(idx)] for idx in range(len(feats))])
         logits = self.classifier(logits)
 
-        out: Dict[str, Any] = {}
+        out: dict[str, Any] = {}
         if self.exportable:
             out["logits"] = logits
             return out
@@ -183,11 +183,16 @@ class LinkNet(nn.Module, _LinkNet):
             out["out_map"] = prob_map
 
         if target is None or return_preds:
-            # Post-process boxes
-            out["preds"] = [
-                dict(zip(self.class_names, preds))
-                for preds in self.postprocessor(prob_map.detach().cpu().permute((0, 2, 3, 1)).numpy())
-            ]
+            # Disable for torch.compile compatibility
+            @torch.compiler.disable  # type: ignore[attr-defined]
+            def _postprocess(prob_map: torch.Tensor) -> list[dict[str, Any]]:
+                return [
+                    dict(zip(self.class_names, preds))
+                    for preds in self.postprocessor(prob_map.detach().cpu().permute((0, 2, 3, 1)).numpy())
+                ]
+
+            # Post-process boxes (keep only text predictions)
+            out["preds"] = _postprocess(prob_map)
 
         if target is not None:
             loss = self.compute_loss(logits, target)
@@ -198,7 +203,7 @@ class LinkNet(nn.Module, _LinkNet):
     def compute_loss(
         self,
         out_map: torch.Tensor,
-        target: List[np.ndarray],
+        target: list[np.ndarray],
         gamma: float = 2.0,
         alpha: float = 0.5,
         eps: float = 1e-8,
@@ -207,7 +212,6 @@ class LinkNet(nn.Module, _LinkNet):
         <https://github.com/tensorflow/addons/>`_.
 
         Args:
-        ----
             out_map: output feature map of the model of shape (N, num_classes, H, W)
             target: list of dictionary where each dict has a `boxes` and a `flags` entry
             gamma: modulating factor in the focal loss formula
@@ -215,7 +219,6 @@ class LinkNet(nn.Module, _LinkNet):
             eps: epsilon factor in dice loss
 
         Returns:
-        -------
             A loss tensor
         """
         _target, _mask = self.build_target(target, out_map.shape[1:], False)  # type: ignore[arg-type]
@@ -252,9 +255,9 @@ def _linknet(
     arch: str,
     pretrained: bool,
     backbone_fn: Callable[[bool], nn.Module],
-    fpn_layers: List[str],
+    fpn_layers: list[str],
     pretrained_backbone: bool = True,
-    ignore_keys: Optional[List[str]] = None,
+    ignore_keys: list[str] | None = None,
     **kwargs: Any,
 ) -> LinkNet:
     pretrained_backbone = pretrained_backbone and not pretrained
@@ -295,12 +298,10 @@ 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(
@@ -327,12 +328,10 @@ 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(
@@ -359,12 +358,10 @@ 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(

doctr/models/detection/linknet/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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.
@@ -6,40 +6,45 @@
 # Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
 
 from copy import deepcopy
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras import Model, Sequential, layers
+from tensorflow.keras import Model, Sequential, layers, losses
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.classification import resnet18, resnet34, resnet50
-from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, conv_sequence, load_pretrained_params
+from doctr.models.utils import (
+    IntermediateLayerGetter,
+    _bf16_to_float32,
+    _build_model,
+    conv_sequence,
+    load_pretrained_params,
+)
 from doctr.utils.repr import NestedObject
 
 from .base import LinkNetPostProcessor, _LinkNet
 
 __all__ = ["LinkNet", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+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.7.0/linknet_resnet18-b9ee56e6.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/linknet_resnet18-615a82c5.weights.h5&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.7.0/linknet_resnet34-51909c56.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/linknet_resnet34-9d772be5.weights.h5&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.7.0/linknet_resnet50-ac9f3829.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/linknet_resnet50-6bf6c8b5.weights.h5&src=0",
     },
 }
 
@@ -68,7 +73,7 @@ class LinkNetFPN(Model, NestedObject):
     def __init__(
         self,
         out_chans: int,
-        in_shapes: List[Tuple[int, ...]],
+        in_shapes: list[tuple[int, ...]],
     ) -> None:
         super().__init__()
         self.out_chans = out_chans
@@ -80,22 +85,21 @@ class LinkNetFPN(Model, NestedObject):
             for in_chan, out_chan, s, in_shape in zip(i_chans, o_chans, strides, in_shapes[::-1])
         ]
 
-    def call(self, x: List[tf.Tensor]) -> tf.Tensor:
+    def call(self, x: list[tf.Tensor], **kwargs: Any) -> tf.Tensor:
         out = 0
         for decoder, fmap in zip(self.decoders, x[::-1]):
-            out = decoder(out + fmap)
+            out = decoder(out + fmap, **kwargs)
         return out
 
     def extra_repr(self) -> str:
         return f"out_chans={self.out_chans}"
 
 
-class LinkNet(_LinkNet, keras.Model):
+class LinkNet(_LinkNet, Model):
     """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
     <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
@@ -106,7 +110,7 @@ class LinkNet(_LinkNet, keras.Model):
         class_names: list of class names
     """
 
-    _children_names: List[str] = ["feat_extractor", "fpn", "classifier", "postprocessor"]
+    _children_names: list[str] = ["feat_extractor", "fpn", "classifier", "postprocessor"]
 
     def __init__(
         self,
@@ -116,8 +120,8 @@ class LinkNet(_LinkNet, keras.Model):
         box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
         exportable: bool = False,
-        cfg: Optional[Dict[str, Any]] = None,
-        class_names: List[str] = [CLASS_NAME],
+        cfg: dict[str, Any] | None = None,
+        class_names: list[str] = [CLASS_NAME],
     ) -> None:
         super().__init__(cfg=cfg)
 
@@ -162,7 +166,7 @@ class LinkNet(_LinkNet, keras.Model):
     def compute_loss(
         self,
         out_map: tf.Tensor,
-        target: List[Dict[str, np.ndarray]],
+        target: list[dict[str, np.ndarray]],
         gamma: float = 2.0,
         alpha: float = 0.5,
         eps: float = 1e-8,
@@ -171,7 +175,6 @@ 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
@@ -179,7 +182,6 @@ 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)
@@ -187,7 +189,7 @@ class LinkNet(_LinkNet, keras.Model):
         seg_mask = tf.convert_to_tensor(seg_mask, dtype=tf.bool)
         seg_mask = tf.cast(seg_mask, tf.float32)
 
-        bce_loss = tf.keras.losses.binary_crossentropy(seg_target[..., None], out_map[..., None], from_logits=True)
+        bce_loss = losses.binary_crossentropy(seg_target[..., None], out_map[..., None], from_logits=True)
         proba_map = tf.sigmoid(out_map)
 
         # Focal loss
@@ -213,16 +215,16 @@ class LinkNet(_LinkNet, keras.Model):
     def call(
         self,
         x: tf.Tensor,
-        target: Optional[List[Dict[str, np.ndarray]]] = None,
+        target: list[dict[str, np.ndarray]] | None = None,
         return_model_output: bool = False,
         return_preds: bool = False,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         feat_maps = self.feat_extractor(x, **kwargs)
         logits = self.fpn(feat_maps, **kwargs)
         logits = self.classifier(logits, **kwargs)
 
-        out: Dict[str, tf.Tensor] = {}
+        out: dict[str, tf.Tensor] = {}
         if self.exportable:
             out["logits"] = logits
             return out
@@ -248,9 +250,9 @@ def _linknet(
     arch: str,
     pretrained: bool,
     backbone_fn,
-    fpn_layers: List[str],
+    fpn_layers: list[str],
     pretrained_backbone: bool = True,
-    input_shape: Optional[Tuple[int, int, int]] = None,
+    input_shape: tuple[int, int, int] | None = None,
     **kwargs: Any,
 ) -> LinkNet:
     pretrained_backbone = pretrained_backbone and not pretrained
@@ -275,9 +277,16 @@ def _linknet(
 
     # Build the model
     model = LinkNet(feat_extractor, cfg=_cfg, **kwargs)
+    _build_model(model)
+
     # Load pretrained parameters
     if pretrained:
-        load_pretrained_params(model, _cfg["url"])
+        # The given class_names differs from the pretrained model => skip the mismatching layers for fine tuning
+        load_pretrained_params(
+            model,
+            _cfg["url"],
+            skip_mismatch=kwargs["class_names"] != default_cfgs[arch].get("class_names", [CLASS_NAME]),
+        )
 
     return model
 
@@ -293,12 +302,10 @@ 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(
@@ -321,12 +328,10 @@ 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(
@@ -349,12 +354,10 @@ 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(

doctr/models/detection/predictor/__init__.py

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

doctr/models/detection/predictor/pytorch.py

@@ -1,9 +1,9 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -20,7 +20,6 @@ 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
     """
@@ -37,10 +36,10 @@ class DetectionPredictor(nn.Module):
     @torch.inference_mode()
     def forward(
         self,
-        pages: List[Union[np.ndarray, torch.Tensor]],
+        pages: list[np.ndarray | torch.Tensor],
         return_maps: bool = False,
         **kwargs: Any,
-    ) -> Union[List[Dict[str, np.ndarray]], Tuple[List[Dict[str, np.ndarray]], List[np.ndarray]]]:
+    ) -> list[dict[str, np.ndarray]] | tuple[list[dict[str, np.ndarray]], list[np.ndarray]]:
         # Extract parameters from the preprocessor
         preserve_aspect_ratio = self.pre_processor.resize.preserve_aspect_ratio
         symmetric_pad = self.pre_processor.resize.symmetric_pad
@@ -60,11 +59,11 @@ class DetectionPredictor(nn.Module):
         ]
         # Remove padding from loc predictions
         preds = _remove_padding(
-            pages,  # type: ignore[arg-type]
+            pages,
             [pred for batch in predicted_batches for pred in batch["preds"]],
             preserve_aspect_ratio=preserve_aspect_ratio,
             symmetric_pad=symmetric_pad,
-            assume_straight_pages=assume_straight_pages,
+            assume_straight_pages=assume_straight_pages,  # type: ignore[arg-type]
         )
 
         if return_maps:

doctr/models/detection/predictor/tensorflow.py

@@ -1,13 +1,13 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import keras
+from tensorflow.keras import Model
 
 from doctr.models.detection._utils import _remove_padding
 from doctr.models.preprocessor import PreProcessor
@@ -20,27 +20,26 @@ 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
     """
 
-    _children_names: List[str] = ["pre_processor", "model"]
+    _children_names: list[str] = ["pre_processor", "model"]
 
     def __init__(
         self,
         pre_processor: PreProcessor,
-        model: keras.Model,
+        model: Model,
     ) -> None:
         self.pre_processor = pre_processor
         self.model = model
 
     def __call__(
         self,
-        pages: List[Union[np.ndarray, tf.Tensor]],
+        pages: list[np.ndarray | tf.Tensor],
         return_maps: bool = False,
         **kwargs: Any,
-    ) -> Union[List[Dict[str, np.ndarray]], Tuple[List[Dict[str, np.ndarray]], List[np.ndarray]]]:
+    ) -> list[dict[str, np.ndarray]] | tuple[list[dict[str, np.ndarray]], list[np.ndarray]]:
         # Extract parameters from the preprocessor
         preserve_aspect_ratio = self.pre_processor.resize.preserve_aspect_ratio
         symmetric_pad = self.pre_processor.resize.symmetric_pad

doctr/models/detection/zoo.py

@@ -1,9 +1,9 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 typing import Any, List
+from typing import Any
 
 from doctr.file_utils import is_tf_available, is_torch_available
 
@@ -14,7 +14,7 @@ from .predictor import DetectionPredictor
 
 __all__ = ["detection_predictor"]
 
-ARCHS: List[str]
+ARCHS: list[str]
 
 
 if is_tf_available():
@@ -56,7 +56,14 @@ def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True,
         if isinstance(_model, detection.FAST):
             _model = reparameterize(_model)
     else:
-        if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):
+        allowed_archs = [detection.DBNet, detection.LinkNet, detection.FAST]
+        if is_torch_available():
+            # Adding the type for torch compiled models to the allowed architectures
+            from doctr.models.utils import _CompiledModule
+
+            allowed_archs.append(_CompiledModule)
+
+        if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")
 
         _model = arch
@@ -79,6 +86,9 @@ def detection_predictor(
     arch: Any = "fast_base",
     pretrained: bool = False,
     assume_straight_pages: bool = True,
+    preserve_aspect_ratio: bool = True,
+    symmetric_pad: bool = True,
+    batch_size: int = 2,
     **kwargs: Any,
 ) -> DetectionPredictor:
     """Text detection architecture.
@@ -90,14 +100,24 @@ 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
+        preserve_aspect_ratio: If True, pad the input document image to preserve the aspect ratio before
+            running the detection model on it
+        symmetric_pad: if True, pad the image symmetrically instead of padding at the bottom-right
+        batch_size: number of samples the model processes in parallel
         **kwargs: optional keyword arguments passed to the architecture
 
     Returns:
-    -------
         Detection predictor
     """
-    return _predictor(arch, pretrained, assume_straight_pages, **kwargs)
+    return _predictor(
+        arch=arch,
+        pretrained=pretrained,
+        assume_straight_pages=assume_straight_pages,
+        preserve_aspect_ratio=preserve_aspect_ratio,
+        symmetric_pad=symmetric_pad,
+        batch_size=batch_size,
+        **kwargs,
+    )

doctr/models/factory/hub.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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.
@@ -20,7 +20,6 @@ from huggingface_hub import (
     get_token_permission,
     hf_hub_download,
     login,
-    snapshot_download,
 )
 
 from doctr import models
@@ -33,7 +32,7 @@ __all__ = ["login_to_hub", "push_to_hf_hub", "from_hub", "_save_model_and_config
 
 
 AVAILABLE_ARCHS = {
-    "classification": models.classification.zoo.ARCHS,
+    "classification": models.classification.zoo.ARCHS + models.classification.zoo.ORIENTATION_ARCHS,
     "detection": models.detection.zoo.ARCHS,
     "recognition": models.recognition.zoo.ARCHS,
 }
@@ -62,7 +61,6 @@ 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
@@ -74,7 +72,7 @@ def _save_model_and_config_for_hf_hub(model: Any, save_dir: str, arch: str, task
         weights_path = save_directory / "pytorch_model.bin"
         torch.save(model.state_dict(), weights_path)
     elif is_tf_available():
-        weights_path = save_directory / "tf_model" / "weights"
+        weights_path = save_directory / "tf_model.weights.h5"
         model.save_weights(str(weights_path))
 
     config_path = save_directory / "config.json"
@@ -98,7 +96,6 @@ 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
@@ -115,9 +112,9 @@ def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:  #
     # default readme
     readme = textwrap.dedent(
         f"""
-    ---
+
     language: en
-    ---
+
 
     <p align="center">
     <img src="https://doctr-static.mindee.com/models?id=v0.3.1/Logo_doctr.gif&src=0" width="60%">
@@ -174,7 +171,7 @@ def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:  #
 
     local_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub", model_name)
     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)
+    repo = Repository(local_dir=local_cache_dir, clone_from=repo_url)
 
     with repo.commit(commit_message):
         _save_model_and_config_for_hf_hub(model, repo.local_dir, arch=arch, task=task)
@@ -191,12 +188,10 @@ 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
@@ -225,7 +220,7 @@ def from_hub(repo_id: str, **kwargs: Any):
         state_dict = torch.load(hf_hub_download(repo_id, filename="pytorch_model.bin", **kwargs), map_location="cpu")
         model.load_state_dict(state_dict)
     else:  # tf
-        repo_path = snapshot_download(repo_id, **kwargs)
-        model.load_weights(os.path.join(repo_path, "tf_model", "weights"))
+        weights = hf_hub_download(repo_id, filename="tf_model.weights.h5", **kwargs)
+        model.load_weights(weights)
 
     return model