feat: add the straighten_pages to the pytorch predictor

mindee · fg-mindee · Dec 6, 2021 · Jul 19, 2021 · Jul 19, 2021 · Sep 21, 2021
commit 1775ebfa90684f993d4e5ae957859df1802cc687
diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py
@@ -10,9 +10,11 @@
 from torch import nn
 
 from doctr.io.elements import Document
+from doctr.models._utils import estimate_orientation
 from doctr.models.builder import DocumentBuilder
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
+from doctr.utils.geometry import rotate_boxes, rotate_image
 
 from .base import _OCRPredictor
 
@@ -33,13 +35,15 @@ def __init__(
         reco_predictor: RecognitionPredictor,
         assume_straight_pages: bool = True,
         export_as_straight_boxes: bool = False,
+        straighten_pages: bool = False,
     ) -> None:
 
         super().__init__()
         self.det_predictor = det_predictor.eval()  # type: ignore[attr-defined]
         self.reco_predictor = reco_predictor.eval()  # type: ignore[attr-defined]
         self.doc_builder = DocumentBuilder(export_as_straight_boxes=export_as_straight_boxes)
         self.assume_straight_pages = assume_straight_pages
+        self.straighten_pages = straighten_pages
 
     @torch.no_grad()
     def forward(
@@ -52,6 +56,13 @@ def forward(
         if any(page.ndim != 3 for page in pages):
             raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
 
+        origin_page_shapes = [page.shape[:2] for page in pages]
+
+        # Detect document rotation and rotate pages
+        if self.straighten_pages:
+            origin_page_orientations = [estimate_orientation(page) for page in pages]
+            pages = [rotate_image(page, -angle, expand=True) for page, angle in zip(pages, origin_page_orientations)]
+
         # Localize text elements
         loc_preds = self.det_predictor(pages, **kwargs)
         # Check whether crop mode should be switched to channels first
@@ -67,6 +78,11 @@ def forward(
             loc_preds, word_preds, allow_rotated_boxes=not self.doc_builder.export_as_straight_boxes
         )
 
+        # 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], mask_shape=mask) for
+                     page_boxes, page, angle, mask in zip(boxes, pages, origin_page_orientations, origin_page_shapes)]
+
         out = self.doc_builder(
             boxes,
             text_preds,