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

doctr/models/modules/transformer/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.
@@ -42,12 +42,15 @@ class PositionalEncoding(layers.Layer, NestedObject):
         x: tf.Tensor,
         **kwargs: Any,
     ) -> tf.Tensor:
-        """
+        """Forward pass
+
         Args:
+        ----
             x: embeddings (batch, max_len, d_model)
             **kwargs: additional arguments
 
-        Returns:
+        Returns
+        -------
             positional embeddings (batch, max_len, d_model)
         """
         if x.dtype == tf.float16:  # amp fix: cast to half
@@ -62,7 +65,6 @@ def scaled_dot_product_attention(
     query: tf.Tensor, key: tf.Tensor, value: tf.Tensor, mask: Optional[tf.Tensor] = None
 ) -> Tuple[tf.Tensor, tf.Tensor]:
     """Scaled Dot-Product Attention"""
-
     scores = tf.matmul(query, tf.transpose(key, perm=[0, 1, 3, 2])) / math.sqrt(query.shape[-1])
     if mask is not None:
         # NOTE: to ensure the ONNX compatibility, tf.where works only with bool type condition

doctr/models/modules/vision_transformer/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.
@@ -28,8 +28,7 @@ class PatchEmbedding(nn.Module):
         self.projection = nn.Conv2d(channels, embed_dim, kernel_size=self.patch_size, stride=self.patch_size)
 
     def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
-        """
-        100 % borrowed from:
+        """100 % borrowed from:
         https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_vit.py
 
         This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
@@ -38,7 +37,6 @@ class PatchEmbedding(nn.Module):
         Source:
         https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py
         """
-
         num_patches = embeddings.shape[1] - 1
         num_positions = self.positions.shape[1] - 1
         if num_patches == num_positions and height == width:

doctr/models/modules/vision_transformer/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.
@@ -45,8 +45,7 @@ class PatchEmbedding(layers.Layer, NestedObject):
         )
 
     def interpolate_pos_encoding(self, embeddings: tf.Tensor, height: int, width: int) -> tf.Tensor:
-        """
-        100 % borrowed from:
+        """100 % borrowed from:
         https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_tf_vit.py
 
         This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
@@ -55,7 +54,6 @@ class PatchEmbedding(layers.Layer, NestedObject):
         Source:
         https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py
         """
-
         seq_len, dim = embeddings.shape[1:]
         num_patches = seq_len - 1
 

doctr/models/obj_detection/faster_rcnn/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.
@@ -70,10 +70,12 @@ def fasterrcnn_mobilenet_v3_large_fpn(pretrained: bool = False, **kwargs: Any) -
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our object detection dataset
+        **kwargs: keyword arguments of the FasterRCNN architecture
 
     Returns:
+    -------
         object detection architecture
     """
-
     return _fasterrcnn("fasterrcnn_mobilenet_v3_large_fpn", pretrained, **kwargs)

doctr/models/predictor/base.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, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import numpy as np
 
@@ -21,6 +21,7 @@ class _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.
@@ -28,7 +29,7 @@ class _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.
-        kwargs: keyword args of `DocumentBuilder`
+        **kwargs: keyword args of `DocumentBuilder`
     """
 
     crop_orientation_predictor: Optional[CropOrientationPredictor]
@@ -47,6 +48,7 @@ class _OCRPredictor:
         self.doc_builder = DocumentBuilder(**kwargs)
         self.preserve_aspect_ratio = preserve_aspect_ratio
         self.symmetric_pad = symmetric_pad
+        self.hooks: List[Callable] = []
 
     @staticmethod
     def _generate_crops(
@@ -148,3 +150,12 @@ class _OCRPredictor:
                 _idx += page_boxes.shape[0]
 
         return loc_preds, text_preds
+
+    def add_hook(self, hook: Callable) -> None:
+        """Add a hook to the predictor
+
+        Args:
+        ----
+            hook: a callable that takes as input the `loc_preds` and returns the modified `loc_preds`
+        """
+        self.hooks.append(hook)

doctr/models/predictor/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, Mindee.
 
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
@@ -13,7 +13,7 @@ from doctr.io.elements import Document
 from doctr.models._utils import estimate_orientation, get_language
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
-from doctr.utils.geometry import rotate_boxes, rotate_image
+from doctr.utils.geometry import rotate_image
 
 from .base import _OCRPredictor
 
@@ -24,6 +24,7 @@ class OCRPredictor(nn.Module, _OCRPredictor):
     """Implements an object able to localize and identify text elements in a set of documents
 
     Args:
+    ----
         det_predictor: detection module
         reco_predictor: recognition module
         assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
@@ -35,7 +36,7 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             page. Doing so will slightly deteriorate the overall latency.
         detect_language: if True, the language prediction will be added to the predictions for each
             page. Doing so will slightly deteriorate the overall latency.
-        kwargs: keyword args of `DocumentBuilder`
+        **kwargs: keyword args of `DocumentBuilder`
     """
 
     def __init__(
@@ -59,7 +60,7 @@ class OCRPredictor(nn.Module, _OCRPredictor):
         self.detect_orientation = detect_orientation
         self.detect_language = detect_language
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         pages: List[Union[np.ndarray, torch.Tensor]],
@@ -71,11 +72,18 @@ class OCRPredictor(nn.Module, _OCRPredictor):
 
         origin_page_shapes = [page.shape[:2] if isinstance(page, np.ndarray) else page.shape[-2:] for page in pages]
 
+        # Localize text elements
+        loc_preds, out_maps = self.det_predictor(pages, return_maps=True, **kwargs)
+
         # Detect document rotation and rotate pages
+        seg_maps = [
+            np.where(out_map > getattr(self.det_predictor.model.postprocessor, "bin_thresh"), 255, 0).astype(np.uint8)
+            for out_map in out_maps
+        ]
         if self.detect_orientation:
-            origin_page_orientations = [estimate_orientation(page) for page in pages]  # type: ignore[arg-type]
+            origin_page_orientations = [estimate_orientation(seq_map) for seq_map in seg_maps]
             orientations = [
-                {"value": orientation_page, "confidence": 1.0} for orientation_page in origin_page_orientations
+                {"value": orientation_page, "confidence": None} for orientation_page in origin_page_orientations
             ]
         else:
             orientations = None
@@ -83,15 +91,12 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             origin_page_orientations = (
                 origin_page_orientations
                 if self.detect_orientation
-                else [estimate_orientation(page) for page in pages]  # type: ignore[arg-type]
+                else [estimate_orientation(seq_map) for seq_map in seg_maps]
             )
-            pages = [
-                rotate_image(page, -angle, expand=True)  # type: ignore[arg-type]
-                for page, angle in zip(pages, origin_page_orientations)
-            ]
+            pages = [rotate_image(page, -angle, expand=False) for page, angle in zip(pages, origin_page_orientations)]
+            # Forward again to get predictions on straight pages
+            loc_preds = self.det_predictor(pages, **kwargs)
 
-        # Localize text elements
-        loc_preds = self.det_predictor(pages, **kwargs)
         assert all(
             len(loc_pred) == 1 for loc_pred in loc_preds
         ), "Detection Model in ocr_predictor should output only one class"
@@ -101,11 +106,15 @@ class OCRPredictor(nn.Module, _OCRPredictor):
         channels_last = len(pages) == 0 or isinstance(pages[0], np.ndarray)
 
         # Rectify crops if aspect ratio
-        loc_preds = self._remove_padding(pages, loc_preds)  # type: ignore[arg-type]
+        loc_preds = self._remove_padding(pages, loc_preds)
+
+        # Apply hooks to loc_preds if any
+        for hook in self.hooks:
+            loc_preds = hook(loc_preds)
 
         # Crop images
         crops, loc_preds = self._prepare_crops(
-            pages,  # type: ignore[arg-type]
+            pages,
             loc_preds,
             channels_last=channels_last,
             assume_straight_pages=self.assume_straight_pages,
@@ -123,24 +132,12 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             languages_dict = [{"value": lang[0], "confidence": lang[1]} for lang in languages]
         else:
             languages_dict = None
-        # Rotate back pages and boxes while keeping original image size
-        if self.straighten_pages:
-            boxes = [
-                rotate_boxes(
-                    page_boxes,
-                    angle,
-                    orig_shape=page.shape[:2]
-                    if isinstance(page, np.ndarray)
-                    else page.shape[1:],  # type: ignore[arg-type]
-                    target_shape=mask,  # type: ignore[arg-type]
-                )
-                for page_boxes, page, angle, mask in zip(boxes, pages, origin_page_orientations, origin_page_shapes)
-            ]
 
         out = self.doc_builder(
+            pages,
             boxes,
             text_preds,
-            [page.shape[:2] if channels_last else page.shape[-2:] for page in pages],  # type: ignore[misc]
+            origin_page_shapes,
             orientations,
             languages_dict,
         )

doctr/models/predictor/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, Mindee.
 
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
@@ -12,7 +12,7 @@ from doctr.io.elements import Document
 from doctr.models._utils import estimate_orientation, get_language
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
-from doctr.utils.geometry import rotate_boxes, rotate_image
+from doctr.utils.geometry import rotate_image
 from doctr.utils.repr import NestedObject
 
 from .base import _OCRPredictor
@@ -24,6 +24,7 @@ class OCRPredictor(NestedObject, _OCRPredictor):
     """Implements an object able to localize and identify text elements in a set of documents
 
     Args:
+    ----
         det_predictor: detection module
         reco_predictor: recognition module
         assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
@@ -35,7 +36,7 @@ class OCRPredictor(NestedObject, _OCRPredictor):
             page. Doing so will slightly deteriorate the overall latency.
         detect_language: if True, the language prediction will be added to the predictions for each
             page. Doing so will slightly deteriorate the overall latency.
-        kwargs: keyword args of `DocumentBuilder`
+        **kwargs: keyword args of `DocumentBuilder`
     """
 
     _children_names = ["det_predictor", "reco_predictor", "doc_builder"]
@@ -71,31 +72,43 @@ class OCRPredictor(NestedObject, _OCRPredictor):
 
         origin_page_shapes = [page.shape[:2] for page in pages]
 
+        # Localize text elements
+        loc_preds_dict, out_maps = self.det_predictor(pages, return_maps=True, **kwargs)
+
         # Detect document rotation and rotate pages
+        seg_maps = [
+            np.where(out_map > getattr(self.det_predictor.model.postprocessor, "bin_thresh"), 255, 0).astype(np.uint8)
+            for out_map in out_maps
+        ]
         if self.detect_orientation:
-            origin_page_orientations = [estimate_orientation(page) for page in pages]
+            origin_page_orientations = [estimate_orientation(seq_map) for seq_map in seg_maps]
             orientations = [
-                {"value": orientation_page, "confidence": 1.0} for orientation_page in origin_page_orientations
+                {"value": orientation_page, "confidence": None} for orientation_page in origin_page_orientations
             ]
         else:
             orientations = None
         if self.straighten_pages:
             origin_page_orientations = (
-                origin_page_orientations if self.detect_orientation else [estimate_orientation(page) for page in pages]
+                origin_page_orientations
+                if self.detect_orientation
+                else [estimate_orientation(seq_map) for seq_map in seg_maps]
             )
-            pages = [rotate_image(page, -angle, expand=True) for page, angle in zip(pages, origin_page_orientations)]
+            pages = [rotate_image(page, -angle, expand=False) for page, angle in zip(pages, origin_page_orientations)]
+            # forward again to get predictions on straight pages
+            loc_preds_dict = self.det_predictor(pages, **kwargs)  # type: ignore[assignment]
 
-        # Localize text elements
-        loc_preds_dict = self.det_predictor(pages, **kwargs)
         assert all(
             len(loc_pred) == 1 for loc_pred in loc_preds_dict
         ), "Detection Model in ocr_predictor should output only one class"
-
-        loc_preds: List[np.ndarray] = [list(loc_pred.values())[0] for loc_pred in loc_preds_dict]
+        loc_preds: List[np.ndarray] = [list(loc_pred.values())[0] for loc_pred in loc_preds_dict]  # type: ignore[union-attr]
 
         # Rectify crops if aspect ratio
         loc_preds = self._remove_padding(pages, loc_preds)
 
+        # Apply hooks to loc_preds if any
+        for hook in self.hooks:
+            loc_preds = hook(loc_preds)
+
         # Crop images
         crops, loc_preds = self._prepare_crops(
             pages, loc_preds, channels_last=True, assume_straight_pages=self.assume_straight_pages
@@ -114,19 +127,9 @@ class OCRPredictor(NestedObject, _OCRPredictor):
             languages_dict = [{"value": lang[0], "confidence": lang[1]} for lang in languages]
         else:
             languages_dict = None
-        # Rotate back pages and boxes while keeping original image size
-        if self.straighten_pages:
-            boxes = [
-                rotate_boxes(
-                    page_boxes,
-                    angle,
-                    orig_shape=page.shape[:2] if isinstance(page, np.ndarray) else page.shape[-2:],
-                    target_shape=mask,  # type: ignore[arg-type]
-                )
-                for page_boxes, page, angle, mask in zip(boxes, pages, origin_page_orientations, origin_page_shapes)
-            ]
 
         out = self.doc_builder(
+            pages,
             boxes,
             text_preds,
             origin_page_shapes,  # type: ignore[arg-type]

doctr/models/preprocessor/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.
@@ -22,6 +22,7 @@ class PreProcessor(nn.Module):
     """Implements an abstract preprocessor object which performs casting, resizing, batching and normalization.
 
     Args:
+    ----
         output_size: expected size of each page in format (H, W)
         batch_size: the size of page batches
         mean: mean value of the training distribution by channel
@@ -34,7 +35,6 @@ class PreProcessor(nn.Module):
         batch_size: int,
         mean: Tuple[float, float, float] = (0.5, 0.5, 0.5),
         std: Tuple[float, float, float] = (1.0, 1.0, 1.0),
-        fp16: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__()
@@ -47,12 +47,13 @@ class PreProcessor(nn.Module):
         """Gather samples into batches for inference purposes
 
         Args:
+        ----
             samples: list of samples of shape (C, H, W)
 
         Returns:
+        -------
             list of batched samples (*, C, H, W)
         """
-
         num_batches = int(math.ceil(len(samples) / self.batch_size))
         batches = [
             torch.stack(samples[idx * self.batch_size : min((idx + 1) * self.batch_size, len(samples))], dim=0)
@@ -78,17 +79,19 @@ class PreProcessor(nn.Module):
         else:
             x = x.to(dtype=torch.float32)  # type: ignore[union-attr]
 
-        return x
+        return x  # type: ignore[return-value]
 
     def __call__(self, x: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]) -> List[torch.Tensor]:
         """Prepare document data for model forwarding
 
         Args:
+        ----
             x: list of images (np.array) or tensors (already resized and batched)
+
         Returns:
+        -------
             list of page batches
         """
-
         # Input type check
         if isinstance(x, (np.ndarray, torch.Tensor)):
             if x.ndim != 4:
@@ -100,8 +103,10 @@ class PreProcessor(nn.Module):
             elif x.dtype not in (torch.uint8, torch.float16, torch.float32):
                 raise TypeError("unsupported data type for torch.Tensor")
             # Resizing
-            if x.shape[-2] != self.resize.size[0] or x.shape[-1] != self.resize.size[1]:
-                x = F.resize(x, self.resize.size, interpolation=self.resize.interpolation)
+            if x.shape[-2] != self.resize.size[0] or x.shape[-1] != self.resize.size[1]:  # type: ignore[union-attr]
+                x = F.resize(
+                    x, self.resize.size, interpolation=self.resize.interpolation, antialias=self.resize.antialias
+                )
             # Data type
             if x.dtype == torch.uint8:  # type: ignore[union-attr]
                 x = x.to(dtype=torch.float32).div(255).clip(0, 1)  # type: ignore[union-attr]
@@ -113,11 +118,11 @@ class PreProcessor(nn.Module):
             # Sample transform (to tensor, resize)
             samples = list(multithread_exec(self.sample_transforms, x))
             # Batching
-            batches = self.batch_inputs(samples)
+            batches = self.batch_inputs(samples)  # type: ignore[assignment]
         else:
             raise TypeError(f"invalid input type: {type(x)}")
 
         # Batch transforms (normalize)
         batches = list(multithread_exec(self.normalize, batches))
 
-        return batches
+        return batches  # type: ignore[return-value]

doctr/models/preprocessor/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.
@@ -20,6 +20,7 @@ class PreProcessor(NestedObject):
     """Implements an abstract preprocessor object which performs casting, resizing, batching and normalization.
 
     Args:
+    ----
         output_size: expected size of each page in format (H, W)
         batch_size: the size of page batches
         mean: mean value of the training distribution by channel
@@ -34,7 +35,6 @@ class PreProcessor(NestedObject):
         batch_size: int,
         mean: Tuple[float, float, float] = (0.5, 0.5, 0.5),
         std: Tuple[float, float, float] = (1.0, 1.0, 1.0),
-        fp16: bool = False,
         **kwargs: Any,
     ) -> None:
         self.batch_size = batch_size
@@ -46,12 +46,13 @@ class PreProcessor(NestedObject):
         """Gather samples into batches for inference purposes
 
         Args:
+        ----
             samples: list of samples (tf.Tensor)
 
         Returns:
+        -------
             list of batched samples
         """
-
         num_batches = int(math.ceil(len(samples) / self.batch_size))
         batches = [
             tf.stack(samples[idx * self.batch_size : min((idx + 1) * self.batch_size, len(samples))], axis=0)
@@ -81,11 +82,13 @@ class PreProcessor(NestedObject):
         """Prepare document data for model forwarding
 
         Args:
+        ----
             x: list of images (np.array) or tensors (already resized and batched)
+
         Returns:
+        -------
             list of page batches
         """
-
         # Input type check
         if isinstance(x, (np.ndarray, tf.Tensor)):
             if x.ndim != 4:
@@ -102,7 +105,9 @@ class PreProcessor(NestedObject):
                 x = tf.image.convert_image_dtype(x, dtype=tf.float32)
             # Resizing
             if (x.shape[1], x.shape[2]) != self.resize.output_size:
-                x = tf.image.resize(x, self.resize.output_size, method=self.resize.method)
+                x = tf.image.resize(
+                    x, self.resize.output_size, method=self.resize.method, antialias=self.resize.antialias
+                )
 
             batches = [x]
 

doctr/models/recognition/core.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, Mindee.
 
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
@@ -27,9 +27,11 @@ class RecognitionModel(NestedObject):
         sequence lengths.
 
         Args:
+        ----
             gts: list of ground-truth labels
 
         Returns:
+        -------
             A tuple of 2 tensors: Encoded labels and sequence lengths (for each entry of the batch)
         """
         encoded = encode_sequences(sequences=gts, vocab=self.vocab, target_size=self.max_length, eos=len(self.vocab))
@@ -41,6 +43,7 @@ class RecognitionPostProcessor(NestedObject):
     """Abstract class to postprocess the raw output of the model
 
     Args:
+    ----
         vocab: string containing the ordered sequence of supported characters
     """