[feature] Add preserve_original_coords option for straighten_pages=True

doctr/models/kie_predictor/base.py

@@ -30,6 +30,8 @@ class _KIEPredictor(_OCRPredictor):
         ignore_regions: optional list of layout class names to ignore during detection/recognition. If provided, the
             layout model will be used to locate the regions of the specified classes, and these regions will
             be masked out (filled with black) before passing the pages to the detection/recognition modules.
+        preserve_original_coords: if True and straighten_pages is True, bounding boxes are mapped back to the
+            original page coordinates. Useful for redaction and annotation.
         kwargs: keyword args of `DocumentBuilder`
     """
 
@@ -44,6 +46,7 @@ def __init__(
         symmetric_pad: bool = True,
         detect_orientation: bool = False,
         ignore_regions: Collection[str] | None = None,
+        preserve_original_coords: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -53,6 +56,7 @@ def __init__(
             symmetric_pad,
             detect_orientation,
             ignore_regions,
+            preserve_original_coords=preserve_original_coords,
             **kwargs,
         )
 

doctr/models/kie_predictor/pytorch.py

@@ -36,6 +36,8 @@ class KIEPredictor(nn.Module, _KIEPredictor):
             page. Doing so will slightly deteriorate the overall latency.
         detect_language: if True, the language prediction will be added to the predictions for each
             page. Doing so will slightly deteriorate the overall latency.
+        preserve_original_coords: if True and straighten_pages is True, bounding boxes are mapped back to the
+            original page coordinates. Useful for redaction and annotation.
         layout_predictor: optional layout detection module
         **kwargs: keyword args of `DocumentBuilder`
     """
@@ -50,6 +52,7 @@ def __init__(
         symmetric_pad: bool = True,
         detect_orientation: bool = False,
         detect_language: bool = False,
+        preserve_original_coords: bool = False,
         layout_predictor: LayoutPredictor | None = None,
         **kwargs: Any,
     ) -> None:
@@ -63,6 +66,7 @@ def __init__(
             preserve_aspect_ratio,
             symmetric_pad,
             detect_orientation,
+            preserve_original_coords=preserve_original_coords,
             **kwargs,
         )
         self.detect_orientation = detect_orientation
@@ -110,9 +114,12 @@ def forward(
             general_pages_orientations = None
             origin_pages_orientations = None
         if self.straighten_pages:
+            _orig_shapes = origin_page_shapes
+            _orig_pages = list(pages)
             pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)
             # update page shapes after straightening
             origin_page_shapes = [page.shape[:2] for page in pages]
+            _straight_shapes = origin_page_shapes
 
             # Detect layout regions on the pages
             regions = self.layout_predictor(pages, **kwargs) if self.layout_predictor is not None else None
@@ -192,6 +199,9 @@ def forward(
             languages_dict,
             regions,
         )
+
+        if self.straighten_pages and self.preserve_original_coords:
+            out = self._remap_to_original_coords(out, _orig_shapes, _straight_shapes, _orig_pages)  # type: ignore[assignment]
         return out
 
     @staticmethod

doctr/models/predictor/base.py

@@ -8,8 +8,15 @@
 
 import numpy as np
 
+from doctr.io.elements import Document, Word
 from doctr.models.builder import DocumentBuilder
-from doctr.utils.geometry import extract_crops, extract_rcrops, remove_image_padding, rotate_image
+from doctr.utils.geometry import (
+    extract_crops,
+    extract_rcrops,
+    remove_image_padding,
+    rotate_image,
+    straighten_page,
+)
 
 from .._utils import estimate_orientation, mask_boxes, rectify_crops, rectify_loc_preds
 from ..classification import crop_orientation_predictor, page_orientation_predictor
@@ -34,6 +41,8 @@
         ignore_regions: optional list of layout class names to ignore during detection/recognition. If provided, the
             layout model will be used to locate the regions of the specified classes, and these regions will
             be masked out (filled with black) before passing the pages to the detection/recognition modules.
+        preserve_original_coords: if True and straighten_pages is True, bounding boxes are mapped back to the
+            original page coordinates. Useful for redaction and annotation.
         **kwargs: keyword args of `DocumentBuilder`
     """
 
@@ -48,10 +57,12 @@
         symmetric_pad: bool = True,
         detect_orientation: bool = False,
         ignore_regions: Collection[str] | None = None,
+        preserve_original_coords: bool = False,
         **kwargs: Any,
     ) -> None:
         self.assume_straight_pages = assume_straight_pages
         self.straighten_pages = straighten_pages
+        self.preserve_original_coords = preserve_original_coords
         self._page_orientation_disabled = kwargs.pop("disable_page_orientation", False)
         self._crop_orientation_disabled = kwargs.pop("disable_crop_orientation", False)
         self.crop_orientation_predictor = (
@@ -134,11 +145,19 @@
                 for seq_map, general_orientation in zip(seg_maps, general_pages_orientations)
             ]
         )
-        return [
-            # expand if height and width are not equal, then remove the padding
-            remove_image_padding(rotate_image(page, angle, expand=page.shape[0] != page.shape[1]))
-            for page, angle in zip(pages, origin_pages_orientations)
-        ]
+        if not self.preserve_original_coords:
+            return [
+                remove_image_padding(rotate_image(page, angle, expand=page.shape[0] != page.shape[1]))
+                for page, angle in zip(pages, origin_pages_orientations)
+            ]
+
+        self._straighten_m_inv = []
+        out = []
+        for page, angle in zip(pages, origin_pages_orientations):
+            straightened, m_inv = straighten_page(page, angle)
+            self._straighten_m_inv.append(m_inv)
+            out.append(straightened)
+        return out
 
     @staticmethod
     def _generate_crops(
@@ -212,6 +231,78 @@
 
         return loc_preds, text_preds, crop_orientation_preds
 
+    def _remap_to_original_coords(
+        self,
+        out: Document,
+        orig_shapes: list[tuple[int, int]],
+        straight_shapes: list[tuple[int, int]],
+        orig_pages: list[np.ndarray] | None = None,
+    ) -> Document:
+        """Remap word geometries from straightened-page coordinates back to original page coordinates.
+
+        This method is called after the document builder when straighten_pages and
+        preserve_original_coords are both True. It applies the inverse straightening
+        transform to each word's geometry so that bounding boxes align with the input image.
+
+        Supports both Page (blocks -> lines -> words) and KIEPage (predictions dict)
+        structures, identified by the presence of a ``predictions`` attribute.
+
+        When ``orig_pages`` is provided, each page's image is also restored to the original
+        (pre-straightening) input so that ``.show()`` renders boxes on the correct image.
+
+        Args:
+            out: the document returned by the builder
+            orig_shapes: original (pre-straightening) page shapes (H, W)
+            straight_shapes: straightened page shapes (H, W)
+            orig_pages: optional list of original page images to restore on ``page.page``
+
+        Returns
+        -------
+        the document with remapped word geometries
+
+        """
+        for pidx, (page, m_inv) in enumerate(zip(out.pages, self._straighten_m_inv)):
+            sh, sw = straight_shapes[pidx]
+            oh, ow = orig_shapes[pidx]
+
+            # Collect all geometry-bearing objects (supports Page & KIEPage)
+            if hasattr(page, "predictions"):
+                objs: list[Word] = []
+                for preds in page.predictions.values():
+                    objs.extend(preds)
+            else:
+                objs = [word for block in page.blocks for line in block.lines for word in line.words]
+
+            for word in objs:
+                pts = np.asarray(word.geometry, dtype=np.float64).reshape(-1, 2)
+                was_straight = pts.shape[0] == 2
+                if was_straight:
+                    (x0, y0), (x1, y1) = pts
+                    pts = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype=np.float64)
+
+                pts[:, 0] *= sw
+                pts[:, 1] *= sh
+                homo = np.column_stack([pts, np.ones(len(pts))])
+                orig = (homo @ m_inv.T)[:, :2]
+                orig[:, 0] = orig[:, 0].clip(0, ow - 1) / ow
+                orig[:, 1] = orig[:, 1].clip(0, oh - 1) / oh
+
+                if was_straight:
+                    word.geometry = (
+                        (float(orig[:, 0].min()), float(orig[:, 1].min())),
+                        (float(orig[:, 0].max()), float(orig[:, 1].max())),
+                    )
+                else:
+                    word.geometry = tuple(tuple(r) for r in orig.tolist())  # type: ignore[assignment]
+
+            # Restore original page image and dimensions so .show() and export()
+            # render boxes in the correct frame (visualize_page scales geometry by dimensions)
+            if orig_pages is not None:
+                out.pages[pidx].page = orig_pages[pidx]
+                out.pages[pidx].dimensions = orig_shapes[pidx]
+
+        return out
+
     def add_hook(self, hook: Callable) -> None:
         """Add a hook to the predictor
 

doctr/models/predictor/pytorch.py

@@ -38,6 +38,8 @@ class OCRPredictor(nn.Module, _OCRPredictor):
             page. Doing so will slightly deteriorate the overall latency.
         detect_language: if True, the language prediction will be added to the predictions for each
             page. Doing so will slightly deteriorate the overall latency.
+        preserve_original_coords: if True and straighten_pages is True, the returned bounding boxes are mapped back
+            to the original page coordinate space. Useful for redaction and annotation.
         layout_predictor: optional layout detection module
         table_predictor: optional table structure recognition module. Requires `layout_predictor`: table
             regions are located by the layout model, cropped, and passed to this module. Words falling inside a
@@ -55,6 +57,7 @@ def __init__(
         symmetric_pad: bool = True,
         detect_orientation: bool = False,
         detect_language: bool = False,
+        preserve_original_coords: bool = False,
         layout_predictor: LayoutPredictor | None = None,
         table_predictor: TablePredictor | None = None,
         **kwargs: Any,
@@ -69,6 +72,7 @@ def __init__(
             preserve_aspect_ratio,
             symmetric_pad,
             detect_orientation,
+            preserve_original_coords=preserve_original_coords,
             **kwargs,
         )
         self.detect_orientation = detect_orientation
@@ -121,9 +125,12 @@ def forward(
             general_pages_orientations = None
             origin_pages_orientations = None
         if self.straighten_pages:
+            _orig_shapes = origin_page_shapes
+            _orig_pages = list(pages)
             pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)
             # update page shapes after straightening
             origin_page_shapes = [page.shape[:2] for page in pages]
+            _straight_shapes = origin_page_shapes
 
             # Detect layout regions on the pages
             regions = self.layout_predictor(pages, **kwargs) if self.layout_predictor is not None else None
@@ -191,6 +198,9 @@ def forward(
             regions,
             tables,
         )
+
+        if self.straighten_pages and self.preserve_original_coords:
+            out = self._remap_to_original_coords(out, _orig_shapes, _straight_shapes, _orig_pages)
         return out
 
     def _tables_from_regions(

doctr/utils/geometry.py

@@ -20,6 +20,7 @@
     "compute_expanded_shape",
     "rotate_image",
     "remove_image_padding",
+    "straighten_page",
     "estimate_page_angle",
     "convert_to_relative_coords",
     "rotate_abs_geoms",
@@ -398,6 +399,71 @@ def remove_image_padding(image: np.ndarray) -> np.ndarray:
     return image[rmin : rmax + 1, cmin : cmax + 1]
 
 
+def straighten_page(
+    page: np.ndarray,
+    angle: float,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Straighten a page by rotating it and cropping to content, returning the straightened image
+    and the inverse affine matrix that maps straightened-page coordinates back to the original page.
+
+    Args:
+        page: the original page image (H, W, C)
+        angle: rotation angle in degrees (between -90 and +90)
+
+    Returns:
+        a tuple (straightened, M_inv) where straightened is the page rotated and cropped to the
+        non-padded region, and M_inv is a (3, 3) array mapping points in the straightened image's
+        pixel space (x, y, 1) back to the original page's pixel space.
+    """
+    h, w = page.shape[:2]
+    expand = h != w
+
+    # Padding to preserve content after rotation
+    if expand:
+        exp = compute_expanded_shape((h, w), angle)
+        h_pad = int(max(0, np.ceil(exp[0] - h)))
+        w_pad = int(max(0, np.ceil(exp[1] - w)))
+    else:
+        h_pad = w_pad = 0
+
+    pt, pb = h_pad // 2, h_pad - h_pad // 2
+    pl, pr = w_pad // 2, w_pad - w_pad // 2
+
+    exp_img = np.pad(page, ((pt, pb), (pl, pr), (0, 0)))
+    ph_pad, pw_pad = exp_img.shape[:2]
+
+    # Rotate around the padded image center
+    rot_mat = cv2.getRotationMatrix2D((pw_pad / 2, ph_pad / 2), angle, 1.0)
+    rotated = cv2.warpAffine(exp_img, rot_mat, (pw_pad, ph_pad))
+
+    # Aspect-ratio padding (applied only on right/bottom so analytic crop stays simple)
+    if expand:
+        if (rotated.shape[0] / rotated.shape[1]) > (h / w):
+            w_pad2 = int(rotated.shape[0] * w / h - rotated.shape[1])
+            rotated = np.pad(rotated, ((0, 0), (0, w_pad2), (0, 0)))
+        else:
+            h_pad2 = int(rotated.shape[1] * h / w - rotated.shape[0])
+            rotated = np.pad(rotated, ((0, h_pad2), (0, 0), (0, 0)))
+
+    # Analytic crop: project the four content corners through the rotation matrix
+    corners = np.array(
+        [[pl, pt, 1], [pl + w, pt, 1], [pl + w, pt + h, 1], [pl, pt + h, 1]],
+        dtype=np.float64,
+    ).T
+    rc = rot_mat @ corners
+    cx, cy = int(np.floor(rc[0].min())), int(np.floor(rc[1].min()))
+    cropped = rotated[cy:, cx:]
+
+    # Composite forward matrix: crop ∘ rotate ∘ pad
+    # cv2.warpAffine treats the rotation matrix as src→dst
+    c3 = np.array([[1, 0, -cx], [0, 1, -cy], [0, 0, 1]], dtype=np.float64)
+    r3 = np.vstack([rot_mat, [0, 0, 1]])
+    p3 = np.array([[1, 0, pl], [0, 1, pt], [0, 0, 1]], dtype=np.float64)
+    m_inv = np.linalg.inv(c3 @ r3 @ p3)
+
+    return cropped, m_inv
+
+
 def estimate_page_angle(polys: np.ndarray) -> float:
     """Takes a batch of rotated previously ORIENTED polys (N, 4, 2) (rectified by the classifier) and return the
     estimated angle ccw in degrees

tests/common/test_utils_geometry.py

@@ -326,3 +326,38 @@ def test_extract_rcrops(mock_pdf, assume_horizontal):
 
     # No box
     assert geometry.extract_rcrops(doc_img, np.zeros((0, 4, 2)), assume_horizontal=assume_horizontal) == []
+
+
+@pytest.mark.parametrize("angle", [5, 12, -5, -12, 90 + 13, 180 + 13, 270 + 13])
+@pytest.mark.parametrize("shape", [(800, 600), (600, 800)])
+def test_straighten_page_inverse(angle, shape):
+    """M_inv returned by straighten_page is the exact inverse of the
+    pad -> rotate -> crop pipeline, accurate to within interpolation noise.
+
+    Sub-degree angles are excluded: at very small rotations every fiducial
+    pixel is blended by interpolation, so exact-colour matching finds nothing.
+    """
+    h, w = shape
+    page = np.ones((h, w, 3), dtype=np.uint8) * 255
+    page[99:102, 99:102] = (255, 0, 0)
+    page[99:102, 499:502] = (0, 255, 0)
+    page[699:702, 99:102] = (0, 0, 255)
+    page[699:702, 499:502] = (255, 255, 0)
+    page[399:402, 299:302] = (255, 0, 255)
+    dots = [(100, 100), (500, 100), (100, 700), (500, 700), (300, 400)]
+    colours = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
+
+    straightened, m_inv = geometry.straighten_page(page, angle)
+
+    errors = []
+    for (dx, dy), colour in zip(dots, colours):
+        mask = np.all(straightened == colour, axis=-1)
+        found = np.argwhere(mask)
+        if len(found) == 0:
+            continue
+        fy, fx = found.mean(axis=0)
+        recovered = (np.array([float(fx), float(fy), 1.0]) @ m_inv.T)[:2]
+        errors.append(np.linalg.norm(recovered - np.array([float(dx), float(dy)])))
+
+    assert len(errors) > 0, "No fiducial dots found -- interpolation blended all pixels"
+    assert max(errors) < 0.6, f"Max remap error {max(errors):.4f}px exceeds 0.6px threshold"