add: test_mask_non_max_merge
This commit is contained in:
parent
ae77077853
commit
4977fd64b1
|
|
@ -8,6 +8,7 @@ from supervision.detection.overlap_filter import (
|
|||
box_non_max_suppression,
|
||||
group_overlapping_boxes,
|
||||
mask_non_max_suppression,
|
||||
mask_non_max_merge,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -447,3 +448,185 @@ def test_mask_non_max_suppression(
|
|||
predictions=predictions, masks=masks, iou_threshold=iou_threshold
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, masks, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty((0, 6)),
|
||||
np.empty((0, 5, 5)),
|
||||
0.5,
|
||||
[],
|
||||
DoesNotRaise(),
|
||||
), # empty predictions and masks
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
[[0]],
|
||||
DoesNotRaise(),
|
||||
), # single mask with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
[[0]],
|
||||
DoesNotRaise(),
|
||||
), # single mask with category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
[[0], [1]],
|
||||
DoesNotRaise(),
|
||||
), # two masks non-overlapping with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.4,
|
||||
[[0, 1]],
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with no category, merge
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.6,
|
||||
[[0, 1]],
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with no category, no merge
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0], [0, 0, 0, 0, 0.9, 1]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.4,
|
||||
[[0], [1]],
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with different categories
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0], [0, 0, 0, 0, 0.9, 0]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.4,
|
||||
[[0, 1]],
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with same category
|
||||
],
|
||||
)
|
||||
def test_mask_non_max_merge(
|
||||
predictions: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: List[List[int]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = mask_non_max_merge(
|
||||
predictions=predictions, masks=masks, iou_threshold=iou_threshold
|
||||
)
|
||||
sorted_result = sorted([sorted(group) for group in result])
|
||||
sorted_expected_result = sorted([sorted(group) for group in expected_result])
|
||||
assert sorted_result == sorted_expected_result
|
||||
|
|
|
|||
Loading…
Reference in New Issue