diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py index 17b04a7391..fea02e375e 100644 --- a/doctr/utils/geometry.py +++ b/doctr/utils/geometry.py @@ -337,7 +337,13 @@ def estimate_page_angle(polys: np.ndarray) -> float: yleft = polys[:, 0, 1] + polys[:, 3, 1] xright = polys[:, 1, 0] + polys[:, 2, 0] yright = polys[:, 1, 1] + polys[:, 2, 1] - return float(np.median(np.arctan((yleft - yright) / (xright - xleft))) * 180 / np.pi) # Y axis from top to bottom! + with np.errstate(divide="raise", invalid="raise"): + try: + return float( + np.median(np.arctan((yleft - yright) / (xright - xleft)) * 180 / np.pi) # Y axis from top to bottom! + ) + except FloatingPointError: + return 0.0 def convert_to_relative_coords(geoms: np.ndarray, img_shape: Tuple[int, int]) -> np.ndarray: diff --git a/tests/common/test_utils_geometry.py b/tests/common/test_utils_geometry.py index e9e8a688ee..6cdd9f0533 100644 --- a/tests/common/test_utils_geometry.py +++ b/tests/common/test_utils_geometry.py @@ -159,6 +159,10 @@ def test_estimate_page_angle(): rotated_polys = geometry.rotate_boxes(straight_polys, angle=20, orig_shape=(512, 512)) angle = geometry.estimate_page_angle(rotated_polys) assert np.isclose(angle, 20) + # Test divide by zero / NaN + invalid_poly = np.array([[[0.5, 0.5], [0.5, 0.5], [0.5, 0.5], [0.5, 0.5]]]) + angle = geometry.estimate_page_angle(invalid_poly) + assert angle == 0.0 def test_extract_crops(mock_pdf):