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

doctr/models/detection/differentiable_binarization/tensorflow.py

@@ -10,12 +10,17 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras import layers
+from tensorflow.keras import Model, Sequential, layers, losses
 from tensorflow.keras.applications import ResNet50
 
 from doctr.file_utils import CLASS_NAME
-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 ...classification import mobilenet_v3_large
@@ -29,13 +34,13 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
         "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/db_resnet50-84171458.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/db_resnet50-649fa22b.weights.h5&src=0",
     },
     "db_mobilenet_v3_large": {
         "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/db_mobilenet_v3_large-da524564.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/db_mobilenet_v3_large-ee2e1dbe.weights.h5&src=0",
     },
 }
 
@@ -81,7 +86,7 @@ class FeaturePyramidNetwork(layers.Layer, NestedObject):
         if dilation_factor > 1:
             _layers.append(layers.UpSampling2D(size=(dilation_factor, dilation_factor), interpolation="nearest"))
 
-        module = keras.Sequential(_layers)
+        module = Sequential(_layers)
 
         return module
 
@@ -104,7 +109,7 @@ class FeaturePyramidNetwork(layers.Layer, NestedObject):
         return layers.concatenate(results)
 
 
-class DBNet(_DBNet, keras.Model, NestedObject):
+class DBNet(_DBNet, Model, NestedObject):
     """DBNet as described in `"Real-time Scene Text Detection with Differentiable Binarization"
     <https://arxiv.org/pdf/1911.08947.pdf>`_.
 
@@ -147,14 +152,14 @@ class DBNet(_DBNet, keras.Model, NestedObject):
         _inputs = [layers.Input(shape=in_shape[1:]) for in_shape in self.feat_extractor.output_shape]
         output_shape = tuple(self.fpn(_inputs).shape)
 
-        self.probability_head = keras.Sequential([
+        self.probability_head = Sequential([
             *conv_sequence(64, "relu", True, kernel_size=3, input_shape=output_shape[1:]),
             layers.Conv2DTranspose(64, 2, strides=2, use_bias=False, kernel_initializer="he_normal"),
             layers.BatchNormalization(),
             layers.Activation("relu"),
             layers.Conv2DTranspose(num_classes, 2, strides=2, kernel_initializer="he_normal"),
         ])
-        self.threshold_head = keras.Sequential([
+        self.threshold_head = Sequential([
             *conv_sequence(64, "relu", True, kernel_size=3, input_shape=output_shape[1:]),
             layers.Conv2DTranspose(64, 2, strides=2, use_bias=False, kernel_initializer="he_normal"),
             layers.BatchNormalization(),
@@ -206,7 +211,7 @@ class DBNet(_DBNet, keras.Model, NestedObject):
 
         # Focal loss
         focal_scale = 10.0
-        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)
 
         # Convert logits to prob, compute gamma factor
         p_t = (seg_target * prob_map) + ((1 - seg_target) * (1 - prob_map))
@@ -305,9 +310,16 @@ def _db_resnet(
 
     # Build the model
     model = DBNet(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
 
@@ -326,6 +338,10 @@ def _db_mobilenet(
     # Patch the config
     _cfg = deepcopy(default_cfgs[arch])
     _cfg["input_shape"] = input_shape or _cfg["input_shape"]
+    if not kwargs.get("class_names", None):
+        kwargs["class_names"] = default_cfgs[arch].get("class_names", [CLASS_NAME])
+    else:
+        kwargs["class_names"] = sorted(kwargs["class_names"])
 
     # Feature extractor
     feat_extractor = IntermediateLayerGetter(
@@ -339,9 +355,15 @@ def _db_mobilenet(
 
     # Build the model
     model = DBNet(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
 

doctr/models/detection/fast/base.py

@@ -31,7 +31,7 @@ class FASTPostProcessor(DetectionPostProcessor):
 
     def __init__(
         self,
-        bin_thresh: float = 0.3,
+        bin_thresh: float = 0.1,
         box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
     ) -> None:
@@ -111,7 +111,7 @@ class FASTPostProcessor(DetectionPostProcessor):
         contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         for contour in contours:
             # Check whether smallest enclosing bounding box is not too small
-            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):  # type: ignore[index]
                 continue
             # Compute objectness
             if self.assume_straight_pages:
@@ -138,10 +138,11 @@ class FASTPostProcessor(DetectionPostProcessor):
                 # compute relative box to get rid of img shape
                 _box[:, 0] /= width
                 _box[:, 1] /= height
-                boxes.append(_box)
+                # Add score to box as (0, score)
+                boxes.append(np.vstack([_box, np.array([0.0, score])]))
 
         if not self.assume_straight_pages:
-            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 4, 2), dtype=pred.dtype)
+            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5, 2), dtype=pred.dtype)
         else:
             return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5), dtype=pred.dtype)
 
@@ -153,7 +154,7 @@ class _FAST(BaseModel):
 
     min_size_box: int = 3
     assume_straight_pages: bool = True
-    shrink_ratio = 0.1
+    shrink_ratio = 0.4
 
     def build_target(
         self,

doctr/models/detection/fast/pytorch.py

@@ -26,19 +26,19 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.8.1/fast_tiny-1acac421.pt&src=0",
     },
     "fast_small": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.8.1/fast_small-10952cc1.pt&src=0",
     },
     "fast_base": {
         "input_shape": (3, 1024, 1024),
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.8.1/fast_base-688a8b34.pt&src=0",
     },
 }
 
@@ -119,7 +119,7 @@ class FAST(_FAST, nn.Module):
     def __init__(
         self,
         feat_extractor: IntermediateLayerGetter,
-        bin_thresh: float = 0.3,
+        bin_thresh: float = 0.1,
         box_thresh: float = 0.1,
         dropout_prob: float = 0.1,
         pooling_size: int = 4,  # different from paper performs better on close text-rich images

doctr/models/detection/fast/tensorflow.py

@@ -10,11 +10,10 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras import Sequential, layers
+from tensorflow.keras import Model, Sequential, layers
 
 from doctr.file_utils import CLASS_NAME
-from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, load_pretrained_params
+from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, _build_model, load_pretrained_params
 from doctr.utils.repr import NestedObject
 
 from ...classification import textnet_base, textnet_small, textnet_tiny
@@ -29,19 +28,19 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
         "input_shape": (1024, 1024, 3),
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/fast_tiny-d7379d7b.weights.h5&src=0",
     },
     "fast_small": {
         "input_shape": (1024, 1024, 3),
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/fast_small-44b27eb6.weights.h5&src=0",
     },
     "fast_base": {
         "input_shape": (1024, 1024, 3),
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/fast_base-f2c6c736.weights.h5&src=0",
     },
 }
 
@@ -100,7 +99,7 @@ class FastHead(Sequential):
         super().__init__(_layers)
 
 
-class FAST(_FAST, keras.Model, NestedObject):
+class FAST(_FAST, Model, NestedObject):
     """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
     <https://arxiv.org/pdf/2111.02394.pdf>`_.
 
@@ -122,7 +121,7 @@ class FAST(_FAST, keras.Model, NestedObject):
     def __init__(
         self,
         feature_extractor: IntermediateLayerGetter,
-        bin_thresh: float = 0.3,
+        bin_thresh: float = 0.1,
         box_thresh: float = 0.1,
         dropout_prob: float = 0.1,
         pooling_size: int = 4,  # different from paper performs better on close text-rich images
@@ -334,12 +333,16 @@ def _fast(
 
     # Build the model
     model = FAST(feat_extractor, cfg=_cfg, **kwargs)
+    _build_model(model)
+
     # Load pretrained parameters
     if pretrained:
-        load_pretrained_params(model, _cfg["url"])
-
-    # Build the model for reparameterization to access the layers
-    _ = model(tf.random.uniform(shape=[1, *_cfg["input_shape"]], maxval=1, dtype=tf.float32), training=False)
+        # 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
 

doctr/models/detection/linknet/base.py

@@ -111,7 +111,7 @@ class LinkNetPostProcessor(DetectionPostProcessor):
         contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         for contour in contours:
             # Check whether smallest enclosing bounding box is not too small
-            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):  # type: ignore[index]
                 continue
             # Compute objectness
             if self.assume_straight_pages:
@@ -138,10 +138,11 @@ class LinkNetPostProcessor(DetectionPostProcessor):
                 # compute relative box to get rid of img shape
                 _box[:, 0] /= width
                 _box[:, 1] /= height
-                boxes.append(_box)
+                # Add score to box as (0, score)
+                boxes.append(np.vstack([_box, np.array([0.0, score])]))
 
         if not self.assume_straight_pages:
-            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 4, 2), dtype=pred.dtype)
+            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5, 2), dtype=pred.dtype)
         else:
             return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5), dtype=pred.dtype)
 

doctr/models/detection/linknet/tensorflow.py

@@ -10,12 +10,17 @@ from typing import Any, Dict, List, Optional, Tuple
 
 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
@@ -27,19 +32,19 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
         "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",
     },
 }
 
@@ -80,17 +85,17 @@ 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>`_.
 
@@ -187,7 +192,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
@@ -275,9 +280,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
 

doctr/models/detection/predictor/pytorch.py

@@ -9,6 +9,7 @@ import numpy as np
 import torch
 from torch import nn
 
+from doctr.models.detection._utils import _remove_padding
 from doctr.models.preprocessor import PreProcessor
 from doctr.models.utils import set_device_and_dtype
 
@@ -40,6 +41,11 @@ class DetectionPredictor(nn.Module):
         return_maps: bool = False,
         **kwargs: Any,
     ) -> Union[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
+        assume_straight_pages = self.model.assume_straight_pages
+
         # 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.")
@@ -52,7 +58,15 @@ class DetectionPredictor(nn.Module):
         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"]]
+        # Remove padding from loc predictions
+        preds = _remove_padding(
+            pages,  # type: ignore[arg-type]
+            [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,
+        )
+
         if return_maps:
             seg_maps = [
                 pred.permute(1, 2, 0).detach().cpu().numpy() for batch in predicted_batches for pred in batch["out_map"]

doctr/models/detection/predictor/tensorflow.py

@@ -7,8 +7,9 @@ from typing import Any, Dict, List, Tuple, Union
 
 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
 from doctr.utils.repr import NestedObject
 
@@ -29,7 +30,7 @@ class DetectionPredictor(NestedObject):
     def __init__(
         self,
         pre_processor: PreProcessor,
-        model: keras.Model,
+        model: Model,
     ) -> None:
         self.pre_processor = pre_processor
         self.model = model
@@ -40,6 +41,11 @@ class DetectionPredictor(NestedObject):
         return_maps: bool = False,
         **kwargs: Any,
     ) -> Union[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
+        assume_straight_pages = self.model.assume_straight_pages
+
         # 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.")
@@ -50,7 +56,15 @@ class DetectionPredictor(NestedObject):
             for batch in processed_batches
         ]
 
-        preds = [pred for batch in predicted_batches for pred in batch["preds"]]
+        # Remove padding from loc predictions
+        preds = _remove_padding(
+            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,
+        )
+
         if return_maps:
             seg_maps = [pred.numpy() for batch in predicted_batches for pred in batch["out_map"]]
             return preds, seg_maps

doctr/models/detection/zoo.py

@@ -8,6 +8,7 @@ from typing import Any, List
 from doctr.file_utils import is_tf_available, is_torch_available
 
 from .. import detection
+from ..detection.fast import reparameterize
 from ..preprocessor import PreProcessor
 from .predictor import DetectionPredictor
 
@@ -51,18 +52,22 @@ def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True,
             pretrained_backbone=kwargs.get("pretrained_backbone", True),
             assume_straight_pages=assume_straight_pages,
         )
+        # Reparameterize FAST models by default to lower inference latency and memory usage
+        if isinstance(_model, detection.FAST):
+            _model = reparameterize(_model)
     else:
         if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):
             raise ValueError(f"unknown architecture: {type(arch)}")
 
         _model = arch
         _model.assume_straight_pages = assume_straight_pages
+        _model.postprocessor.assume_straight_pages = assume_straight_pages
 
     kwargs.pop("pretrained_backbone", None)
 
     kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"])
     kwargs["std"] = kwargs.get("std", _model.cfg["std"])
-    kwargs["batch_size"] = kwargs.get("batch_size", 1)
+    kwargs["batch_size"] = kwargs.get("batch_size", 2)
     predictor = DetectionPredictor(
         PreProcessor(_model.cfg["input_shape"][:-1] if is_tf_available() else _model.cfg["input_shape"][1:], **kwargs),
         _model,
@@ -71,7 +76,7 @@ def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True,
 
 
 def detection_predictor(
-    arch: Any = "db_resnet50",
+    arch: Any = "fast_base",
     pretrained: bool = False,
     assume_straight_pages: bool = True,
     **kwargs: Any,

doctr/models/factory/hub.py

@@ -20,7 +20,6 @@ from huggingface_hub import (
     get_token_permission,
     hf_hub_download,
     login,
-    snapshot_download,
 )
 
 from doctr import models
@@ -33,10 +32,9 @@ __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,
-    "obj_detection": ["fasterrcnn_mobilenet_v3_large_fpn"] if is_torch_available() else None,
 }
 
 
@@ -75,7 +73,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"
@@ -110,8 +108,8 @@ def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:  #
 
     if run_config is None and arch is None:
         raise ValueError("run_config or arch must be specified")
-    if task not in ["classification", "detection", "recognition", "obj_detection"]:
-        raise ValueError("task must be one of classification, detection, recognition, obj_detection")
+    if task not in ["classification", "detection", "recognition"]:
+        raise ValueError("task must be one of classification, detection, recognition")
 
     # default readme
     readme = textwrap.dedent(
@@ -165,7 +163,7 @@ def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:  #
                                   \n{json.dumps(vars(run_config), indent=2, ensure_ascii=False)}"""
         )
 
-    if arch not in AVAILABLE_ARCHS[task]:  # type: ignore
+    if arch not in AVAILABLE_ARCHS[task]:
         raise ValueError(
             f"Architecture: {arch} for task: {task} not found.\
                          \nAvailable architectures: {AVAILABLE_ARCHS}"
@@ -175,7 +173,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)
@@ -217,14 +215,6 @@ def from_hub(repo_id: str, **kwargs: Any):
         model = models.detection.__dict__[arch](pretrained=False)
     elif task == "recognition":
         model = models.recognition.__dict__[arch](pretrained=False, input_shape=cfg["input_shape"], vocab=cfg["vocab"])
-    elif task == "obj_detection" and is_torch_available():
-        model = models.obj_detection.__dict__[arch](
-            pretrained=False,
-            image_mean=cfg["mean"],
-            image_std=cfg["std"],
-            max_size=cfg["input_shape"][-1],
-            num_classes=len(cfg["classes"]),
-        )
 
     # update model cfg
     model.cfg = cfg
@@ -234,7 +224,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

doctr/models/kie_predictor/base.py

@@ -7,7 +7,7 @@ from typing import Any, Optional
 
 from doctr.models.builder import KIEDocumentBuilder
 
-from ..classification.predictor import CropOrientationPredictor
+from ..classification.predictor import OrientationPredictor
 from ..predictor.base import _OCRPredictor
 
 __all__ = ["_KIEPredictor"]
@@ -25,10 +25,13 @@ class _KIEPredictor(_OCRPredictor):
             accordingly. Doing so will improve performances for documents with page-uniform rotations.
         preserve_aspect_ratio: if True, resize preserving the aspect ratio (with padding)
         symmetric_pad: if True and preserve_aspect_ratio is True, pas the image symmetrically.
+        detect_orientation: if True, the estimated general page orientation will be added to the predictions for each
+            page. Doing so will slightly deteriorate the overall latency.
         kwargs: keyword args of `DocumentBuilder`
     """
 
-    crop_orientation_predictor: Optional[CropOrientationPredictor]
+    crop_orientation_predictor: Optional[OrientationPredictor]
+    page_orientation_predictor: Optional[OrientationPredictor]
 
     def __init__(
         self,
@@ -36,8 +39,15 @@ class _KIEPredictor(_OCRPredictor):
         straighten_pages: bool = False,
         preserve_aspect_ratio: bool = True,
         symmetric_pad: bool = True,
+        detect_orientation: bool = False,
         **kwargs: Any,
     ) -> None:
-        super().__init__(assume_straight_pages, straighten_pages, preserve_aspect_ratio, symmetric_pad, **kwargs)
+        super().__init__(
+            assume_straight_pages, straighten_pages, preserve_aspect_ratio, symmetric_pad, detect_orientation, **kwargs
+        )
+
+        # Remove the following arguments from kwargs after initialization of the parent class
+        kwargs.pop("disable_page_orientation", None)
+        kwargs.pop("disable_crop_orientation", None)
 
         self.doc_builder: KIEDocumentBuilder = KIEDocumentBuilder(**kwargs)