xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Updated algorithm

This commit is contained in:
sharingan000 2024-04-03 20:40:54 +10:00
parent 7ded0bd493
commit cd81c1cef4
2 changed files with 29 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
supervision/geometry/utils.py
View File

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

40
test/geometry/test_utils.py
View File

@ -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
result = get_polygon_center(polygon)
assert result == expected_result