Rotate page

doctr/models/_utils.py

@@ -150,7 +150,11 @@ def estimate_orientation(img: np.ndarray, n_ct: int = 50, ratio_threshold_for_li
             angles.append(angle)
         elif w / h < 1 / ratio_threshold_for_lines:  # if lines are vertical, substract 90 degree
             angles.append(angle - 90)
-    return -median_low(angles)
+
+    if len(angles) == 0:
+        return 0  # in case no angles is found
+    else:
+        return -median_low(angles)
 
 
 def get_bitmap_angle(bitmap: np.ndarray, n_ct: int = 20, std_max: float = 3.) -> float:

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
 
@@ -29,6 +31,9 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             without rotated textual elements.
         export_as_straight_boxes: when assume_straight_pages is set to False, export final predictions
             (potentially rotated) as straight bounding boxes.
+        straighten_pages: if True, estimates the page general orientation based on the median line orientation.
+            Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped
+            accordingly. Doing so will improve performances for documents with page-uniform rotations.
 
     """
 
@@ -38,13 +43,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(
@@ -57,6 +64,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] if isinstance(page, np.ndarray) else 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
@@ -70,6 +84,14 @@ def forward(
 
         boxes, text_preds = self._process_predictions(loc_preds, word_preds)
 
+        # 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) for
+                     page_boxes, page, angle, mask in zip(boxes, pages, origin_page_orientations, origin_page_shapes)]
+
         out = self.doc_builder(
             boxes,
             text_preds,

doctr/models/predictor/tensorflow.py

@@ -9,9 +9,11 @@
 import tensorflow as tf
 
 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 doctr.utils.repr import NestedObject
 
 from .base import _OCRPredictor
@@ -29,6 +31,9 @@ class OCRPredictor(NestedObject, _OCRPredictor):
             without rotated textual elements.
         export_as_straight_boxes: when assume_straight_pages is set to False, export final predictions
             (potentially rotated) as straight bounding boxes.
+        straighten_pages: if True, estimates the page general orientation based on the median line orientation.
+            Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped
+            accordingly. Doing so will improve performances for documents with page-uniform rotations.
     """
     _children_names = ['det_predictor', 'reco_predictor']
 
@@ -38,13 +43,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
         self.reco_predictor = reco_predictor
         self.doc_builder = DocumentBuilder(export_as_straight_boxes=export_as_straight_boxes)
         self.assume_straight_pages = assume_straight_pages
+        self.straighten_pages = straighten_pages
 
     def __call__(
         self,
@@ -56,6 +63,13 @@ def __call__(
         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)
         # Crop images
@@ -67,5 +81,10 @@ def __call__(
 
         boxes, text_preds = self._process_predictions(loc_preds, word_preds)
 
-        out = self.doc_builder(boxes, text_preds, [page.shape[:2] for page in pages])  # type: ignore[misc]
+        # 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], target_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, origin_page_shapes)  # type: ignore[misc]
         return out

doctr/utils/common_types.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 from typing import List, Tuple, Union
 
-__all__ = ['Point2D', 'BoundingBox', 'RotatedBbox', 'Polygon4P', 'Polygon']
+__all__ = ['Point2D', 'BoundingBox', 'RotatedBbox', 'Polygon4P', 'Polygon', 'Bbox']
 
 
 Point2D = Tuple[float, float]

doctr/utils/geometry.py

@@ -3,8 +3,8 @@
 # This program is licensed under the Apache License version 2.
 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
 
-import math
-from typing import List, Tuple, Union
+from math import ceil
+from typing import List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
@@ -145,54 +145,105 @@ def rotate_abs_boxes(boxes: np.ndarray, angle: float, img_shape: Tuple[int, int]
     return rotated_boxes
 
 
+def remap_boxes(loc_preds: np.ndarray, orig_shape: Tuple[int, int], dest_shape: Tuple[int, int]) -> np.ndarray:
+    """ Remaps a batch of rotated locpred (x, y, w, h, alpha, c) expressed for an origin_shape to a destination_shape.
+    This does not impact the absolute shape of the boxes, but allow to calculate the new relative RotatedBbox
+    coordinates after a resizing of the image.
+
+    Args:
+        loc_preds: (N, 6) array of RELATIVE locpred (x, y, w, h, alpha, c)
+        orig_shape: shape of the origin image
+        dest_shape: shape of the destination image
+
+    Returns:
+        A batch of rotated loc_preds (N, 6): (x, y, w, h, alpha, c) expressed in the destination referencial
+
+    """
+
+    if len(dest_shape) != 2:
+        raise ValueError(f"Mask length should be 2, was found at: {len(dest_shape)}")
+    if len(orig_shape) != 2:
+        raise ValueError(f"Image_shape length should be 2, was found at: {len(orig_shape)}")
+    orig_height, orig_width = orig_shape
+    dest_height, dest_width = dest_shape
+    mboxes = loc_preds.copy()
+    # remaps position of the box center for the destination image shape
+    mboxes[:, 0] = ((loc_preds[:, 0] * orig_width) + (dest_width - orig_width) / 2) / dest_width
+    mboxes[:, 1] = ((loc_preds[:, 1] * orig_height) + (dest_height - orig_height) / 2) / dest_height
+    # remaps box dimension for the destination image shape
+    mboxes[:, 2] = loc_preds[:, 2] * orig_width / dest_width
+    mboxes[:, 3] = loc_preds[:, 3] * orig_height / dest_height
+    return mboxes
+
+
 def rotate_boxes(
-    boxes: np.ndarray,
-    angle: float = 0.,
-    min_angle: float = 1.
+    loc_preds: np.ndarray,
+    angle: float,
+    orig_shape: Tuple[int, int],
+    min_angle: float = 1.,
+    target_shape: Optional[Tuple[int, int]] = None,
 ) -> np.ndarray:
-    """Rotate a batch of straight bounding boxes (xmin, ymin, xmax, ymax) of an angle,
-    if angle > min_angle, around the center of the page.
+    """Rotate a batch of straight bounding boxes (xmin, ymin, xmax, ymax, c) or rotated bounding boxes
+    (x, y, w, h, alpha, c) of an angle, if angle > min_angle, around the center of the page.
+    If target_shape is specified, the boxes are remapped to the target shape after the rotation. This
+    is done to remove the padding that is created by rotate_page(expand=True)
 
     Args:
-        boxes: (N, 4) array of RELATIVE boxes
+        loc_preds: (N, 5) or (N, 6) array of RELATIVE boxes
         angle: angle between -90 and +90 degrees
+        orig_shape: shape of the origin image
         min_angle: minimum angle to rotate boxes
+        target_shape: shape of the target image
 
     Returns:
         A batch of rotated boxes (N, 5): (x, y, w, h, alpha) or a batch of straight bounding boxes
     """
+
+    # Change format of the boxes to rotated boxes
+    _boxes = loc_preds.copy()
+    if _boxes.shape[1] == 5:
+        _boxes = np.column_stack(((_boxes[:, 0] + _boxes[:, 2]) / 2,
+                                  (_boxes[:, 1] + _boxes[:, 3]) / 2,
+                                  _boxes[:, 2] - _boxes[:, 0],
+                                  _boxes[:, 3] - _boxes[:, 1],
+                                  np.zeros(_boxes.shape[0]),
+                                  _boxes[:, 4]))
     # If small angle, return boxes (no rotation)
     if abs(angle) < min_angle or abs(angle) > 90 - min_angle:
-        return boxes
+        return _boxes
     # Compute rotation matrix
     angle_rad = angle * np.pi / 180.  # compute radian angle for np functions
     rotation_mat = np.array([
         [np.cos(angle_rad), -np.sin(angle_rad)],
         [np.sin(angle_rad), np.cos(angle_rad)]
-    ], dtype=boxes.dtype)
-    # Compute unrotated boxes
-    x_unrotated, y_unrotated = (boxes[:, 0] + boxes[:, 2]) / 2, (boxes[:, 1] + boxes[:, 3]) / 2
-    width, height = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]
-    # Rotate centers
-    centers = np.stack((x_unrotated, y_unrotated), axis=-1)
-    rotated_centers = .5 + np.matmul(centers - .5, np.transpose(rotation_mat))
-    x_center, y_center = rotated_centers[:, 0], rotated_centers[:, 1]
+    ], dtype=_boxes.dtype)
+    # Rotate absolute centers
+    centers = np.stack((_boxes[:, 0] * orig_shape[1], _boxes[:, 1] * orig_shape[0]), axis=-1)
+    image_center = (orig_shape[1] / 2, orig_shape[0] / 2)
+    rotated_centers = image_center + np.matmul(centers - image_center, rotation_mat)
+    x_center, y_center = rotated_centers[:, 0] / orig_shape[1], rotated_centers[:, 1] / orig_shape[0]
     # Compute rotated boxes
-    rotated_boxes = np.stack((x_center, y_center, width, height, angle * np.ones_like(boxes[:, 0])), axis=1)
+    rotated_boxes = np.stack((x_center, y_center, _boxes[:, 2], _boxes[:, 3], angle * np.ones_like(_boxes[:, 0]),
+                              _boxes[:, 5]), axis=1)
+    # Apply a mask if requested
+    if target_shape is not None:
+        rotated_boxes = remap_boxes(rotated_boxes, orig_shape=orig_shape, dest_shape=target_shape)
     return rotated_boxes
 
 
 def rotate_image(
     image: np.ndarray,
     angle: float,
-    expand=False,
+    expand: bool = False,
+    preserve_origin_shape: bool = False,
 ) -> np.ndarray:
     """Rotate an image counterclockwise by an given angle.
 
     Args:
         image: numpy tensor to rotate
         angle: rotation angle in degrees, between -90 and +90
         expand: whether the image should be padded before the rotation
+        preserve_origin_shape: if expand is set to True, resizes the final output to the original image size
 
     Returns:
         Rotated array, padded by 0 by default.
@@ -201,14 +252,15 @@ def rotate_image(
     # Compute the expanded padding
     if expand:
         exp_shape = compute_expanded_shape(image.shape[:-1], angle)
-        h_pad, w_pad = int(math.ceil(exp_shape[0] - image.shape[0])), int(math.ceil(exp_shape[1] - image.shape[1]))
+        h_pad, w_pad = int(max(0, ceil(exp_shape[0] - image.shape[0]))), int(
+            max(0, ceil(exp_shape[1] - image.shape[1])))
         exp_img = np.pad(image, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
     else:
         exp_img = image
 
     height, width = exp_img.shape[:2]
     rot_mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
-    rot_img = cv2.warpAffine(exp_img.astype(np.float32), rot_mat, (width, height))
+    rot_img = cv2.warpAffine(exp_img, rot_mat, (width, height))
     if expand:
         # Pad to get the same aspect ratio
         if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):
@@ -219,7 +271,8 @@ def rotate_image(
             else:
                 h_pad, w_pad = int(rot_img.shape[1] * image.shape[0] / image.shape[1] - rot_img.shape[0]), 0
             rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
-        # rescale
-        rot_img = cv2.resize(rot_img, image.shape[:-1][::-1], interpolation=cv2.INTER_LINEAR)
+        if preserve_origin_shape:
+            # rescale
+            rot_img = cv2.resize(rot_img, image.shape[:-1][::-1], interpolation=cv2.INTER_LINEAR)
 
     return rot_img

doctr/utils/visualization.py

@@ -100,7 +100,7 @@ def polygon_patch(
     # Switch to absolute coords
     x, w = x * width, w * width
     y, h = y * height, h * height
-    points = cv2.boxPoints(((x, y), (w, h), a))
+    points = cv2.boxPoints(((x, y), (w, h), -a))
 
     return patches.Polygon(
         points,

tests/common/test_models.py

@@ -7,6 +7,7 @@
 
 from doctr.io import DocumentFile, reader
 from doctr.models._utils import estimate_orientation, extract_crops, extract_rcrops, get_bitmap_angle
+from doctr.utils import geometry
 
 
 def test_extract_crops(mock_pdf):  # noqa: F811
@@ -91,5 +92,11 @@ def test_get_bitmap_angle(mock_bitmap):
 
 
 def test_estimate_orientation(mock_image):
+    assert estimate_orientation(mock_image * 0) == 0
+
     angle = estimate_orientation(mock_image)
     assert abs(angle - 30.) < 1.
+
+    rotated = geometry.rotate_image(mock_image, -angle)
+    angle_rotated = estimate_orientation(rotated)
+    assert abs(angle_rotated) < 1.