From ae4a7926ad33ac0955ffb00024333e7d2ca23352 Mon Sep 17 00:00:00 2001 From: AryanRG Date: Fri, 16 Oct 2026 23:31:21 +0530 Subject: [PATCH 1/2] [fix] Handle zero-width horizontal crops --- doctr/utils/geometry.py | 3 ++ tests/common/test_utils_geometry.py | 30 +++++++++++++ tests/pytorch/test_models_zoo_pt.py | 68 +++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+) diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py index db16830f1e..a49cd0d406 100644 --- a/doctr/utils/geometry.py +++ b/doctr/utils/geometry.py @@ -570,6 +570,9 @@ def extract_rcrops( # Divide the points into left and right left_points = box[box[:, 0] < centroid[0]] right_points = box[box[:, 0] >= centroid[0]] + if len(left_points) == 0 or len(right_points) == 0: + crops.append(np.zeros((0, 0, *img.shape[2:]), dtype=img.dtype)) + continue # Sort the left points according to the y-axis left_points = left_points[np.argsort(left_points[:, 1])] diff --git a/tests/common/test_utils_geometry.py b/tests/common/test_utils_geometry.py index 9191c69c5f..29cf59d4d4 100644 --- a/tests/common/test_utils_geometry.py +++ b/tests/common/test_utils_geometry.py @@ -328,6 +328,36 @@ def test_extract_rcrops(mock_pdf, assume_horizontal): assert geometry.extract_rcrops(doc_img, np.zeros((0, 4, 2)), assume_horizontal=assume_horizontal) == [] +@pytest.mark.parametrize("case", ["x0", "xmid", "x1", "repeated"]) +@pytest.mark.parametrize("assume_horizontal", [True, False], ids=["horizontal", "oriented"]) +@pytest.mark.parametrize("shape, dtype", [((100, 200, 3), np.uint8), ((100, 200), np.float32)], ids=["rgb8", "gray32"]) +def test_extract_rcrops_zero_width(case, assume_horizontal, shape, dtype): + img = (np.arange(np.prod(shape)).reshape(shape) % 256).astype(dtype) + left = [[0.125, 0.25], [0.25, 0.25], [0.25, 0.75], [0.125, 0.75]] + right = [[0.625, 0.25], [0.75, 0.25], [0.75, 0.75], [0.625, 0.75]] + if case in ("x0", "xmid", "x1"): + x = {"x0": 0, "xmid": 0.5, "x1": 1}[case] + middle = [[x, 0.25], [x, 0.25], [x, 0.75], [x, 0.75]] + else: + middle = [[0.5, 0.5]] * 4 + polys = np.array([left, middle, right], dtype=np.float32) + original_img, original_polys = img.copy(), polys.copy() + references = [ + geometry.extract_rcrops(img, polys[idx : idx + 1], assume_horizontal=assume_horizontal)[0] for idx in (0, 2) + ] + crops = geometry.extract_rcrops(img, polys, assume_horizontal=assume_horizontal) + assert len(crops) == 3 + for crop, reference in zip((crops[0], crops[2]), references): + assert crop.size > 0 + assert crop.shape == reference.shape + assert crop.dtype == reference.dtype == img.dtype + np.testing.assert_array_equal(crop, reference) + assert crops[1].shape == (0, 0, *img.shape[2:]) + assert crops[1].dtype == img.dtype + np.testing.assert_array_equal(img, original_img) + np.testing.assert_array_equal(polys, original_polys) + + @pytest.mark.parametrize("angle", [5, 12, -5, -12, 90 + 13, 180 + 13, 270 + 13]) @pytest.mark.parametrize("shape", [(800, 600), (600, 800), (700, 700)]) def test_straighten_page_inverse(angle, shape): diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index a0a79618b3..02a024dfc2 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -184,6 +184,74 @@ def forward(self, crops, **kwargs): assert [word.objectness_score for word in words] == pytest.approx([0.9, 0.8]) +@pytest.mark.parametrize( + "predictor_cls, all_collapsed", + [ + pytest.param(OCRPredictor, False, id="ocr-mixed"), + pytest.param(OCRPredictor, True, id="ocr-all-collapsed"), + pytest.param(KIEPredictor, False, id="kie-mixed"), + pytest.param(KIEPredictor, True, id="kie-all-collapsed"), + ], +) +def test_predictors_drop_zero_width_quadrilaterals(predictor_cls, all_collapsed): + polygons = np.array( + [ + [[0.125, 0.25], [0.25, 0.25], [0.25, 0.75], [0.125, 0.75]], + [[0.5, 0.25], [0.5, 0.25], [0.5, 0.75], [0.5, 0.75]], + [[0.625, 0.25], [0.75, 0.25], [0.75, 0.75], [0.625, 0.75]], + ], + dtype=np.float32, + ) + score_points = np.zeros((3, 1, 2), dtype=np.float32) + score_points[:, 0, 1] = [0.9, 0.1, 0.8] + detections = np.concatenate([polygons, score_points], axis=1) + if all_collapsed: + detections = detections[1:2].copy() + + class _FixedDetPredictor(nn.Module): + def forward(self, pages, **kwargs): + return [{CLASS_NAME: detections.copy()} for _ in pages] + + class _CheckingRecoPredictor(nn.Module): + def __init__(self): + super().__init__() + self.counts = [] + + def forward(self, crops, **kwargs): + assert all(crop.size > 0 for crop in crops) + self.counts.append(len(crops)) + return [(f"word{idx}", 1.0) for idx in range(len(crops))] + + reco = _CheckingRecoPredictor() + predictor = predictor_cls( + _FixedDetPredictor(), + reco, + assume_straight_pages=False, + disable_page_orientation=True, + disable_crop_orientation=True, + detect_orientation=False, + straighten_pages=False, + resolve_lines=False, + resolve_blocks=False, + ) + result = predictor([np.full((100, 200, 3), 255, dtype=np.uint8)]) + assert type(result) is (KIEDocument if predictor_cls is KIEPredictor else Document) + assert len(result.pages) == 1 + page = result.pages[0] + words = ( + page.predictions.get(CLASS_NAME, []) + if predictor_cls is KIEPredictor + else [word for block in page.blocks for line in block.lines for word in line.words] + ) + words = sorted(words, key=lambda word: min(point[0] for point in word.geometry)) + assert [min(point[0] for point in word.geometry) for word in words] == pytest.approx( + [] if all_collapsed else [0.125, 0.625] + ) + assert [word.objectness_score for word in words] == pytest.approx([] if all_collapsed else [0.9, 0.8]) + assert [word.value for word in words] == ([] if all_collapsed else ["word0", "word1"]) + assert reco.counts == ([0] if all_collapsed else [2]) + + def test_ocrpredictor_layout(mock_pdf, mock_vocab, mock_payslip): det_predictor = DetectionPredictor( PreProcessor(output_size=(512, 512), batch_size=2), From 126e127391936fdcddc792d82259f163be8c86aa Mon Sep 17 00:00:00 2001 From: AryanRG Date: Wed, 21 Oct 2026 12:38:39 +0530 Subject: [PATCH 2/2] [tests] Consolidate zero-width crop regressions --- tests/common/test_utils_geometry.py | 42 ++++-------- tests/pytorch/test_models_zoo_pt.py | 102 ++++++++-------------------- 2 files changed, 39 insertions(+), 105 deletions(-) diff --git a/tests/common/test_utils_geometry.py b/tests/common/test_utils_geometry.py index 29cf59d4d4..844bbf6d03 100644 --- a/tests/common/test_utils_geometry.py +++ b/tests/common/test_utils_geometry.py @@ -324,40 +324,22 @@ def test_extract_rcrops(mock_pdf, assume_horizontal): assert all(isinstance(crop, np.ndarray) for crop in croped_imgs) assert all(crop.ndim == 3 for crop in croped_imgs) + # Zero-width boxes yield an empty crop, without shifting the other crops + collapsed_boxes = np.array( + [[[x, 0.25], [x, 0.25], [x, 0.75], [x, 0.75]] for x in (0.0, 0.5, 1.0)] + [[[0.5, 0.5]] * 4], + dtype=np.float32, + ) + croped_imgs = geometry.extract_rcrops( + doc_img, np.concatenate([rel_boxes, collapsed_boxes]), assume_horizontal=assume_horizontal + ) + assert len(croped_imgs) == num_crops + len(collapsed_boxes) + assert all(crop.size > 0 for crop in croped_imgs[:num_crops]) + assert all(crop.shape == (0, 0, 3) and crop.dtype == doc_img.dtype for crop in croped_imgs[num_crops:]) + # No box assert geometry.extract_rcrops(doc_img, np.zeros((0, 4, 2)), assume_horizontal=assume_horizontal) == [] -@pytest.mark.parametrize("case", ["x0", "xmid", "x1", "repeated"]) -@pytest.mark.parametrize("assume_horizontal", [True, False], ids=["horizontal", "oriented"]) -@pytest.mark.parametrize("shape, dtype", [((100, 200, 3), np.uint8), ((100, 200), np.float32)], ids=["rgb8", "gray32"]) -def test_extract_rcrops_zero_width(case, assume_horizontal, shape, dtype): - img = (np.arange(np.prod(shape)).reshape(shape) % 256).astype(dtype) - left = [[0.125, 0.25], [0.25, 0.25], [0.25, 0.75], [0.125, 0.75]] - right = [[0.625, 0.25], [0.75, 0.25], [0.75, 0.75], [0.625, 0.75]] - if case in ("x0", "xmid", "x1"): - x = {"x0": 0, "xmid": 0.5, "x1": 1}[case] - middle = [[x, 0.25], [x, 0.25], [x, 0.75], [x, 0.75]] - else: - middle = [[0.5, 0.5]] * 4 - polys = np.array([left, middle, right], dtype=np.float32) - original_img, original_polys = img.copy(), polys.copy() - references = [ - geometry.extract_rcrops(img, polys[idx : idx + 1], assume_horizontal=assume_horizontal)[0] for idx in (0, 2) - ] - crops = geometry.extract_rcrops(img, polys, assume_horizontal=assume_horizontal) - assert len(crops) == 3 - for crop, reference in zip((crops[0], crops[2]), references): - assert crop.size > 0 - assert crop.shape == reference.shape - assert crop.dtype == reference.dtype == img.dtype - np.testing.assert_array_equal(crop, reference) - assert crops[1].shape == (0, 0, *img.shape[2:]) - assert crops[1].dtype == img.dtype - np.testing.assert_array_equal(img, original_img) - np.testing.assert_array_equal(polys, original_polys) - - @pytest.mark.parametrize("angle", [5, 12, -5, -12, 90 + 13, 180 + 13, 270 + 13]) @pytest.mark.parametrize("shape", [(800, 600), (600, 800), (700, 700)]) def test_straighten_page_inverse(angle, shape): diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index 02a024dfc2..276ff367d3 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -156,13 +156,26 @@ def test_predictors_on_empty_batch(mock_vocab): assert out.export() == {"pages": []} +@pytest.mark.parametrize("assume_straight_pages", [True, False]) @pytest.mark.parametrize("predictor_cls", [OCRPredictor, KIEPredictor]) -def test_predictors_keep_objectness_scores_aligned(predictor_cls): - # The middle box lies on the right border of the page: its crop is empty, so it is dropped before recognition - boxes = np.array( - [[0.1, 0.1, 0.3, 0.2, 0.9], [1.0, 0.5, 1.0, 0.6, 0.1], [0.5, 0.1, 0.7, 0.2, 0.8]], - dtype=np.float32, - ) +def test_predictors_keep_objectness_scores_aligned(predictor_cls, assume_straight_pages): + # The middle box lies on the right border of the page and has no width: its crop is empty, + # so it is dropped before recognition + if assume_straight_pages: + boxes = np.array( + [[0.1, 0.1, 0.3, 0.2, 0.9], [1.0, 0.5, 1.0, 0.6, 0.1], [0.5, 0.1, 0.7, 0.2, 0.8]], + dtype=np.float32, + ) + else: + # Polygons with the objectness score in the last row + boxes = np.array( + [ + [[0.1, 0.1], [0.3, 0.1], [0.3, 0.2], [0.1, 0.2], [0.9, 0.9]], + [[1.0, 0.5], [1.0, 0.5], [1.0, 0.6], [1.0, 0.6], [0.1, 0.1]], + [[0.5, 0.1], [0.7, 0.1], [0.7, 0.2], [0.5, 0.2], [0.8, 0.8]], + ], + dtype=np.float32, + ) class _FixedDetPredictor(nn.Module): def forward(self, pages, **kwargs): @@ -172,7 +185,14 @@ class _DummyRecoPredictor(nn.Module): def forward(self, crops, **kwargs): return [(f"word{idx}", 1.0) for idx in range(len(crops))] - predictor = predictor_cls(_FixedDetPredictor(), _DummyRecoPredictor(), assume_straight_pages=True) + # Without page orientation, rotated boxes are cropped with `assume_horizontal=True` + predictor = predictor_cls( + _FixedDetPredictor(), + _DummyRecoPredictor(), + assume_straight_pages=assume_straight_pages, + disable_page_orientation=True, + disable_crop_orientation=True, + ) page = predictor([np.full((100, 200, 3), 255, dtype=np.uint8)]).pages[0] if predictor_cls is KIEPredictor: words = page.predictions[CLASS_NAME] @@ -184,74 +204,6 @@ def forward(self, crops, **kwargs): assert [word.objectness_score for word in words] == pytest.approx([0.9, 0.8]) -@pytest.mark.parametrize( - "predictor_cls, all_collapsed", - [ - pytest.param(OCRPredictor, False, id="ocr-mixed"), - pytest.param(OCRPredictor, True, id="ocr-all-collapsed"), - pytest.param(KIEPredictor, False, id="kie-mixed"), - pytest.param(KIEPredictor, True, id="kie-all-collapsed"), - ], -) -def test_predictors_drop_zero_width_quadrilaterals(predictor_cls, all_collapsed): - polygons = np.array( - [ - [[0.125, 0.25], [0.25, 0.25], [0.25, 0.75], [0.125, 0.75]], - [[0.5, 0.25], [0.5, 0.25], [0.5, 0.75], [0.5, 0.75]], - [[0.625, 0.25], [0.75, 0.25], [0.75, 0.75], [0.625, 0.75]], - ], - dtype=np.float32, - ) - score_points = np.zeros((3, 1, 2), dtype=np.float32) - score_points[:, 0, 1] = [0.9, 0.1, 0.8] - detections = np.concatenate([polygons, score_points], axis=1) - if all_collapsed: - detections = detections[1:2].copy() - - class _FixedDetPredictor(nn.Module): - def forward(self, pages, **kwargs): - return [{CLASS_NAME: detections.copy()} for _ in pages] - - class _CheckingRecoPredictor(nn.Module): - def __init__(self): - super().__init__() - self.counts = [] - - def forward(self, crops, **kwargs): - assert all(crop.size > 0 for crop in crops) - self.counts.append(len(crops)) - return [(f"word{idx}", 1.0) for idx in range(len(crops))] - - reco = _CheckingRecoPredictor() - predictor = predictor_cls( - _FixedDetPredictor(), - reco, - assume_straight_pages=False, - disable_page_orientation=True, - disable_crop_orientation=True, - detect_orientation=False, - straighten_pages=False, - resolve_lines=False, - resolve_blocks=False, - ) - result = predictor([np.full((100, 200, 3), 255, dtype=np.uint8)]) - assert type(result) is (KIEDocument if predictor_cls is KIEPredictor else Document) - assert len(result.pages) == 1 - page = result.pages[0] - words = ( - page.predictions.get(CLASS_NAME, []) - if predictor_cls is KIEPredictor - else [word for block in page.blocks for line in block.lines for word in line.words] - ) - words = sorted(words, key=lambda word: min(point[0] for point in word.geometry)) - assert [min(point[0] for point in word.geometry) for word in words] == pytest.approx( - [] if all_collapsed else [0.125, 0.625] - ) - assert [word.objectness_score for word in words] == pytest.approx([] if all_collapsed else [0.9, 0.8]) - assert [word.value for word in words] == ([] if all_collapsed else ["word0", "word1"]) - assert reco.counts == ([0] if all_collapsed else [2]) - - def test_ocrpredictor_layout(mock_pdf, mock_vocab, mock_payslip): det_predictor = DetectionPredictor( PreProcessor(output_size=(512, 512), batch_size=2),