308 lines
9.2 KiB
Python
308 lines
9.2 KiB
Python
from contextlib import ExitStack as DoesNotRaise
|
|
from typing import List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from supervision.detection.utils import (
|
|
clip_boxes,
|
|
filter_polygons_by_area,
|
|
non_max_suppression,
|
|
process_roboflow_result,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"predictions, iou_threshold, expected_result, exception",
|
|
[
|
|
(
|
|
np.empty(shape=(0, 5)),
|
|
0.5,
|
|
np.array([]),
|
|
DoesNotRaise(),
|
|
), # single box with no category
|
|
(
|
|
np.array([[10.0, 10.0, 40.0, 40.0, 0.8]]),
|
|
0.5,
|
|
np.array([True]),
|
|
DoesNotRaise(),
|
|
), # single box with no category
|
|
(
|
|
np.array([[10.0, 10.0, 40.0, 40.0, 0.8, 0]]),
|
|
0.5,
|
|
np.array([True]),
|
|
DoesNotRaise(),
|
|
), # single box with category
|
|
(
|
|
np.array(
|
|
[
|
|
[10.0, 10.0, 40.0, 40.0, 0.8],
|
|
[15.0, 15.0, 40.0, 40.0, 0.9],
|
|
]
|
|
),
|
|
0.5,
|
|
np.array([False, True]),
|
|
DoesNotRaise(),
|
|
), # two boxes with no category
|
|
(
|
|
np.array(
|
|
[
|
|
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
|
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
|
|
]
|
|
),
|
|
0.5,
|
|
np.array([True, True]),
|
|
DoesNotRaise(),
|
|
), # two boxes with different category
|
|
(
|
|
np.array(
|
|
[
|
|
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
|
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
|
|
]
|
|
),
|
|
0.5,
|
|
np.array([False, True]),
|
|
DoesNotRaise(),
|
|
), # two boxes with same category
|
|
(
|
|
np.array(
|
|
[
|
|
[0.0, 0.0, 30.0, 40.0, 0.8],
|
|
[5.0, 5.0, 35.0, 45.0, 0.9],
|
|
[10.0, 10.0, 40.0, 50.0, 0.85],
|
|
]
|
|
),
|
|
0.5,
|
|
np.array([False, True, False]),
|
|
DoesNotRaise(),
|
|
), # three boxes with no category
|
|
(
|
|
np.array(
|
|
[
|
|
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
|
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
|
|
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
|
|
]
|
|
),
|
|
0.5,
|
|
np.array([True, True, True]),
|
|
DoesNotRaise(),
|
|
), # three boxes with same category
|
|
(
|
|
np.array(
|
|
[
|
|
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
|
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
|
|
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
|
|
]
|
|
),
|
|
0.5,
|
|
np.array([False, True, True]),
|
|
DoesNotRaise(),
|
|
), # three boxes with different category
|
|
],
|
|
)
|
|
def test_non_max_suppression(
|
|
predictions: np.ndarray,
|
|
iou_threshold: float,
|
|
expected_result: Optional[np.ndarray],
|
|
exception: Exception,
|
|
) -> None:
|
|
with exception:
|
|
result = non_max_suppression(
|
|
predictions=predictions, iou_threshold=iou_threshold
|
|
)
|
|
assert np.array_equal(result, expected_result)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"boxes_xyxy, frame_resolution_wh, expected_result",
|
|
[
|
|
(
|
|
np.empty(shape=(0, 4)),
|
|
(1280, 720),
|
|
np.empty(shape=(0, 4)),
|
|
),
|
|
(
|
|
np.array([[1.0, 1.0, 1279.0, 719.0]]),
|
|
(1280, 720),
|
|
np.array([[1.0, 1.0, 1279.0, 719.0]]),
|
|
),
|
|
(
|
|
np.array([[-1.0, 1.0, 1279.0, 719.0]]),
|
|
(1280, 720),
|
|
np.array([[0.0, 1.0, 1279.0, 719.0]]),
|
|
),
|
|
(
|
|
np.array([[1.0, -1.0, 1279.0, 719.0]]),
|
|
(1280, 720),
|
|
np.array([[1.0, 0.0, 1279.0, 719.0]]),
|
|
),
|
|
(
|
|
np.array([[1.0, 1.0, 1281.0, 719.0]]),
|
|
(1280, 720),
|
|
np.array([[1.0, 1.0, 1280.0, 719.0]]),
|
|
),
|
|
(
|
|
np.array([[1.0, 1.0, 1279.0, 721.0]]),
|
|
(1280, 720),
|
|
np.array([[1.0, 1.0, 1279.0, 720.0]]),
|
|
),
|
|
],
|
|
)
|
|
def test_clip_boxes(
|
|
boxes_xyxy: np.ndarray,
|
|
frame_resolution_wh: Tuple[int, int],
|
|
expected_result: np.ndarray,
|
|
) -> None:
|
|
result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
|
|
assert np.array_equal(result, expected_result)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"polygons, min_area, max_area, expected_result, exception",
|
|
[
|
|
(
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
None,
|
|
None,
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
DoesNotRaise(),
|
|
), # single polygon without area constraints
|
|
(
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
50,
|
|
None,
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
DoesNotRaise(),
|
|
), # single polygon with min_area constraint
|
|
(
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
None,
|
|
50,
|
|
[],
|
|
DoesNotRaise(),
|
|
), # single polygon with max_area constraint
|
|
(
|
|
[
|
|
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
|
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
|
],
|
|
200,
|
|
None,
|
|
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
|
DoesNotRaise(),
|
|
), # two polygons with min_area constraint
|
|
(
|
|
[
|
|
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
|
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
|
],
|
|
None,
|
|
200,
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
DoesNotRaise(),
|
|
), # two polygons with max_area constraint
|
|
(
|
|
[
|
|
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
|
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
|
],
|
|
200,
|
|
200,
|
|
[],
|
|
DoesNotRaise(),
|
|
), # two polygons with both area constraints
|
|
(
|
|
[
|
|
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
|
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
|
],
|
|
100,
|
|
100,
|
|
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
|
DoesNotRaise(),
|
|
), # two polygons with min_area and
|
|
# max_area equal to the area of the first polygon
|
|
(
|
|
[
|
|
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
|
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
|
],
|
|
400,
|
|
400,
|
|
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
|
DoesNotRaise(),
|
|
), # two polygons with min_area and
|
|
# max_area equal to the area of the second polygon
|
|
],
|
|
)
|
|
def test_filter_polygons_by_area(
|
|
polygons: List[np.ndarray],
|
|
min_area: Optional[float],
|
|
max_area: Optional[float],
|
|
expected_result: List[np.ndarray],
|
|
exception: Exception,
|
|
) -> None:
|
|
with exception:
|
|
result = filter_polygons_by_area(
|
|
polygons=polygons, min_area=min_area, max_area=max_area
|
|
)
|
|
assert len(result) == len(expected_result)
|
|
for result_polygon, expected_result_polygon in zip(result, expected_result):
|
|
assert np.array_equal(result_polygon, expected_result_polygon)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"roboflow_result, class_list, expected_result, exception",
|
|
[
|
|
(
|
|
{"predictions": [], "image": {"width": 1000, "height": 1000}},
|
|
["person", "car", "truck"],
|
|
(np.empty((0, 4)), np.empty(0), np.empty(0), None),
|
|
DoesNotRaise(),
|
|
), # empty result
|
|
(
|
|
{
|
|
"predictions": [
|
|
{
|
|
"x": 200.0,
|
|
"y": 300.0,
|
|
"width": 50.0,
|
|
"height": 50.0,
|
|
"confidence": 0.9,
|
|
"class": "person",
|
|
}
|
|
],
|
|
"image": {"width": 1000, "height": 1000},
|
|
},
|
|
["person", "car", "truck"],
|
|
(
|
|
np.array([[175.0, 275.0, 225.0, 325.0]]),
|
|
np.array([0.9]),
|
|
np.array([0]),
|
|
None,
|
|
),
|
|
DoesNotRaise(),
|
|
), # single bounding box
|
|
],
|
|
)
|
|
def test_process_roboflow_result(
|
|
roboflow_result: dict,
|
|
class_list: List[str],
|
|
expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
|
|
exception: Exception,
|
|
) -> None:
|
|
with exception:
|
|
result = process_roboflow_result(
|
|
roboflow_result=roboflow_result, class_list=class_list
|
|
)
|
|
assert np.array_equal(result[0], expected_result[0])
|
|
assert np.array_equal(result[1], expected_result[1])
|
|
assert np.array_equal(result[2], expected_result[2])
|
|
assert (result[3] is None and expected_result[3] is None) or (
|
|
np.array_equal(result[3], expected_result[3])
|
|
)
|