[feature] Add preserve_original_coords option for `straighten_pages=True` by saad-rd11 · Pull Request #2108 · mindee/doctr

saad-rd11 · 2026-07-08T10:26:08Z

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 · 2026-07-08T22:32:08Z

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 · 2026-07-09T07:00:15Z

+                for page, angle in zip(pages, origin_pages_orientations)
+            ]
+
+        self._straighten_m_inv = []


This should be outsourced to utils/geometry.py and only run if straighten_pages & preserve_original_coords == True (standalone test in test_utils_geometry.py) otherwise keep the current behaviour

Done in 7fd11bc moved to straighten_page() in geometry.py. Fiducial test moved to tests/common/test_utils_geometry.py.

felixdittrich92 · 2026-07-09T07:08:10Z

            tables,
        )
+
+        if self.straighten_pages and self.preserve_original_coords and self._straighten_m_inv:


This logic should be moved to _OCRPredictor as method where KIEPredictor also inerhits from - so same logic needs to be applied in kie_predictor/pytorch.py

looks like the and self._straighten_m_inv check can be dropped should always exist if self.straighten_pages and self.preserve_original_coords

Done in 16f6276 _remap_to_original_coords() on _OCRPredictor, called from both OCR and KIE forward(). See the top-level note re: KIE storing geometry in a predictions dict rather than blocks/lines/words.

looks like the and self._straighten_m_inv check can be dropped should always exist if self.straighten_pages and self.preserve_original_coords

Done in 47558c5

felixdittrich92 · 2026-07-09T07:10:05Z

            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


passing as kwarg in ocr_predictor is fine here ftm 👍

Thanks, left it as is.

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 · 2026-07-09T07:14:37Z

            general_pages_orientations = None
            origin_pages_orientations = None
        if self.straighten_pages:
+            _orig_shapes = [p.shape[:2] for p in pages]


_orig_shapes = origin_page_shapes 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]

Done in 47558c5 dropped the redundant guard and applied your origin_page_shapes capture pattern.

felixdittrich92 · 2026-07-09T07:17:02Z

            tables,
        )
+
+        if self.straighten_pages and self.preserve_original_coords and self._straighten_m_inv:


looks like the and self._straighten_m_inv check can be dropped should always exist if self.straighten_pages and self.preserve_original_coords

felixdittrich92 · 2026-07-09T07:26:55Z

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 · 2026-07-09T07:45:50Z

@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 · 2026-07-09T08:21:01Z

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.

…ed to original page coordinates

GT boxes were measured in the pre-skew frame but compared against detections in the skewed frame — M_skew transform added. Replaced fragile CC+merge GT with getTextSize+ink-tightening. Module-scoped fixture amortizes model load across cases. Halved parametrization.

…_coords When assume_straight_pages=True, word.geometry is stored as a 2-point box ((xmin,ymin),(xmax,ymax)). The remap loop must detect this, expand to 4 corners before the transform, and return the axis-aligned envelope.

When the flag is off (default), _straighten_pages() returns the exact original remove_image_padding(rotate_image(...)) one-liner — zero behavior change for existing users. The analytic crop with matrix capture only activates when preserve_original_coords=True, making the guarantee easy to verify at a glance.

Extract the analytic pad -> rotate -> aspect-pad -> crop -> matrix-capture pipeline into a standalone straighten_page() function in utils/geometry.py so it has a single source of truth independent of the OCR predictor. _OCRPredictor._straighten_pages() calls straighten_page() when preserve_original_coords=True; the flag-off path remains the original one-liner remove_image_padding(rotate_image(...)). The fiducial inverse test moves from test_preserve_original_coords.py to test_utils_geometry.py and exercises straighten_page() directly (no model weights, 14 parametrized cases, max error < 0.6px).

…m 5 / gap fix) Add the preserve_original_coords parameter to both _KIEPredictor and KIEPredictor so the flag propagates through the KIE init chain to _OCRPredictor, enabling the feature for KIE usage.

…ew item 2) The shared method on _OCRPredictor now handles both Page (blocks -> lines -> words) and KIEPage (predictions dict) structures via duck typing. Both OCRPredictor.forward() and KIEPredictor.forward() call it after the document builder.

…& 4) Drop the redundant and self._straighten_m_inv guard (guaranteed non-empty when both flags are true). Use origin_page_shapes directly for _orig_shapes instead of recomputing from pages, exactly as Felix suggested.

…ew item 6) _remap_to_original_coords now accepts an optional orig_pages parameter. When provided, each page.page is restored to the original input image after the geometry remap so that .show() renders boxes on the correct frame. Verified by an explicit np.array_equal assertion in the roundtrip test (all 4 cases).

…nch (review item 6) Two KIEPredictor runs (flag-on vs flag-off) on the same skewed input must produce different geometries. L2 diff confirmed 0.72-1.34, real skew shift.

visualize_page scales geometry by page.dimensions, which still held the straightened dimensions after the page swap, causing a 1.26x coordinate shift. Now dimensions is updated alongside page.page in the shared _remap_to_original_coords method so both .show() and hOCR export see correct scaling. Verified by explicit assertion in the roundtrip test.

D213: multi-line summary on second line D406: Returns without colon D407: dashed underline after Returns D413: blank line after last section

_OCRPredictor.__init__ gained ignore_regions as 7th positional param from upstream. OCRPredictor was passing preserve_original_coords positionally, which landed in the wrong slot. Switch to keyword arg.

saad-rd11 · 2026-07-09T18:58:32Z

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 · 2026-07-09T19:12:57Z

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 mentioned this pull request Jul 8, 2026

straighten_pages=True returns bounding boxes in straightened-page coordinates, making them unusable for redaction or annotation on the original document #2107

Open

felixdittrich92 reviewed Jul 9, 2026

View reviewed changes

felixdittrich92 requested changes Jul 9, 2026

View reviewed changes

felixdittrich92 added this to the 1.1.0 milestone Jul 9, 2026

felixdittrich92 self-assigned this Jul 9, 2026

saad-rd11 added 15 commits July 9, 2026 23:17

fix: add preserve_original_coords flag to remap boxes from straighten…

8c6df56

…ed to original page coordinates

test: add preserve_original_coords inverse matrix accuracy test

ed8194d

chore: remove unused rotate_image import

3941f0c

Add preserve_original_coords kwarg to KIE predictor chain (review ite…

2136377

…m 5 / gap fix) Add the preserve_original_coords parameter to both _KIEPredictor and KIEPredictor so the flag propagates through the KIE init chain to _OCRPredictor, enabling the feature for KIE usage.

Add KIE smoke test verifying remap fires for hasattr(predictions) bra…

b844302

…nch (review item 6) Two KIEPredictor runs (flag-on vs flag-off) on the same skewed input must produce different geometries. L2 diff confirmed 0.72-1.34, real skew shift.

Fix Codacy docstring formatting issues on _remap_to_original_coords

0fdceb0

D213: multi-line summary on second line D406: Returns without colon D407: dashed underline after Returns D413: blank line after last section

Fix positional arg mismatch for preserve_original_coords after rebase

8980535

_OCRPredictor.__init__ gained ignore_regions as 7th positional param from upstream. OCRPredictor was passing preserve_original_coords positionally, which landed in the wrong slot. Switch to keyword arg.

saad-rd11 force-pushed the bug branch from 4b60310 to 8980535 Compare July 9, 2026 18:46

Fix Codacy D212: multi-line summary on first line

3fb1b75

Uh oh!

Conversation

saad-rd11 commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

How it works

Validation

Notes for reviewers

Uh oh!

saad-rd11 commented Jul 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

felixdittrich92 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

felixdittrich92 commented Jul 9, 2026

Uh oh!

saad-rd11 commented Jul 9, 2026

Uh oh!

codecov Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

saad-rd11 commented Jul 9, 2026

Uh oh!

saad-rd11 commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

saad-rd11 commented Jul 8, 2026 •

edited

Loading

codecov Bot commented Jul 9, 2026 •

edited

Loading

saad-rd11 commented Jul 9, 2026 •

edited

Loading