[feature] Add preserve_original_coords option for straighten_pages=True

[saad-rd11]

Closes #2107

Left: skewed input. Middle: current behavior with straighten_pages=True, boxes come back flat in the internally straightened page's coordinate space, misaligned with the input. Right: with preserve_original_coords=True, boxes are mapped back to the input image, usable for redaction, annotation, and overlays.

Summary

When straighten_pages=True, detection runs on internally deskewed pages, and the returned Word.geometry is relative to that straightened image, a coordinate space the user never sees. For text extraction this doesn't matter, but for anything that needs the boxes to line up with the input image (redaction, annotation, highlighting), they are unusable, and the straightening transform is discarded inside _straighten_pages() so there is no way to recover it.

This PR adds an opt-in flag, preserve_original_coords (default False, zero behavior change when off), that remaps word geometries back to the coordinate space of the image the user passed in.

How it works

Two files changed plus one new test file. builder.py untouched.

models/predictor/base.py: When preserve_original_coords=False (default), _straighten_pages() runs the exact original remove_image_padding(rotate_image(...)) path, verified byte-identical to main's output across portrait/landscape pages at ±12°. When the flag is on, an analytic pad → rotate → crop path runs instead, recording the straightening transform as one composite affine matrix per page (inv(C @ R @ P), built from the actual cv2.getRotationMatrix2D matrix). The flag-on path needs its own analytic crop because a content-dependent pixel scan can't be inverted exactly.

models/predictor/pytorch.py: original page shapes are captured before straightening. After DocumentBuilder returns the Document, an optional post-processing pass converts each word polygon to absolute straightened-page pixels, applies the stored inverse matrix, clips to the original page bounds, and renormalizes.

The remap handles both geometry formats: 4-point polygons (assume_straight_pages=False) and 2-point boxes (assume_straight_pages=True). In the 2-point case the box is expanded to all 4 corners before the transform (rotating only the two stored diagonal points would give a wrong envelope) and returned as the axis-aligned envelope of the rotated corners. The envelope is a conservative superset of the true rotated region, which is the right behavior for redaction, and it preserves the 2-point format contract that downstream consumers expect.

Remapping after document building is deliberate. DocumentBuilder._sort_boxes() re-estimates the page angle from box geometry, so boxes remapped any earlier are detected as skewed and silently rotated straight again, undoing the correction. I verified this failure mode directly before settling on the post-builder approach. Doing the remap last means the detection, recognition, and building pipeline runs completely unmodified and the two mechanisms never interact.

Validation

Three independent checks, ordered from pure math to full pipeline. The test-based checks are included as pytest cases in this PR (tests/pytorch/test_preserve_original_coords.py). Full suite: 19 passed in 45s.

1. Analytic round-trip. Points pushed through the forward matrix C @ R @ P and back through the stored inverse recover to 2.8e-13 px. The inverse is exact by construction; this check confirms no composition-order or convention error.

2. Real-pixel fiducial test (test_straighten_inverse_fiducial, 14 parametrized cases, no model weights, runs in about 2 seconds). Colored 3x3 dots at known positions go through the same pad, warpAffine, and crop path as _straighten_pages, are located in the output by exact color match, and remapped through the stored matrix. This test makes no assumptions about OpenCV's matrix conventions; it measures where real pixels actually land, which is what caught two convention bugs during development. Result: max error 0.49 px across angles of plus and minus 5 and 12 degrees plus 103, 193, and 283 degrees (covering the 90/180/270 base-orientation compositions with fine skew), on both portrait and landscape pages. Sub-degree angles are excluded with a comment in the test: at those rotations interpolation blends every fiducial pixel, so exact-color matching finds nothing to measure.

3. End-to-end tests (pretrained db_resnet50 + crnn_vgg16_bn). Text is rendered at known positions, ground-truth word boxes are measured from the ink pixels of the clean render, the page is skewed by plus or minus 12 degrees, and the full predictor runs with preserve_original_coords=True. The GT boxes are transformed into the skewed frame (the frame the flag returns coordinates in) and compared against the remapped detections by IoU.

• test_preserve_original_coords_roundtrip (4 cases, assume_straight_pages=False, module-scoped predictor): mean IoU 0.894 to 0.910 across both page shapes and both skew signs, against a 0.4 threshold.
• test_preserve_original_coords_2point (1 case, assume_straight_pages=True): asserts the returned geometry stays 2-point and the envelope clears the same IoU threshold, exercising the 2-point expansion path end to end.

The remaining gap to IoU 1.0 is the detection model's own box localization, not the transform: check 2 bounds the transform's contribution at under half a pixel.

Notes for reviewers

I kept this deliberately minimal and non-invasive, but I'm happy to restructure toward whichever shape fits the codebase better. Two alternatives I considered:

• Extending rotate_image() and remove_image_padding() in utils/geometry.py to optionally return their transform matrices, so the transform logic has a single source of truth there instead of being partially inlined in _straighten_pages(). More invasive, but removes the duplicated pad, rotate, and crop logic this PR currently carries.
• Exposing the matrix as Page-level metadata instead of rewriting word.geometry in place, leaving the remap to users. Cleaner separation, but less useful out of the box for the redaction use case that motivated this.

The validation suite transfers unchanged to either variant.

Separately, while tracing the coordinate chain I found that _sort_boxes() hardcodes orig_shape=(1024, 1024), which distorts the reading-order rotation for non-square pages when assume_straight_pages=False. It doesn't affect this PR since the remap happens after the builder, so I'll file it as its own issue rather than mixing it in here.

[saad-rd11]

Pushed c4f27f4: the analytic crop is now gated behind the flag. With preserve_original_coords=False, _straighten_pages() takes the exact original remove_image_padding(rotate_image(...)) path — I verified this empirically by running both branches' _straighten_pages on identical inputs (portrait and landscape pages, ±12°, text near edges): outputs are byte-identical (np.array_equal per page). The analytic path only activates when the flag is on, since the matrix remap needs a deterministic, exactly invertible crop.

[felixdittrich92]

Hi @saad-rd11 👋

Thanks for the PR
I left some initial comments

Additional we should verify that this combination works with .show() afterwards correct

[felixdittrich92]

Hi @saad-rd11 👋

Thanks for the PR I left some initial comments

Additional we should verify that this combination works with .show() afterwards correct

Currently this will still forward the "rotation corrected pages" here: pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)

with your change you also need to override this to the original input images which then are passed to the builder if both straigthen_pages & preserve_original_coords are True

[saad-rd11]

@felixdittrich92 Thanks for the review!

All makes sense. I'll move the transform to geometry.py with a standalone test, lift the remap into _OCRPredictor so KIE inherits it, apply the simplifications, and verify .show() renders correctly with the flag on.

Will have it pushed shortly.

[codecov]

Codecov Report

❌ Patch coverage is 98.73418% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.71%. Comparing base (1a1018c) to head (3fb1b75).

Files with missing lines	Patch %	Lines
doctr/utils/geometry.py	96.55%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2108      +/-   ##
==========================================
+ Coverage   96.69%   96.71%   +0.01%     
==========================================
  Files         168      168              
  Lines        8969     9046      +77     
==========================================
+ Hits         8673     8749      +76     
- Misses        296      297       +1     

Flag	Coverage Δ
unittests	96.71% <98.73%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[saad-rd11]

Thanks for the review @felixdittrich92 all six comments addressed, commits linked on each thread. Rebased onto current main. I found two things during implementation worth mentioning:

KIE had no remap path at all. Moving the logic into _OCRPredictor (per your comment) meant it needed to work for KIE too, but KIEPage stores geometry in a predictions dict rather than the OCR blocks → lines → words tree. Handled it with duck-typing in the shared _remap_to_original_coords method (hasattr(page, "predictions")), and added a smoke test that runs KIE flag-on vs flag-off and asserts the geometries actually differ (L2 ~0.72–1.34), since a structural mismatch there would silently no-op rather than raise.

Verifying .show() caught a real bug as page.dimensions was left at the straightened value after the coordinate swap, so visualize_page scaled boxes by the wrong dimensions (~1.26×, shifting every box down-left). hOCR/XML export would have hit the same. Fixed by restoring dimensions alongside page.page, with an assertion in the roundtrip test. Before/after can been seen down this thread.

Rebase note: upstream dfb8f7c added ignore_regions as a new positional param to _OCRPredictor.init, which collided with the positional preserve_original_coords pass fixed with keyword passing.

Full suite green: 20 feature tests + 17 zoo, ruff clean, mypy unchanged from baseline. preserve_original_coords=False (default) is a no-op path, byte-identical to current behavior.

[saad-rd11]

Hi @saad-rd11 👋
Thanks for the PR I left some initial comments
Additional we should verify that this combination works with .show() afterwards correct

Currently this will still forward the "rotation corrected pages" here: pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)

with your change you also need to override this to the original input images which then are passed to the builder if both straigthen_pages & preserve_original_coords are True

Verified renders correctly with the flag on (image below), via docTR's visualize_page path.

This also surfaced a stale page.dimensions bug: after the coordinate swap, page.dimensions still held the straightened size, so visualize_page scaled boxes ~1.26× (down-left shift). Fixed in 12c9918, restored alongside page.page with a roundtrip assertion. hOCR/XML export would have hit the same.

BEFORE:

AFTER:

[saad-rd11]

Hi @felixdittrich92 quick note on the remaining Codacy flag: it's now raising D213 on the same docstring where it previously raised D212. These two rules are mutually exclusive in pydocstyle D212 requires the summary on the first line, D213 on the second so satisfying one always triggers the other.

[felixdittrich92]

if isinstance(page, KIEPage) nd hasattr(page, "predictions"):

[saad-rd11]

Fixed in 2ad882a changed to isinstance(page, KIEPage) with the import added to base.py. Verified no circular import: elements.py only imports from doctr.file_utils, doctr.utils.*, and top-level doctr, none of which import from the predictor modules.

[felixdittrich92]

Replace all `` with single ` for all the added docstrings / comments

[saad-rd11]

Fixed in fcc3273: all double backticks replaced with single across the added docstrings and comments.

[felixdittrich92]

Returns:
     the document with remapped word geometries

[saad-rd11]

Fixed in fcc3273: changed to Returns: with indented body, matching the file's existing convention.

[felixdittrich92]

A comment would be good

# manipulate the already built Document to restore the original pages / shapes and geometries

Additional self._remap_to_original_coords must also handle the geometry from the Layout Elements and table elements not only the word geometry

[saad-rd11]

Comment added verbatim in 46e6d56.

The layout/table/artefact geometry handling was addressed in 19869a8: extracted a _remap_geometry static helper and extended _remap_to_original_coords to walk every geometry-bearing object — Word, Line, Block (envelope), Artefact, LayoutElement, Table, TableCell, and Prediction. Verified against elements.py: KIEPage takes layout but not tables (line 618), so the walk uses guarded access and skips tables gracefully on KIE pages; Page handles both.

Coverage in be3db15: a synthetic-Document test constructs one of every geometry type (no pretrained layout models needed), runs the real _remap_to_original_coords with a realistic 12° rotation m_inv, and asserts every geometry changed plus a containment assertion that remapped Block envelopes still bound their remapped Words.

[felixdittrich92]

Square pages (h == w) will break the analytic crop.

[saad-rd11]

Confirmed and fixed in 2ad882a. Reproduced first: for (700, 700) at 12°, the expand = h != w guard skipped padding entirely, content corners rotated outside the canvas, and the fiducial remap error was ~90px.

Fix: removed the h != w guard, compute_expanded_shape is now called unconditionally. At angle ≈ 0 it returns the original shape (zero pad), so near-zero rotations are unchanged. The aspect-ratio padding branch guard was updated to h_pad > 0 or w_pad > 0 accordingly. The flag-off path was not touched: the "False is byte-identical" guarantee holds.

Square shape (700, 700) added to the fiducial parametrization: all 21 cases (3 shapes × 7 angles, including the 90/180/270 base-orientation compositions) pass with max error < 0.6px. One note: the (700,700) × 193° case initially found zero fiducials, the 180°+13° composition's resampling blended the 3×3 dots below exact-color-match threshold. Resolved by using 5×5 dots throughout; all cases now pass with no skips.

[felixdittrich92]

This makes the predictor stateful we should avoid saving it in self

[saad-rd11]

Fixed in 8f0c936: _straighten_pages now returns (pages, m_invs) instead of writing to self, both OCR and KIE forward() unpack the tuple and pass m_invs as a parameter to _remap_to_original_coords. The flag-off path returns an empty list to keep the return type uniform. No instance state remains: grep for _straighten_m_inv across doctr/ returns zero results.

[felixdittrich92]

Remove docstrings after the test function names - it's enough to add small comments at places where it's not clear / required

[saad-rd11]

Fixed in 4cad1c9: test docstrings removed; the one non-obvious rationale (sub-degree angle exclusion) kept as an inline comment.

[felixdittrich92]

This file should be included in tests/pytorch/test_models_zoo_pt.py extend the existing tests parameterization should be enough / maybe a geometry check afterwards but not a standalone test file :)

[saad-rd11]

Done in 2d96c4d and 6bcbc59: standalone file deleted, coverage folded into test_models_zoo_pt.py following its existing patterns: module-scoped fixtures for the pretrained predictors, test_preserve_original_coords_roundtrip (4 cases, mean IoU > 0.4, page.page and page.dimensions assertions), test_preserve_original_coords_2point (assume_straight_pages=True variant), test_preserve_original_coords_kie_smoke (flag-on vs flag-off geometry difference), plus the _remap_geometry unit test and the synthetic-Document walk test. All 25 zoo tests pass locally (321s with cached weights). The fiducial inverse test stays in tests/common/test_utils_geometry.py.

[saad-rd11]

@felixdittrich92 Sure! I will apply the style items, remove the predictor state, fold the e2e tests into the zoo parametrization, handle square pages (will add a square case to the fiducial test first), and extend the remap to Layout/Table element geometry. Should have it pushed in a day or two.

[saad-rd11]

Thanks for the second round @felixdittrich92 all eight comments addressed, replies with commits on each thread. The short version:

Regarding your square-page catch: I reproduced it first: (700, 700) at 12° skipped padding entirely because of the inherited h != w expand guard, content corners rotated off the canvas, and the remap error was around 90px. Fixed by removing the guard; compute_expanded_shape now decides unconditionally, which costs nothing at near-zero angles. Square shapes are now part of the fiducial test parametrization, all 21 cases pass under 0.6px.
The predictor is now stateless. _straighten_pages returns the inverse matrices instead of writing them to self, and they flow through forward() as locals into _remap_to_original_coords as a parameter. No instance state left grep for the old attribute comes back empty.

The remap now covers every geometry-bearing object, not just words: Line, Block, Artefact, LayoutElement, Table, TableCell, and Prediction, for both Page and KIEPage (checked against elements.py. KIEPage carries layout but not tables, so access is guarded). A synthetic-Document test builds one of every type and asserts each geometry actually changes after the remap, plus a containment check that remapped Block envelopes still bound their remapped Words. No pretrained layout models needed for that coverage.
Tests are folded into test_models_zoo_pt.py per your preference: module-scoped fixtures, the roundtrip/2-point/KIE cases, the walk test, and the _remap_geometry unit test. The standalone file is gone. The fiducial inverse test stays in test_utils_geometry.py. All 25 zoo tests pass locally, full suite is 61 green, ruff clean, mypy unchanged from baseline.

Style items (isinstance, backticks, Returns: format, the comment in kie_predictor, test docstrings) all applied as requested.
As before: preserve_original_coords=False is untouched: the default path stays byte-identical to main.

One note on Codacy: it now flags D406/D407 on the Returns: docstring, those rules enforce numpy-style formatting, i.e. they flag the exact format requested in review. Along with the earlier D213, all three notices come from Codacy's Prospector profile enforcing conventions that aren't in the repo's ruff config. Nothing actionable on the PR side as far as I can tell.