diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py index c69453a78c..8742a3f0ce 100644 --- a/doctr/utils/geometry.py +++ b/doctr/utils/geometry.py @@ -464,7 +464,7 @@ def extract_crops(img: np.ndarray, boxes: np.ndarray) -> list[np.ndarray]: _boxes[:, [1, 3]] *= h _boxes = _boxes.round().astype(int) # Add last index - _boxes[2:] += 1 + _boxes[:, 2:] += 1 return deepcopy([img[box[1] : box[3], box[0] : box[2]] for box in _boxes]) diff --git a/tests/common/test_utils_geometry.py b/tests/common/test_utils_geometry.py index d595875728..c00f90049b 100644 --- a/tests/common/test_utils_geometry.py +++ b/tests/common/test_utils_geometry.py @@ -272,6 +272,12 @@ def test_extract_crops(mock_pdf): # Identity assert np.all(doc_img == geometry.extract_crops(doc_img, np.array([[0, 0, 1, 1]], dtype=np.float32))[0]) + # Identical boxes must yield identical crops regardless of their position in the batch + gradient_img = np.tile(np.arange(100, dtype=np.uint8).reshape(100, 1, 1), (1, 100, 3)) + same_box = [0.1, 0.1, 0.2, 0.2] + identical_crops = geometry.extract_crops(gradient_img, np.array([same_box, same_box, same_box], dtype=np.float32)) + assert all(np.array_equal(crop, identical_crops[0]) for crop in identical_crops) + # No box assert geometry.extract_crops(doc_img, np.zeros((0, 4))) == []