Skip to content

[feature] Add preserve_original_coords option for straighten_pages=True#2108

Open
saad-rd11 wants to merge 5 commits into
mindee:mainfrom
saad-rd11:bug
Open

[feature] Add preserve_original_coords option for straighten_pages=True#2108
saad-rd11 wants to merge 5 commits into
mindee:mainfrom
saad-rd11:bug

Conversation

@saad-rd11

Copy link
Copy Markdown

Closes #2107

side_by_side

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: _straighten_pages() now records the straightening transform as one composite affine matrix per page. The forward chain is expansion pad, rotation, aspect-ratio pad, crop, and the stored inverse is inv(C @ R @ P) where R is the actual cv2.getRotationMatrix2D matrix used for the warp, so there is no reconstructed rotation center that can drift. The crop offset is computed analytically by projecting the page's content corners through the rotation matrix, replacing the previous remove_image_padding pixel scan. This makes the crop deterministic and independent of page content (the pixel scan can over-crop on pages with dark edges, and interpolation fringe pixels make it jitter by a pixel across environments).

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 added 4 commits July 8, 2026 14:05
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant