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

doctr/models/builder.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 DocumentBuilder(NestedObject):
     """Implements a document builder
 
     Args:
+    ----
         resolve_lines: whether words should be automatically grouped into lines
         resolve_blocks: whether lines should be automatically grouped into blocks
         paragraph_break: relative length of the minimum space separating paragraphs
@@ -44,9 +45,11 @@ class DocumentBuilder(NestedObject):
         """Sort bounding boxes from top to bottom, left to right
 
         Args:
+        ----
             boxes: bounding boxes of shape (N, 4) or (N, 4, 2) (in case of rotated bbox)
 
         Returns:
+        -------
             tuple: indices of ordered boxes of shape (N,), boxes
                 If straight boxes are passed tpo the function, boxes are unchanged
                 else: boxes returned are straight boxes fitted to the straightened rotated boxes
@@ -66,10 +69,12 @@ class DocumentBuilder(NestedObject):
         """Split a line in sub_lines
 
         Args:
+        ----
             boxes: bounding boxes of shape (N, 4)
             word_idcs: list of indexes for the words of the line
 
         Returns:
+        -------
             A list of (sub-)lines computed from the original line (words)
         """
         lines = []
@@ -104,12 +109,13 @@ class DocumentBuilder(NestedObject):
         """Order boxes to group them in lines
 
         Args:
+        ----
             boxes: bounding boxes of shape (N, 4) or (N, 4, 2) in case of rotated bbox
 
         Returns:
+        -------
             nested list of box indices
         """
-
         # Sort boxes, and straighten the boxes if they are rotated
         idxs, boxes = self._sort_boxes(boxes)
 
@@ -151,25 +157,23 @@ class DocumentBuilder(NestedObject):
         """Order lines to group them in blocks
 
         Args:
+        ----
             boxes: bounding boxes of shape (N, 4) or (N, 4, 2)
             lines: list of lines, each line is a list of idx
 
         Returns:
+        -------
             nested list of box indices
         """
         # Resolve enclosing boxes of lines
         if boxes.ndim == 3:
-            box_lines: np.ndarray = np.asarray(
-                [
-                    resolve_enclosing_rbbox([tuple(boxes[idx, :, :]) for idx in line])  # type: ignore[misc]
-                    for line in lines
-                ]
-            )
+            box_lines: np.ndarray = np.asarray([
+                resolve_enclosing_rbbox([tuple(boxes[idx, :, :]) for idx in line])  # type: ignore[misc]
+                for line in lines
+            ])
         else:
             _box_lines = [
-                resolve_enclosing_bbox(
-                    [(tuple(boxes[idx, :2]), tuple(boxes[idx, 2:])) for idx in line]  # type: ignore[misc]
-                )
+                resolve_enclosing_bbox([(tuple(boxes[idx, :2]), tuple(boxes[idx, 2:])) for idx in line])
                 for line in lines
             ]
             box_lines = np.asarray([(x1, y1, x2, y2) for ((x1, y1), (x2, y2)) in _box_lines])
@@ -220,13 +224,14 @@ class DocumentBuilder(NestedObject):
         """Gather independent words in structured blocks
 
         Args:
+        ----
             boxes: bounding boxes of all detected words of the page, of shape (N, 5) or (N, 4, 2)
             word_preds: list of all detected words of the page, of shape N
 
         Returns:
+        -------
             list of block elements
         """
-
         if boxes.shape[0] != len(word_preds):
             raise ValueError(f"Incompatible argument lengths: {boxes.shape[0]}, {len(word_preds)}")
 
@@ -248,24 +253,18 @@ class DocumentBuilder(NestedObject):
             _blocks = [lines]
 
         blocks = [
-            Block(
-                [
-                    Line(
-                        [
-                            Word(
-                                *word_preds[idx],
-                                tuple([tuple(pt) for pt in boxes[idx].tolist()]),  # type: ignore[arg-type]
-                            )
-                            if boxes.ndim == 3
-                            else Word(
-                                *word_preds[idx], ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3]))
-                            )
-                            for idx in line
-                        ]
+            Block([
+                Line([
+                    Word(
+                        *word_preds[idx],
+                        tuple([tuple(pt) for pt in boxes[idx].tolist()]),  # type: ignore[arg-type]
                     )
-                    for line in lines
-                ]
-            )
+                    if boxes.ndim == 3
+                    else Word(*word_preds[idx], ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3])))
+                    for idx in line
+                ])
+                for line in lines
+            ])
             for lines in _blocks
         ]
 
@@ -280,6 +279,7 @@ class DocumentBuilder(NestedObject):
 
     def __call__(
         self,
+        pages: List[np.ndarray],
         boxes: List[np.ndarray],
         text_preds: List[List[Tuple[str, float]]],
         page_shapes: List[Tuple[int, int]],
@@ -289,12 +289,19 @@ class DocumentBuilder(NestedObject):
         """Re-arrange detected words into structured blocks
 
         Args:
+        ----
+            pages: list of N elements, where each element represents the page image
             boxes: list of N elements, where each element represents the localization predictions, of shape (*, 5)
                 or (*, 6) for all words for a given page
             text_preds: list of N elements, where each element is the list of all word prediction (text + confidence)
-            page_shape: shape of each page, of size N
+            page_shapes: shape of each page, of size N
+            orientations: optional, list of N elements,
+                where each element is a dictionary containing the orientation (orientation + confidence)
+            languages: optional, list of N elements,
+                where each element is a dictionary containing the language (language + confidence)
 
         Returns:
+        -------
             document object
         """
         if len(boxes) != len(text_preds) or len(boxes) != len(page_shapes):
@@ -307,15 +314,12 @@ class DocumentBuilder(NestedObject):
         if self.export_as_straight_boxes and len(boxes) > 0:
             # If boxes are already straight OK, else fit a bounding rect
             if boxes[0].ndim == 3:
-                straight_boxes: List[np.ndarray] = []
-                # Iterate over pages
-                for p_boxes in boxes:
-                    # Iterate over boxes of the pages
-                    straight_boxes.append(np.concatenate((p_boxes.min(1), p_boxes.max(1)), 1))
-                boxes = straight_boxes
+                # Iterate over pages and boxes
+                boxes = [np.concatenate((p_boxes.min(1), p_boxes.max(1)), 1) for p_boxes in boxes]
 
         _pages = [
             Page(
+                page,
                 self._build_blocks(
                     page_boxes,
                     word_preds,
@@ -325,8 +329,8 @@ class DocumentBuilder(NestedObject):
                 orientation,
                 language,
             )
-            for _idx, shape, page_boxes, word_preds, orientation, language in zip(
-                range(len(boxes)), page_shapes, boxes, text_preds, _orientations, _languages
+            for page, _idx, shape, page_boxes, word_preds, orientation, language in zip(
+                pages, range(len(boxes)), page_shapes, boxes, text_preds, _orientations, _languages
             )
         ]
 
@@ -337,6 +341,7 @@ class KIEDocumentBuilder(DocumentBuilder):
     """Implements a KIE document builder
 
     Args:
+    ----
         resolve_lines: whether words should be automatically grouped into lines
         resolve_blocks: whether lines should be automatically grouped into blocks
         paragraph_break: relative length of the minimum space separating paragraphs
@@ -346,6 +351,7 @@ class KIEDocumentBuilder(DocumentBuilder):
 
     def __call__(  # type: ignore[override]
         self,
+        pages: List[np.ndarray],
         boxes: List[Dict[str, np.ndarray]],
         text_preds: List[Dict[str, List[Tuple[str, float]]]],
         page_shapes: List[Tuple[int, int]],
@@ -355,12 +361,19 @@ class KIEDocumentBuilder(DocumentBuilder):
         """Re-arrange detected words into structured predictions
 
         Args:
+        ----
+            pages: list of N elements, where each element represents the page image
             boxes: list of N dictionaries, where each element represents the localization predictions for a class,
-            of shape (*, 5) or (*, 6) for all predictions
+                of shape (*, 5) or (*, 6) for all predictions
             text_preds: list of N dictionaries, where each element is the list of all word prediction
-            page_shape: shape of each page, of size N
+            page_shapes: shape of each page, of size N
+            orientations: optional, list of N elements,
+                where each element is a dictionary containing the orientation (orientation + confidence)
+            languages: optional, list of N elements,
+                where each element is a dictionary containing the language (language + confidence)
 
         Returns:
+        -------
             document object
         """
         if len(boxes) != len(text_preds) or len(boxes) != len(page_shapes):
@@ -384,6 +397,7 @@ class KIEDocumentBuilder(DocumentBuilder):
 
         _pages = [
             KIEPage(
+                page,
                 {
                     k: self._build_blocks(
                         page_boxes[k],
@@ -396,8 +410,8 @@ class KIEDocumentBuilder(DocumentBuilder):
                 orientation,
                 language,
             )
-            for _idx, shape, page_boxes, word_preds, orientation, language in zip(
-                range(len(boxes)), page_shapes, boxes, text_preds, _orientations, _languages
+            for page, _idx, shape, page_boxes, word_preds, orientation, language in zip(
+                pages, range(len(boxes)), page_shapes, boxes, text_preds, _orientations, _languages
             )
         ]
 
@@ -411,13 +425,14 @@ class KIEDocumentBuilder(DocumentBuilder):
         """Gather independent words in structured blocks
 
         Args:
+        ----
             boxes: bounding boxes of all detected words of the page, of shape (N, 5) or (N, 4, 2)
             word_preds: list of all detected words of the page, of shape N
 
         Returns:
+        -------
             list of block elements
         """
-
         if boxes.shape[0] != len(word_preds):
             raise ValueError(f"Incompatible argument lengths: {boxes.shape[0]}, {len(word_preds)}")
 

doctr/models/classification/__init__.py

@@ -3,4 +3,5 @@ from .resnet import *
 from .vgg import *
 from .magc_resnet import *
 from .vit import *
+from .textnet import *
 from .zoo import *

doctr/models/classification/magc_resnet/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, Mindee.
 
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
@@ -36,6 +36,7 @@ class MAGC(nn.Module):
     <https://arxiv.org/pdf/1910.02562.pdf>`_.
 
     Args:
+    ----
         inplanes: input channels
         headers: number of headers to split channels
         attn_scale: if True, re-scale attention to counteract the variance distibutions
@@ -153,12 +154,14 @@ def magc_resnet31(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A feature extractor model
     """
-
     return _magc_resnet(
         "magc_resnet31",
         pretrained,

doctr/models/classification/magc_resnet/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, Mindee.
 
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
@@ -36,6 +36,7 @@ class MAGC(layers.Layer):
     <https://arxiv.org/pdf/1910.02562.pdf>`_.
 
     Args:
+    ----
         inplanes: input channels
         headers: number of headers to split channels
         attn_scale: if True, re-scale attention to counteract the variance distibutions
@@ -169,12 +170,14 @@ def magc_resnet31(pretrained: bool = False, **kwargs: Any) -> ResNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the ResNet architecture
 
     Returns:
+    -------
         A feature extractor model
     """
-
     return _magc_resnet(
         "magc_resnet31",
         pretrained,

doctr/models/classification/mobilenet/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.
@@ -113,12 +113,14 @@ def mobilenet_v3_small(pretrained: bool = False, **kwargs: Any) -> mobilenetv3.M
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a torch.nn.Module
     """
-
     return _mobilenet_v3(
         "mobilenet_v3_small", pretrained, ignore_keys=["classifier.3.weight", "classifier.3.bias"], **kwargs
     )
@@ -136,12 +138,14 @@ def mobilenet_v3_small_r(pretrained: bool = False, **kwargs: Any) -> mobilenetv3
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a torch.nn.Module
     """
-
     return _mobilenet_v3(
         "mobilenet_v3_small_r",
         pretrained,
@@ -163,9 +167,12 @@ def mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> mobilenetv3.M
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a torch.nn.Module
     """
     return _mobilenet_v3(
@@ -188,9 +195,12 @@ def mobilenet_v3_large_r(pretrained: bool = False, **kwargs: Any) -> mobilenetv3
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a torch.nn.Module
     """
     return _mobilenet_v3(
@@ -214,12 +224,14 @@ def mobilenet_v3_small_orientation(pretrained: bool = False, **kwargs: Any) -> m
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a torch.nn.Module
     """
-
     return _mobilenet_v3(
         "mobilenet_v3_small_orientation",
         pretrained,

doctr/models/classification/mobilenet/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.
@@ -82,14 +82,12 @@ class SqueezeExcitation(Sequential):
     """Squeeze and Excitation."""
 
     def __init__(self, chan: int, squeeze_factor: int = 4) -> None:
-        super().__init__(
-            [
-                layers.GlobalAveragePooling2D(),
-                layers.Dense(chan // squeeze_factor, activation="relu"),
-                layers.Dense(chan, activation="hard_sigmoid"),
-                layers.Reshape((1, 1, chan)),
-            ]
-        )
+        super().__init__([
+            layers.GlobalAveragePooling2D(),
+            layers.Dense(chan // squeeze_factor, activation="relu"),
+            layers.Dense(chan, activation="hard_sigmoid"),
+            layers.Reshape((1, 1, chan)),
+        ])
 
     def call(self, inputs: tf.Tensor, **kwargs: Any) -> tf.Tensor:
         x = super().call(inputs, **kwargs)
@@ -126,6 +124,7 @@ class InvertedResidual(layers.Layer):
     """InvertedResidual for mobilenet
 
     Args:
+    ----
         conf: configuration object for inverted residual
     """
 
@@ -220,14 +219,12 @@ class MobileNetV3(Sequential):
         )
 
         if include_top:
-            _layers.extend(
-                [
-                    layers.GlobalAveragePooling2D(),
-                    layers.Dense(head_chans, activation=hard_swish),
-                    layers.Dropout(0.2),
-                    layers.Dense(num_classes),
-                ]
-            )
+            _layers.extend([
+                layers.GlobalAveragePooling2D(),
+                layers.Dense(head_chans, activation=hard_swish),
+                layers.Dropout(0.2),
+                layers.Dense(num_classes),
+            ])
 
         super().__init__(_layers)
         self.cfg = cfg
@@ -309,12 +306,14 @@ def mobilenet_v3_small(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a keras.Model
     """
-
     return _mobilenet_v3("mobilenet_v3_small", pretrained, False, **kwargs)
 
 
@@ -330,12 +329,14 @@ def mobilenet_v3_small_r(pretrained: bool = False, **kwargs: Any) -> MobileNetV3
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a keras.Model
     """
-
     return _mobilenet_v3("mobilenet_v3_small_r", pretrained, True, **kwargs)
 
 
@@ -351,9 +352,12 @@ def mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a keras.Model
     """
     return _mobilenet_v3("mobilenet_v3_large", pretrained, False, **kwargs)
@@ -371,9 +375,12 @@ def mobilenet_v3_large_r(pretrained: bool = False, **kwargs: Any) -> MobileNetV3
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a keras.Model
     """
     return _mobilenet_v3("mobilenet_v3_large_r", pretrained, True, **kwargs)
@@ -391,10 +398,12 @@ def mobilenet_v3_small_orientation(pretrained: bool = False, **kwargs: Any) -> M
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained: boolean, True if model is pretrained
+        **kwargs: keyword arguments of the MobileNetV3 architecture
 
     Returns:
+    -------
         a keras.Model
     """
-
     return _mobilenet_v3("mobilenet_v3_small_orientation", pretrained, include_top=True, **kwargs)

doctr/models/classification/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.
@@ -20,6 +20,7 @@ class CropOrientationPredictor(nn.Module):
     4 possible orientations: 0, 90, 180, 270 degrees counter clockwise.
 
     Args:
+    ----
         pre_processor: transform inputs for easier batched model inference
         model: core classification architecture (backbone + classification head)
     """
@@ -33,7 +34,7 @@ class CropOrientationPredictor(nn.Module):
         self.pre_processor = pre_processor
         self.model = model.eval()
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         crops: List[Union[np.ndarray, torch.Tensor]],

doctr/models/classification/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.
@@ -20,6 +20,7 @@ class CropOrientationPredictor(NestedObject):
     4 possible orientations: 0, 90, 180, 270 degrees counter clockwise.
 
     Args:
+    ----
         pre_processor: transform inputs for easier batched model inference
         model: core classification architecture (backbone + classification head)
     """