[feature] Add preserve_original_coords option for straighten_pages=True#2108
Open
saad-rd11 wants to merge 5 commits into
Open
[feature] Add preserve_original_coords option for straighten_pages=True#2108saad-rd11 wants to merge 5 commits into
straighten_pages=True#2108saad-rd11 wants to merge 5 commits into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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: withpreserve_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 returnedWord.geometryis 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(defaultFalse, 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.pyuntouched.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 isinv(C @ R @ P)where R is the actualcv2.getRotationMatrix2Dmatrix 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 previousremove_image_paddingpixel 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. AfterDocumentBuilderreturns theDocument, 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 @ Pand 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 withpreserve_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:
rotate_image()andremove_image_padding()inutils/geometry.pyto 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.Page-level metadata instead of rewritingword.geometryin 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()hardcodesorig_shape=(1024, 1024), which distorts the reading-order rotation for non-square pages whenassume_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.