[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.93617% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.71%. Comparing base (1a1018c) to head (c860794).

Files with missing lines	Patch %	Lines
doctr/models/predictor/base.py	98.21%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2108      +/-   ##
==========================================
+ Coverage   96.69%   96.71%   +0.01%     
==========================================
  Files         168      168              
  Lines        8969     9059      +90     
==========================================
+ Hits         8673     8761      +88     
- Misses        296      298       +2     

Flag	Coverage Δ
unittests	96.71% <98.93%> (+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.

[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.

[felixdittrich92]

Thanks again @saad-rd11 👍

One last round and we should be good to merge 👍

[felixdittrich92]

comment like before should be keept:

# 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]))

[saad-rd11]

Restored verbatim in e582b32.

[felixdittrich92]

Move all imports to the top of the file where the other imports are placed

[saad-rd11]

Resolved in 24d121e; the functions carrying local imports were removed in the test consolidation, so no function-level imports remain in the file.

[felixdittrich92]

two tests should be enough here:

• The two "smoke" tests (one for ocr_predictor / one for kie_predictor) - test_ocr_predictor_straighten_with_preserve_original_coords and test_kie_predictor_straighten_with_preserve_original_coords
• As a test file we could use the mock_payslip fixture (

  doctr/tests/conftest.py

  Line 46 in 1a1018c

  def mock_payslip(tmpdir_factory):

  )

[saad-rd11]

Done in 24d121e: standalone tests and both module-scoped fixtures removed, replaced with test_ocr_predictor_straighten_with_preserve_original_coords (mock_payslip: values, reference geometry per the test_trained_ocr_predictor pattern, page/dimensions restoration) and test_kie_predictor_straighten_with_preserve_original_coords. For the KIE case I used mock_tilted_payslip so the remap is exercised on genuine skew end to end — and that fixture immediately paid off: its real-world dimensions (2291×2025 at 30°) surfaced a floating-point bug in the analytic crop that the synthetic fiducial shapes never triggered. A content corner projected to x = −0.033, floor gave −1, and Python's negative slice rotated[:, -1:] silently collapsed the crop to a 1-pixel strip. Fixed in 6f3f737 with a max(0, ...) clamp on cx/cy, and the failing shape/angle is pinned in the fiducial parametrization (now 22 cases, max error < 0.6px). Both smoke tests pass locally (~25s KIE, comparable to test_trained_kie_predictor).

[felixdittrich92]

I was still thinking about a possible way to avoid the duplicated logic now in straighten_page and remove_image_padding(rotate_image(..)) but let's keep it as is ftm 👍

[saad-rd11]

👍

[felixdittrich92]

comment

# straightened_height, straightened_width
sh, sw =
# original_height, original_width
oh, ow = 

or rename variables original_height, original_width, straightened_height, straightened_width

[saad-rd11]

Added as inline comments in e582b32 (# straightened_height, straightened_width / # original_height, original_width).

[saad-rd11]

Thanks @felixdittrich92! I am on it, all four are quick, will push today. 👍

[saad-rd11]

Thanks @felixdittrich92: all items addressed, latest round pushed (three commits):

e582b32: restore flag-off path comment + variable name clarification
6f3f737: fix analytic crop: clamp cx/cy non-negative (floating-point corner projection; details on the thread) + regression case in the fiducial parametrization
24d121e: tests consolidated to the two smoke tests (−310/+54)

Test suite: 59 green locally (22 fiducial inverse cases < 0.6px, both new smoke tests passing), ruff clean, mypy unchanged from baseline. As throughout: preserve_original_coords=False remains byte-identical to main.

On Codacy: same situation as before, the 3 notices are D213/D406/D407 from its Prospector profile, which enforce docstring conventions contradicting the repo's requested format. Nothing actionable that I can see.