diff --git a/supervision/geometry/utils.py b/supervision/geometry/utils.py index 3d5668fb..401f672e 100644 --- a/supervision/geometry/utils.py +++ b/supervision/geometry/utils.py @@ -9,7 +9,7 @@ def get_polygon_center(polygon: np.ndarray) -> Point: This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. - + The center is calculated as center of frame. polygon -> polygon, where p[i] = p[i + 1] + p[i] / 2, with mass = length of vector p[i + 1] - p[i] @@ -32,10 +32,11 @@ def get_polygon_center(polygon: np.ndarray) -> Point: Point(x=1, y=1) ``` """ + polygon = polygon.astype(np.float32) shifted_polygon = np.roll(polygon, 1, axis=0) points = (shifted_polygon + polygon) / 2 vectors = shifted_polygon - polygon - mass = (vectors[:, 0] ** 2 + vectors[:, 1] ** 2) ** 0.5 - mass = np.array([mass, mass]).T - center = (np.sum(points * mass, axis=0) / np.sum(mass) * 2).round() + mass = np.sum(vectors ** 2, axis=1) ** 0.5 + center = ((mass @ points) / np.sum(mass)).round() + return Point(x=center[0], y=center[1]) diff --git a/test/geometry/test_utils.py b/test/geometry/test_utils.py index 39c45216..5457463b 100644 --- a/test/geometry/test_utils.py +++ b/test/geometry/test_utils.py @@ -1,24 +1,32 @@ +import numpy as np import pytest -import numpy as np -from supervision.geometry.core import Point, Vector +from supervision.geometry.core import Point from supervision.geometry.utils import get_polygon_center @pytest.mark.parametrize( - "polygon, expected_result", - [ - (np.array([[0, 0], [0, 2], [2, 2], [2, 0]]), Point(x=1, y=1)), - (np.array([[0, 0], [3, 4], [6, 0]]), Point(x=3, y=1)), - (np.array([[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [5, 2]]), Point(x=2, y=2)), - (np.array([[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [4, 4], [4, 0]]), Point(x=2, y=2)), - (np.array([[0, 2], [2, 4], [4, 2], [2, 0]]), Point(x=2, y=2)), - (np.array([[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [0, 1000]]), Point(x=0, y=500)), - (np.array([[0, 0], [13, 200], [0, 150]]), Point(x=4, y=100)), - (np.array([[0, 0], [0, 1], [1, 1], [1, 2], [2, 2], [2, 3], [3, 3], [3, 0]]), Point(x=2, y=1)), - ], + "polygon, expected_result", + [ + (np.array([[0, 0], [0, 2], [2, 2], [2, 0]]), Point(x=1, y=1)), + (np.array([[0, 0], [3, 4], [6, 0]]), Point(x=3, y=1)), + (np.array([[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [5, 2]]), Point(x=2, y=2)), + ( + np.array([[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [4, 4], [4, 0]]), + Point(x=2, y=2), + ), + (np.array([[0, 2], [2, 4], [4, 2], [2, 0]]), Point(x=2, y=2)), + ( + np.array([[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [0, 1000]]), + Point(x=0, y=500), + ), + (np.array([[0, 0], [13, 200], [0, 150]]), Point(x=4, y=100)), + ( + np.array([[0, 0], [0, 1], [1, 1], [1, 2], [2, 2], [2, 3], [3, 3], [3, 0]]), + Point(x=2, y=1), + ), + ], ) def test_get_polygon_center(polygon: np.ndarray, expected_result: Point) -> None: - result = get_polygon_center(polygon) - assert result == expected_result - \ No newline at end of file + result = get_polygon_center(polygon) + assert result == expected_result