[fix] Handle zero-width horizontal crops

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])]

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) == []
 

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]