Rotate page

doctr/models/_utils.py

@@ -5,11 +5,13 @@
 
 import numpy as np
 import cv2
-from math import floor
-from typing import List
+from math import floor, ceil
+from typing import List, Optional, Tuple
 from statistics import median_low
 
-__all__ = ['estimate_orientation', 'extract_crops', 'extract_rcrops', 'get_bitmap_angle']
+from doctr.utils import compute_expanded_shape
+
+__all__ = ['estimate_orientation', 'extract_crops', 'extract_rcrops', 'get_bitmap_angle', 'rotate_image']
 
 
 def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True) -> List[np.ndarray]:
@@ -188,3 +190,47 @@ def get_bitmap_angle(bitmap: np.ndarray, n_ct: int = 20, std_max: float = 3.) ->
             angle = 90 + angle
 
     return angle
+
+
+def rotate_image(
+        image: np.ndarray,
+        angle: float,
+        expand: bool = False,
+        mask_shape: Optional[Tuple[int, int]] = None
+) -> 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
+        mask_shape: applies a mask on the image of the specified shape given in absolute pixels
+
+    Returns:
+        Rotated array, padded by 0 by default.
+    """
+
+    # Compute the expanded padding
+    if expand:
+        exp_shape = compute_expanded_shape(image.shape[:-1], angle)
+        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, rot_mat, (width, height))
+
+    if mask_shape is not None:
+        if len(mask_shape) != 2:
+            raise ValueError(f"Mask length should be 2, was found at: {len(mask_shape)}")
+        h_crop, w_crop = int(height - ceil(mask_shape[0])), int(ceil(width - mask_shape[1]))
+        if h_crop > 0 and w_crop > 0:
+            rot_img = rot_img[h_crop // 2: - h_crop // 2, w_crop // 2: - w_crop // 2]
+        elif w_crop <= 0:
+            rot_img = rot_img[h_crop // 2: - h_crop // 2, ]
+        elif h_crop <= 0:
+            rot_img = rot_img[:, w_crop // 2: - w_crop // 2]
+
+    return rot_img

doctr/models/predictor/tensorflow.py

@@ -8,10 +8,12 @@
 from typing import List, Any, Union
 
 from doctr.io.elements import Document
+from doctr.utils.geometry import rotate_boxes
 from doctr.utils.repr import NestedObject
 from doctr.models.builder import DocumentBuilder
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
+from doctr.models._utils import estimate_orientation, rotate_image
 from .base import _OCRPredictor
 
 
@@ -30,12 +32,14 @@ def __init__(
         self,
         det_predictor: DetectionPredictor,
         reco_predictor: RecognitionPredictor,
-        rotated_bbox: bool = False
+        rotated_bbox: bool = False,
+        straighten_pages: bool = False,
     ) -> None:
 
         super().__init__()
         self.det_predictor = det_predictor
         self.reco_predictor = reco_predictor
+        self.straighten_pages = straighten_pages
         self.doc_builder = DocumentBuilder(rotated_bbox=rotated_bbox)
 
     def __call__(
@@ -48,6 +52,12 @@ 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.")
 
+        # Detect document rotation and rotate pages
+        if self.straighten_pages:
+            page_orientations = [estimate_orientation(page) for page in pages]
+            page_shapes = [page.shape[:-1] for page in pages]
+            pages = [rotate_image(page, -angle, expand=True) for page, angle in zip(pages, page_orientations)]
+
         # Localize text elements
         loc_preds = self.det_predictor(pages, **kwargs)
         # Crop images, rotate page if necessary
@@ -56,5 +66,15 @@ def __call__(
         word_preds = self.reco_predictor([crop for page_crops in crops for crop in page_crops], **kwargs)
 
         boxes, text_preds = self._process_predictions(loc_preds, word_preds, self.doc_builder.rotated_bbox)
+
+        # Rotate back pages and boxes while keeping original image size
+        if self.straighten_pages:
+            pages = [rotate_image(page, angle, expand=True, mask_shape=mask) for page, angle, mask in
+                     zip(pages, page_orientations, page_shapes)]
+            rboxes = [rotate_boxes(page_boxes, angle, expand=True, orig_shape=page.shape[:2], mask_shape=mask) for
+                      page_boxes, page, angle, mask in zip(boxes, pages, page_orientations, page_shapes)]
+            boxes = rboxes
+            self.doc_builder = DocumentBuilder(rotated_bbox=True)  # override the current doc_builder
+
         out = self.doc_builder(boxes, text_preds, [page.shape[:2] for page in pages])  # type: ignore[misc]
         return out

doctr/utils/common_types.py

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

doctr/utils/geometry.py

@@ -4,10 +4,10 @@
 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
 
 import math
-from typing import List, Union, Tuple
+from typing import List, Union, Tuple, Optional
 import numpy as np
 import cv2
-from .common_types import BoundingBox, Polygon4P, RotatedBbox
+from .common_types import BoundingBox, Polygon4P, RotatedBbox, Bbox
 
 __all__ = ['rbbox_to_polygon', 'bbox_to_polygon', 'polygon_to_bbox', 'polygon_to_rbbox',
            'resolve_enclosing_bbox', 'resolve_enclosing_bbox', 'fit_rbbox', 'rotate_boxes', 'rotate_abs_boxes',
@@ -38,6 +38,10 @@ def polygon_to_rbbox(polygon: Polygon4P) -> RotatedBbox:
     return fit_rbbox(cnt)
 
 
+def bbox_to_rbbox(bbox: Bbox) -> RotatedBbox:
+    return (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2, bbox[2] - bbox[0], bbox[3] - bbox[1], 0
+
+
 def resolve_enclosing_bbox(bboxes: Union[List[BoundingBox], np.ndarray]) -> Union[BoundingBox, np.ndarray]:
     """Compute enclosing bbox either from:
 
@@ -129,10 +133,41 @@ def rotate_abs_boxes(boxes: np.ndarray, angle: float, img_shape: Tuple[int, int]
     return rotated_boxes
 
 
+def remap_boxes(boxes: np.ndarray, orig_shape: Tuple[int, int], dest_shape: Tuple[int, int]) -> np.ndarray:
+    """ Remaps a batch of RotatedBbox (x, y, w, h, alpha) expressed for an origin_shape to a destination_shape,
+    This does not impact the absolute shape of the boxes
+
+    Args:
+        boxes: (N, 5) array of RELATIVE RotatedBbox (x, y, w, h, alpha)
+        orig_shape: shape of the origin image
+        dest_shape: shape of the destination image
+
+    Returns:
+        A batch of rotated boxes (N, 5): (x, y, w, h, alpha) 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_width, orig_height = orig_shape
+    dest_width, dest_height = dest_shape
+    mboxes = boxes.copy()
+    mboxes[:, 0] = ((boxes[:, 0] * orig_height) + (dest_height - orig_height) / 2) / dest_height
+    mboxes[:, 1] = ((boxes[:, 1] * orig_width) + (dest_width - orig_width) / 2) / dest_width
+    mboxes[:, 2] = boxes[:, 2] * orig_height / dest_height
+    mboxes[:, 3] = boxes[:, 3] * orig_width / dest_width
+    return mboxes
+
+
 def rotate_boxes(
     boxes: np.ndarray,
     angle: float = 0.,
-    min_angle: float = 1.
+    min_angle: float = 1.,
+    expand: bool = False,
+    orig_shape: Optional[Tuple[int, int]] = None,
+    mask_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.
@@ -141,28 +176,37 @@ def rotate_boxes(
         boxes: (N, 4) array of RELATIVE boxes
         angle: angle between -90 and +90 degrees
         min_angle: minimum angle to rotate boxes
+        expand: whether the image should be padded before the rotation
+        orig_shape: shape of the origin image
+        mask_shape: shape of the mask if the image is cropped after the rotation
 
     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 = np.apply_along_axis(bbox_to_rbbox, 1, boxes)
     # If small angle, return boxes (no rotation)
     if abs(angle) < min_angle or abs(angle) > 90 - min_angle:
         return boxes
+    if expand:
+        exp_shape = compute_expanded_shape(orig_shape, angle)
+        boxes = remap_boxes(boxes, orig_shape=orig_shape, dest_shape=exp_shape)
+        orig_shape = exp_shape  # in case a mask is used afterwards
     # 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))
+    centers = np.stack((boxes[:, 0], boxes[:, 1]), axis=-1)
+    rotated_centers = .5 + np.matmul(centers - .5, rotation_mat)
     x_center, y_center = rotated_centers[:, 0], rotated_centers[:, 1]
     # 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])), axis=1)
+    # Apply a mask if requested
+    if mask_shape is not None:
+        rotated_boxes = remap_boxes(rotated_boxes, orig_shape=orig_shape, dest_shape=mask_shape)
     return rotated_boxes
 
 

doctr/utils/visualization.py

@@ -99,7 +99,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,

test/common/test_models.py

@@ -5,7 +5,7 @@
 import cv2
 
 from doctr.io import reader, DocumentFile
-from doctr.models._utils import extract_crops, extract_rcrops, get_bitmap_angle, estimate_orientation
+from doctr.models._utils import extract_crops, extract_rcrops, get_bitmap_angle, estimate_orientation, rotate_image
 
 
 def test_extract_crops(mock_pdf):  # noqa: F811
@@ -89,6 +89,34 @@ def test_get_bitmap_angle(mock_bitmap):
     assert abs(angle - 30.) < 1.
 
 
+def test_rotate_image():
+    img = np.ones((32, 64, 3), dtype=np.float32)
+    rotated = rotate_image(img, 30.)
+    assert rotated.shape[:-1] == (32, 64)
+    assert rotated[0, 0, 0] == 0
+    assert rotated[0, :, 0].sum() > 1
+
+    # Expand
+    rotated = rotate_image(img, 30., expand=True)
+    assert rotated.shape[:-1] == (60, 72)
+    assert rotated[0, :, 0].sum() <= 1
+
+    # Expand with 90° rotation
+    rotated = rotate_image(img, 90., expand=True)
+    assert rotated.shape[:-1] == (64, 64)
+    assert rotated[0, :, 0].sum() <= 1
+
+    # Expand with mask
+    rotated = rotate_image(img, 30., expand=True, mask_shape=(40, 72))
+    assert rotated.shape[:-1] == (40, 72)
+    assert rotated[0, :, 0].sum() > 1
+
+
 def test_estimate_orientation(mock_image):
     angle = estimate_orientation(mock_image)
     assert abs(angle - 30.) < 1.
+
+    angle = estimate_orientation(mock_image)
+    rotated = rotate_image(mock_image, -angle)
+    angle_rotated = estimate_orientation(rotated)
+    assert abs(angle_rotated - 0.) < 1.

test/common/test_utils_geometry.py

@@ -28,25 +28,54 @@ def test_polygon_to_rbbox():
     assert all(abs(i - j) <= 1e-7 for (i, j) in zip(pred, target))
 
 
+def test_bbox_to_rbbox():
+    pred = geometry.bbox_to_rbbox((0, 0, 0.6, 0.4))
+    target = (0.3, 0.2, 0.6, 0.4, 0)
+    assert all(abs(i - j) <= 1e-7 for (i, j) in zip(pred, target))
+
+
 def test_resolve_enclosing_rbbox():
     pred = geometry.resolve_enclosing_rbbox([(.2, .2, .05, .05, 0), (.2, .2, .2, .2, 0)])[:4]
     target = (.2, .2, .2, .2)
     assert all(abs(i - j) <= 1e-7 for (i, j) in zip(pred, target))
 
 
+def test_remap_boxes():
+    pred = geometry.remap_boxes(np.array([[0.5, 0.5, 0.1, 0.1, 0.]]), (10, 10), (20, 20))
+    target = np.array([[0.5, 0.5, 0.05, 0.05, 0.]])
+    assert pred.all() == target.all()
+
+    pred = geometry.remap_boxes(np.array([[0.5, 0.5, 0.1, 0.1, 0.]]), (10, 10), (20, 10))
+    target = np.array([[0.5, 0.5, 0.1, 0.05, 0.]])
+    assert pred.all() == target.all()
+
+    pred = geometry.remap_boxes(np.array([[0.25, 0.5, 0.5, 0.33, 0.]]), (80, 30), (160, 30))
+    target = np.array([[0.375, 0.5, 0.25, 0.1, 0.]])
+    assert pred.all() == target.all()
+
+
 def test_rotate_boxes():
     boxes = np.array([[0.1, 0.1, 0.8, 0.3]])
+    rboxes = np.apply_along_axis(geometry.bbox_to_rbbox, 1, boxes)
     # Angle = 0
     rotated = geometry.rotate_boxes(boxes, angle=0.)
-    assert rotated.all() == boxes.all()
+    assert rotated.all() == rboxes.all()
     # Angle < 1:
     rotated = geometry.rotate_boxes(boxes, angle=0.5)
-    assert rotated.all() == boxes.all()
+    assert rotated.all() == rboxes.all()
     # Angle = 30
     rotated = geometry.rotate_boxes(boxes, angle=30)
     assert rotated.shape == (1, 5)
     assert rotated[0, 4] == 30.
 
+    boxes = np.array([[0., 0., 0.6, 0.2]])
+    # Angle = -90:
+    rotated = geometry.rotate_boxes(boxes, angle=-90, min_angle=0)
+    assert rotated.all() == np.array([[0.1, 0.7, 0.6, 0.2, -90.]]).all()
+    # Angle = 90
+    rotated = geometry.rotate_boxes(boxes, angle=+90, min_angle=0)
+    assert rotated.all() == np.array([[0.9, 0.3, 0.6, 0.2, 90.]]).all()
+
 
 def test_rotate_image():
     img = np.ones((32, 64, 3), dtype=np.float32)