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..844bbf6d03 100644 --- a/tests/common/test_utils_geometry.py +++ b/tests/common/test_utils_geometry.py @@ -324,6 +324,18 @@ 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) == [] diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index a0a79618b3..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]