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

doctr/models/predictor/base.py

@@ -3,16 +3,16 @@
 # 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, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 
 from doctr.models.builder import DocumentBuilder
-from doctr.utils.geometry import extract_crops, extract_rcrops
+from doctr.utils.geometry import extract_crops, extract_rcrops, rotate_image
 
-from .._utils import rectify_crops, rectify_loc_preds
-from ..classification import crop_orientation_predictor
-from ..classification.predictor import CropOrientationPredictor
+from .._utils import estimate_orientation, rectify_crops, rectify_loc_preds
+from ..classification import crop_orientation_predictor, page_orientation_predictor
+from ..classification.predictor import OrientationPredictor
 
 __all__ = ["_OCRPredictor"]
 
@@ -29,10 +29,13 @@ 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.
+        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,
@@ -40,16 +43,69 @@ class _OCRPredictor:
         straighten_pages: bool = False,
         preserve_aspect_ratio: bool = True,
         symmetric_pad: bool = True,
+        detect_orientation: bool = False,
         **kwargs: Any,
     ) -> None:
         self.assume_straight_pages = assume_straight_pages
         self.straighten_pages = straighten_pages
         self.crop_orientation_predictor = None if assume_straight_pages else crop_orientation_predictor(pretrained=True)
+        self.page_orientation_predictor = (
+            page_orientation_predictor(pretrained=True)
+            if detect_orientation or straighten_pages or not assume_straight_pages
+            else None
+        )
         self.doc_builder = DocumentBuilder(**kwargs)
         self.preserve_aspect_ratio = preserve_aspect_ratio
         self.symmetric_pad = symmetric_pad
         self.hooks: List[Callable] = []
 
+    def _general_page_orientations(
+        self,
+        pages: List[np.ndarray],
+    ) -> List[Tuple[int, float]]:
+        _, classes, probs = zip(self.page_orientation_predictor(pages))  # type: ignore[misc]
+        # Flatten to list of tuples with (value, confidence)
+        page_orientations = [
+            (orientation, prob)
+            for page_classes, page_probs in zip(classes, probs)
+            for orientation, prob in zip(page_classes, page_probs)
+        ]
+        return page_orientations
+
+    def _get_orientations(
+        self, pages: List[np.ndarray], seg_maps: List[np.ndarray]
+    ) -> Tuple[List[Tuple[int, float]], List[int]]:
+        general_pages_orientations = self._general_page_orientations(pages)
+        origin_page_orientations = [
+            estimate_orientation(seq_map, general_orientation)
+            for seq_map, general_orientation in zip(seg_maps, general_pages_orientations)
+        ]
+        return general_pages_orientations, origin_page_orientations
+
+    def _straighten_pages(
+        self,
+        pages: List[np.ndarray],
+        seg_maps: List[np.ndarray],
+        general_pages_orientations: Optional[List[Tuple[int, float]]] = None,
+        origin_pages_orientations: Optional[List[int]] = None,
+    ) -> List[np.ndarray]:
+        general_pages_orientations = (
+            general_pages_orientations if general_pages_orientations else self._general_page_orientations(pages)
+        )
+        origin_pages_orientations = (
+            origin_pages_orientations
+            if origin_pages_orientations
+            else [
+                estimate_orientation(seq_map, general_orientation)
+                for seq_map, general_orientation in zip(seg_maps, general_pages_orientations)
+            ]
+        )
+        return [
+            # We exapnd if the page is wider than tall and the angle is 90 or -90
+            rotate_image(page, angle, expand=page.shape[1] > page.shape[0] and abs(angle) == 90)
+            for page, angle in zip(pages, origin_pages_orientations)
+        ]
+
     @staticmethod
     def _generate_crops(
         pages: List[np.ndarray],
@@ -88,68 +144,39 @@ class _OCRPredictor:
         self,
         crops: List[List[np.ndarray]],
         loc_preds: List[np.ndarray],
-    ) -> Tuple[List[List[np.ndarray]], List[np.ndarray]]:
+    ) -> Tuple[List[List[np.ndarray]], List[np.ndarray], List[Tuple[int, float]]]:
         # Work at a page level
-        orientations = [self.crop_orientation_predictor(page_crops) for page_crops in crops]  # type: ignore[misc]
+        orientations, classes, probs = zip(*[self.crop_orientation_predictor(page_crops) for page_crops in crops])  # type: ignore[misc]
         rect_crops = [rectify_crops(page_crops, orientation) for page_crops, orientation in zip(crops, orientations)]
         rect_loc_preds = [
             rectify_loc_preds(page_loc_preds, orientation) if len(page_loc_preds) > 0 else page_loc_preds
             for page_loc_preds, orientation in zip(loc_preds, orientations)
         ]
-        return rect_crops, rect_loc_preds  # type: ignore[return-value]
-
-    def _remove_padding(
-        self,
-        pages: List[np.ndarray],
-        loc_preds: List[np.ndarray],
-    ) -> List[np.ndarray]:
-        if self.preserve_aspect_ratio:
-            # Rectify loc_preds to remove padding
-            rectified_preds = []
-            for page, loc_pred in zip(pages, loc_preds):
-                h, w = page.shape[0], page.shape[1]
-                if h > w:
-                    # y unchanged, dilate x coord
-                    if self.symmetric_pad:
-                        if self.assume_straight_pages:
-                            loc_pred[:, [0, 2]] = np.clip((loc_pred[:, [0, 2]] - 0.5) * h / w + 0.5, 0, 1)
-                        else:
-                            loc_pred[:, :, 0] = np.clip((loc_pred[:, :, 0] - 0.5) * h / w + 0.5, 0, 1)
-                    else:
-                        if self.assume_straight_pages:
-                            loc_pred[:, [0, 2]] *= h / w
-                        else:
-                            loc_pred[:, :, 0] *= h / w
-                elif w > h:
-                    # x unchanged, dilate y coord
-                    if self.symmetric_pad:
-                        if self.assume_straight_pages:
-                            loc_pred[:, [1, 3]] = np.clip((loc_pred[:, [1, 3]] - 0.5) * w / h + 0.5, 0, 1)
-                        else:
-                            loc_pred[:, :, 1] = np.clip((loc_pred[:, :, 1] - 0.5) * w / h + 0.5, 0, 1)
-                    else:
-                        if self.assume_straight_pages:
-                            loc_pred[:, [1, 3]] *= w / h
-                        else:
-                            loc_pred[:, :, 1] *= w / h
-                rectified_preds.append(loc_pred)
-            return rectified_preds
-        return loc_preds
+        # Flatten to list of tuples with (value, confidence)
+        crop_orientations = [
+            (orientation, prob)
+            for page_classes, page_probs in zip(classes, probs)
+            for orientation, prob in zip(page_classes, page_probs)
+        ]
+        return rect_crops, rect_loc_preds, crop_orientations  # type: ignore[return-value]
 
     @staticmethod
     def _process_predictions(
         loc_preds: List[np.ndarray],
         word_preds: List[Tuple[str, float]],
-    ) -> Tuple[List[np.ndarray], List[List[Tuple[str, float]]]]:
+        crop_orientations: List[Dict[str, Any]],
+    ) -> Tuple[List[np.ndarray], List[List[Tuple[str, float]]], List[List[Dict[str, Any]]]]:
         text_preds = []
+        crop_orientation_preds = []
         if len(loc_preds) > 0:
-            # Text
+            # Text & crop orientation predictions at page level
             _idx = 0
             for page_boxes in loc_preds:
                 text_preds.append(word_preds[_idx : _idx + page_boxes.shape[0]])
+                crop_orientation_preds.append(crop_orientations[_idx : _idx + page_boxes.shape[0]])
                 _idx += page_boxes.shape[0]
 
-        return loc_preds, text_preds
+        return loc_preds, text_preds, crop_orientation_preds
 
     def add_hook(self, hook: Callable) -> None:
         """Add a hook to the predictor

doctr/models/predictor/pytorch.py

@@ -10,10 +10,10 @@ import torch
 from torch import nn
 
 from doctr.io.elements import Document
-from doctr.models._utils import estimate_orientation, get_language
+from doctr.models._utils import get_language
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
-from doctr.utils.geometry import rotate_image
+from doctr.utils.geometry import detach_scores
 
 from .base import _OCRPredictor
 
@@ -55,7 +55,13 @@ class OCRPredictor(nn.Module, _OCRPredictor):
         self.det_predictor = det_predictor.eval()  # type: ignore[attr-defined]
         self.reco_predictor = reco_predictor.eval()  # type: ignore[attr-defined]
         _OCRPredictor.__init__(
-            self, assume_straight_pages, straighten_pages, preserve_aspect_ratio, symmetric_pad, **kwargs
+            self,
+            assume_straight_pages,
+            straighten_pages,
+            preserve_aspect_ratio,
+            symmetric_pad,
+            detect_orientation,
+            **kwargs,
         )
         self.detect_orientation = detect_orientation
         self.detect_language = detect_language
@@ -81,19 +87,16 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             for out_map in out_maps
         ]
         if self.detect_orientation:
-            origin_page_orientations = [estimate_orientation(seq_map) for seq_map in seg_maps]
+            general_pages_orientations, origin_pages_orientations = self._get_orientations(pages, seg_maps)  # type: ignore[arg-type]
             orientations = [
-                {"value": orientation_page, "confidence": None} for orientation_page in origin_page_orientations
+                {"value": orientation_page, "confidence": None} for orientation_page in origin_pages_orientations
             ]
         else:
             orientations = None
+            general_pages_orientations = None
+            origin_pages_orientations = None
         if self.straighten_pages:
-            origin_page_orientations = (
-                origin_page_orientations
-                if self.detect_orientation
-                else [estimate_orientation(seq_map) for seq_map in seg_maps]
-            )
-            pages = [rotate_image(page, -angle, expand=False) for page, angle in zip(pages, origin_page_orientations)]
+            pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore
             # Forward again to get predictions on straight pages
             loc_preds = self.det_predictor(pages, **kwargs)
 
@@ -102,30 +105,36 @@ class OCRPredictor(nn.Module, _OCRPredictor):
         ), "Detection Model in ocr_predictor should output only one class"
 
         loc_preds = [list(loc_pred.values())[0] for loc_pred in loc_preds]
+        # Detach objectness scores from loc_preds
+        loc_preds, objectness_scores = detach_scores(loc_preds)
         # Check whether crop mode should be switched to channels first
         channels_last = len(pages) == 0 or isinstance(pages[0], np.ndarray)
 
-        # 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,
+            pages,  # type: ignore[arg-type]
             loc_preds,
             channels_last=channels_last,
             assume_straight_pages=self.assume_straight_pages,
         )
-        # Rectify crop orientation
+        # Rectify crop orientation and get crop orientation predictions
+        crop_orientations: Any = []
         if not self.assume_straight_pages:
-            crops, loc_preds = self._rectify_crops(crops, loc_preds)
+            crops, loc_preds, _crop_orientations = self._rectify_crops(crops, loc_preds)
+            crop_orientations = [
+                {"value": orientation[0], "confidence": orientation[1]} for orientation in _crop_orientations
+            ]
+
         # Identify character sequences
         word_preds = self.reco_predictor([crop for page_crops in crops for crop in page_crops], **kwargs)
+        if not crop_orientations:
+            crop_orientations = [{"value": 0, "confidence": None} for _ in word_preds]
 
-        boxes, text_preds = self._process_predictions(loc_preds, word_preds)
+        boxes, text_preds, crop_orientations = self._process_predictions(loc_preds, word_preds, crop_orientations)
 
         if self.detect_language:
             languages = [get_language(" ".join([item[0] for item in text_pred])) for text_pred in text_preds]
@@ -134,10 +143,12 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             languages_dict = None
 
         out = self.doc_builder(
-            pages,
+            pages,  # type: ignore[arg-type]
             boxes,
+            objectness_scores,
             text_preds,
-            origin_page_shapes,
+            origin_page_shapes,  # type: ignore[arg-type]
+            crop_orientations,
             orientations,
             languages_dict,
         )

doctr/models/predictor/tensorflow.py

@@ -9,10 +9,10 @@ import numpy as np
 import tensorflow as tf
 
 from doctr.io.elements import Document
-from doctr.models._utils import estimate_orientation, get_language
+from doctr.models._utils import get_language
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
-from doctr.utils.geometry import rotate_image
+from doctr.utils.geometry import detach_scores
 from doctr.utils.repr import NestedObject
 
 from .base import _OCRPredictor
@@ -56,7 +56,13 @@ class OCRPredictor(NestedObject, _OCRPredictor):
         self.det_predictor = det_predictor
         self.reco_predictor = reco_predictor
         _OCRPredictor.__init__(
-            self, assume_straight_pages, straighten_pages, preserve_aspect_ratio, symmetric_pad, **kwargs
+            self,
+            assume_straight_pages,
+            straighten_pages,
+            preserve_aspect_ratio,
+            symmetric_pad,
+            detect_orientation,
+            **kwargs,
         )
         self.detect_orientation = detect_orientation
         self.detect_language = detect_language
@@ -81,19 +87,16 @@ class OCRPredictor(NestedObject, _OCRPredictor):
             for out_map in out_maps
         ]
         if self.detect_orientation:
-            origin_page_orientations = [estimate_orientation(seq_map) for seq_map in seg_maps]
+            general_pages_orientations, origin_pages_orientations = self._get_orientations(pages, seg_maps)
             orientations = [
-                {"value": orientation_page, "confidence": None} for orientation_page in origin_page_orientations
+                {"value": orientation_page, "confidence": None} for orientation_page in origin_pages_orientations
             ]
         else:
             orientations = None
+            general_pages_orientations = None
+            origin_pages_orientations = None
         if self.straighten_pages:
-            origin_page_orientations = (
-                origin_page_orientations
-                if self.detect_orientation
-                else [estimate_orientation(seq_map) for seq_map in seg_maps]
-            )
-            pages = [rotate_image(page, -angle, expand=False) for page, angle in zip(pages, origin_page_orientations)]
+            pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)
             # forward again to get predictions on straight pages
             loc_preds_dict = self.det_predictor(pages, **kwargs)  # type: ignore[assignment]
 
@@ -101,9 +104,8 @@ class OCRPredictor(NestedObject, _OCRPredictor):
             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]  # type: ignore[union-attr]
-
-        # Rectify crops if aspect ratio
-        loc_preds = self._remove_padding(pages, loc_preds)
+        # Detach objectness scores from loc_preds
+        loc_preds, objectness_scores = detach_scores(loc_preds)
 
         # Apply hooks to loc_preds if any
         for hook in self.hooks:
@@ -113,14 +115,20 @@ class OCRPredictor(NestedObject, _OCRPredictor):
         crops, loc_preds = self._prepare_crops(
             pages, loc_preds, channels_last=True, assume_straight_pages=self.assume_straight_pages
         )
-        # Rectify crop orientation
+        # Rectify crop orientation and get crop orientation predictions
+        crop_orientations: Any = []
         if not self.assume_straight_pages:
-            crops, loc_preds = self._rectify_crops(crops, loc_preds)
+            crops, loc_preds, _crop_orientations = self._rectify_crops(crops, loc_preds)
+            crop_orientations = [
+                {"value": orientation[0], "confidence": orientation[1]} for orientation in _crop_orientations
+            ]
 
         # Identify character sequences
         word_preds = self.reco_predictor([crop for page_crops in crops for crop in page_crops], **kwargs)
+        if not crop_orientations:
+            crop_orientations = [{"value": 0, "confidence": None} for _ in word_preds]
 
-        boxes, text_preds = self._process_predictions(loc_preds, word_preds)
+        boxes, text_preds, crop_orientations = self._process_predictions(loc_preds, word_preds, crop_orientations)
 
         if self.detect_language:
             languages = [get_language(" ".join([item[0] for item in text_pred])) for text_pred in text_preds]
@@ -131,8 +139,10 @@ class OCRPredictor(NestedObject, _OCRPredictor):
         out = self.doc_builder(
             pages,
             boxes,
+            objectness_scores,
             text_preds,
             origin_page_shapes,  # type: ignore[arg-type]
+            crop_orientations,
             orientations,
             languages_dict,
         )

doctr/models/preprocessor/pytorch.py

@@ -79,7 +79,7 @@ class PreProcessor(nn.Module):
         else:
             x = x.to(dtype=torch.float32)  # type: ignore[union-attr]
 
-        return x  # type: ignore[return-value]
+        return x
 
     def __call__(self, x: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]) -> List[torch.Tensor]:
         """Prepare document data for model forwarding
@@ -103,7 +103,7 @@ 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]:  # type: ignore[union-attr]
+            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, antialias=self.resize.antialias
                 )
@@ -118,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)  # type: ignore[assignment]
+            batches = self.batch_inputs(samples)
         else:
             raise TypeError(f"invalid input type: {type(x)}")
 
         # Batch transforms (normalize)
         batches = list(multithread_exec(self.normalize, batches))
 
-        return batches  # type: ignore[return-value]
+        return batches

doctr/models/preprocessor/tensorflow.py

@@ -41,6 +41,7 @@ class PreProcessor(NestedObject):
         self.resize = Resize(output_size, **kwargs)
         # Perform the division by 255 at the same time
         self.normalize = Normalize(mean, std)
+        self._runs_on_cuda = tf.test.is_gpu_available()
 
     def batch_inputs(self, samples: List[tf.Tensor]) -> List[tf.Tensor]:
         """Gather samples into batches for inference purposes
@@ -113,13 +114,13 @@ class PreProcessor(NestedObject):
 
         elif isinstance(x, list) and all(isinstance(sample, (np.ndarray, tf.Tensor)) for sample in x):
             # Sample transform (to tensor, resize)
-            samples = list(multithread_exec(self.sample_transforms, x))
+            samples = list(multithread_exec(self.sample_transforms, x, threads=1 if self._runs_on_cuda else None))
             # Batching
             batches = self.batch_inputs(samples)
         else:
             raise TypeError(f"invalid input type: {type(x)}")
 
         # Batch transforms (normalize)
-        batches = list(multithread_exec(self.normalize, batches))
+        batches = list(multithread_exec(self.normalize, batches, threads=1 if self._runs_on_cuda else None))
 
         return batches

doctr/models/recognition/master/pytorch.py

@@ -107,7 +107,7 @@ class MASTER(_MASTER, nn.Module):
         # NOTE: nn.TransformerDecoder takes the inverse from this implementation
         # [True, True, True, ..., False, False, False] -> False is masked
         # (N, 1, 1, max_length)
-        target_pad_mask = (target != self.vocab_size + 2).unsqueeze(1).unsqueeze(1)  # type: ignore[attr-defined]
+        target_pad_mask = (target != self.vocab_size + 2).unsqueeze(1).unsqueeze(1)
         target_length = target.size(1)
         # sub mask filled diagonal with True = see and False = masked (max_length, max_length)
         # NOTE: onnxruntime tril/triu works only with float currently (onnxruntime 1.11.1 - opset 14)
@@ -142,7 +142,7 @@ class MASTER(_MASTER, nn.Module):
         # Input length : number of timesteps
         input_len = model_output.shape[1]
         # Add one for additional <eos> token (sos disappear in shift!)
-        seq_len = seq_len + 1  # type: ignore[assignment]
+        seq_len = seq_len + 1
         # Compute loss: don't forget to shift gt! Otherwise the model learns to output the gt[t-1]!
         # The "masked" first gt char is <sos>. Delete last logit of the model output.
         cce = F.cross_entropy(model_output[:, :-1, :].permute(0, 2, 1), gt[:, 1:], reduction="none")

doctr/models/recognition/parseq/pytorch.py

@@ -212,7 +212,7 @@ class PARSeq(_PARSeq, nn.Module):
 
         sos_idx = torch.zeros(len(final_perms), 1, device=seqlen.device)
         eos_idx = torch.full((len(final_perms), 1), max_num_chars + 1, device=seqlen.device)
-        combined = torch.cat([sos_idx, final_perms + 1, eos_idx], dim=1).int()  # type: ignore
+        combined = torch.cat([sos_idx, final_perms + 1, eos_idx], dim=1).int()
         if len(combined) > 1:
             combined[1, 1:] = max_num_chars + 1 - torch.arange(max_num_chars + 1, device=seqlen.device)
         return combined
@@ -282,7 +282,8 @@ class PARSeq(_PARSeq, nn.Module):
                 ys[:, i + 1] = pos_prob.squeeze().argmax(-1)
 
                 # Stop decoding if all sequences have reached the EOS token
-                if max_len is None and (ys == self.vocab_size).any(dim=-1).all():  # type: ignore[attr-defined]
+                # NOTE: `break` isn't correctly translated to Onnx so we don't break here if we want to export
+                if not self.exportable and max_len is None and (ys == self.vocab_size).any(dim=-1).all():
                     break
 
         logits = torch.cat(pos_logits, dim=1)  # (N, max_length, vocab_size + 1)
@@ -297,7 +298,7 @@ class PARSeq(_PARSeq, nn.Module):
 
         # Create padding mask for refined target input maskes all behind EOS token as False
         # (N, 1, 1, max_length)
-        target_pad_mask = ~((ys == self.vocab_size).int().cumsum(-1) > 0).unsqueeze(1).unsqueeze(1)  # type: ignore[attr-defined]
+        target_pad_mask = ~((ys == self.vocab_size).int().cumsum(-1) > 0).unsqueeze(1).unsqueeze(1)
         mask = (target_pad_mask.bool() & query_mask[:, : ys.shape[1]].bool()).int()
         logits = self.head(self.decode(ys, features, mask, target_query=pos_queries))
 

doctr/models/recognition/parseq/tensorflow.py

@@ -288,10 +288,11 @@ class PARSeq(_PARSeq, Model):
                 )
 
                 # Stop decoding if all sequences have reached the EOS token
-                # We need to check it on True to be compatible with ONNX
+                # NOTE: `break` isn't correctly translated to Onnx so we don't break here if we want to export
                 if (
-                    max_len is None
-                    and tf.reduce_any(tf.reduce_all(tf.equal(ys, tf.constant(self.vocab_size)), axis=-1)) is True
+                    not self.exportable
+                    and max_len is None
+                    and tf.reduce_any(tf.reduce_all(tf.equal(ys, tf.constant(self.vocab_size)), axis=-1))
                 ):
                     break
 

doctr/models/recognition/sar/pytorch.py

@@ -125,25 +125,26 @@ class SARDecoder(nn.Module):
             if t == 0:
                 # step to init the first states of the LSTMCell
                 hidden_state_init = cell_state_init = torch.zeros(
-                    features.size(0), features.size(1), device=features.device
+                    features.size(0), features.size(1), device=features.device, dtype=features.dtype
                 )
                 hidden_state, cell_state = hidden_state_init, cell_state_init
                 prev_symbol = holistic
             elif t == 1:
                 # step to init a 'blank' sequence of length vocab_size + 1 filled with zeros
                 # (N, vocab_size + 1) --> (N, embedding_units)
-                prev_symbol = torch.zeros(features.size(0), self.vocab_size + 1, device=features.device)
+                prev_symbol = torch.zeros(
+                    features.size(0), self.vocab_size + 1, device=features.device, dtype=features.dtype
+                )
                 prev_symbol = self.embed(prev_symbol)
             else:
-                if gt is not None:
+                if gt is not None and self.training:
                     # (N, embedding_units) -2 because of <bos> and <eos> (same)
                     prev_symbol = self.embed(gt_embedding[:, t - 2])
                 else:
                     # -1 to start at timestep where prev_symbol was initialized
                     index = logits_list[t - 1].argmax(-1)
                     # update prev_symbol with ones at the index of the previous logit vector
-                    # (N, embedding_units)
-                    prev_symbol = prev_symbol.scatter_(1, index.unsqueeze(1), 1)
+                    prev_symbol = self.embed(self.embed_tgt(index))
 
             # (N, C), (N, C)  take the last hidden state and cell state from current timestep
             hidden_state_init, cell_state_init = self.lstm_cell(prev_symbol, (hidden_state_init, cell_state_init))
@@ -292,7 +293,7 @@ class SAR(nn.Module, RecognitionModel):
         # Input length : number of timesteps
         input_len = model_output.shape[1]
         # Add one for additional <eos> token
-        seq_len = seq_len + 1  # type: ignore[assignment]
+        seq_len = seq_len + 1
         # Compute loss
         # (N, L, vocab_size + 1)
         cce = F.cross_entropy(model_output.permute(0, 2, 1), gt, reduction="none")

doctr/models/recognition/sar/tensorflow.py

@@ -177,23 +177,17 @@ class SARDecoder(layers.Layer, NestedObject):
             elif t == 1:
                 # step to init a 'blank' sequence of length vocab_size + 1 filled with zeros
                 # (N, vocab_size + 1) --> (N, embedding_units)
-                prev_symbol = tf.zeros([features.shape[0], self.vocab_size + 1])
+                prev_symbol = tf.zeros([features.shape[0], self.vocab_size + 1], dtype=features.dtype)
                 prev_symbol = self.embed(prev_symbol, **kwargs)
             else:
-                if gt is not None:
+                if gt is not None and kwargs.get("training", False):
                     # (N, embedding_units) -2 because of <bos> and <eos> (same)
                     prev_symbol = self.embed(gt_embedding[:, t - 2], **kwargs)
                 else:
                     # -1 to start at timestep where prev_symbol was initialized
                     index = tf.argmax(logits_list[t - 1], axis=-1)
                     # update prev_symbol with ones at the index of the previous logit vector
-                    # (N, embedding_units)
-                    index = tf.ones_like(index)
-                    prev_symbol = tf.scatter_nd(
-                        tf.expand_dims(index, axis=1),
-                        prev_symbol,
-                        tf.constant([features.shape[0], features.shape[-1]], dtype=tf.int64),
-                    )
+                    prev_symbol = self.embed(self.embed_tgt(index, **kwargs), **kwargs)
 
             # (N, C), (N, C)  take the last hidden state and cell state from current timestep
             _, states = self.lstm_cells(prev_symbol, states, **kwargs)

doctr/models/recognition/vitstr/pytorch.py

@@ -137,7 +137,7 @@ class ViTSTR(_ViTSTR, nn.Module):
         # Input length : number of steps
         input_len = model_output.shape[1]
         # Add one for additional <eos> token (sos disappear in shift!)
-        seq_len = seq_len + 1  # type: ignore[assignment]
+        seq_len = seq_len + 1
         # Compute loss: don't forget to shift gt! Otherwise the model learns to output the gt[t-1]!
         # The "masked" first gt char is <sos>.
         cce = F.cross_entropy(model_output.permute(0, 2, 1), gt[:, 1:], reduction="none")

doctr/models/recognition/zoo.py

@@ -45,7 +45,7 @@ def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredict
 
     kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"])
     kwargs["std"] = kwargs.get("std", _model.cfg["std"])
-    kwargs["batch_size"] = kwargs.get("batch_size", 32)
+    kwargs["batch_size"] = kwargs.get("batch_size", 128)
     input_shape = _model.cfg["input_shape"][:2] if is_tf_available() else _model.cfg["input_shape"][-2:]
     predictor = RecognitionPredictor(PreProcessor(input_shape, preserve_aspect_ratio=True, **kwargs), _model)
 

doctr/models/zoo.py

@@ -61,7 +61,7 @@ def _predictor(
 
 
 def ocr_predictor(
-    det_arch: Any = "db_resnet50",
+    det_arch: Any = "fast_base",
     reco_arch: Any = "crnn_vgg16_bn",
     pretrained: bool = False,
     pretrained_backbone: bool = True,
@@ -175,7 +175,7 @@ def _kie_predictor(
 
 
 def kie_predictor(
-    det_arch: Any = "db_resnet50",
+    det_arch: Any = "fast_base",
     reco_arch: Any = "crnn_vgg16_bn",
     pretrained: bool = False,
     pretrained_backbone: bool = True,

doctr/transforms/functional/base.py

@@ -200,4 +200,4 @@ def create_shadow_mask(
     mask: np.ndarray = np.zeros((*target_shape, 1), dtype=np.uint8)
     mask = cv2.fillPoly(mask, [final_contour], (255,), lineType=cv2.LINE_AA)[..., 0]
 
-    return (mask / 255).astype(np.float32).clip(0, 1) * intensity_mask.astype(np.float32)  # type: ignore[operator]
+    return (mask / 255).astype(np.float32).clip(0, 1) * intensity_mask.astype(np.float32)