xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
supervision/test/detection/utils/test_masks.py

729 lines
21 KiB
Python
Raw Blame History

from __future__ import annotations
from contextlib import ExitStack as DoesNotRaise
import numpy as np
import numpy.typing as npt
import pytest
from supervision.detection.utils.masks import (
calculate_masks_centroids,
contains_holes,
contains_multiple_segments,
filter_segments_by_distance,
move_masks,
)
@pytest.mark.parametrize(
"masks, offset, resolution_wh, expected_result, exception",
[
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([0, 0]),
(4, 4),
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([-1, -1]),
(4, 4),
np.array(
[
[
[True, True, False, False],
[True, True, False, False],
[False, False, False, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([-2, -2]),
(4, 4),
np.array(
[
[
[True, False, False, False],
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([-3, -3]),
(4, 4),
np.array(
[
[
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([-2, -1]),
(4, 4),
np.array(
[
[
[True, False, False, False],
[True, False, False, False],
[False, False, False, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([-1, -2]),
(4, 4),
np.array(
[
[
[True, True, False, False],
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([-2, 2]),
(4, 4),
np.array(
[
[
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
[True, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([3, 3]),
(4, 4),
np.array(
[
[
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
[False, False, False, False],
]
],
dtype=bool,
),
DoesNotRaise(),
),
(
np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
np.array([3, 3]),
(6, 6),
np.array(
[
[
[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, False, False, False, True, True],
[False, False, False, False, True, True],
]
],
dtype=bool,
),
DoesNotRaise(),
),
],
)
def test_move_masks(
masks: np.ndarray,
offset: np.ndarray,
resolution_wh: tuple[int, int],
expected_result: np.ndarray,
exception: Exception,
) -> None:
with exception:
result = move_masks(masks=masks, offset=offset, resolution_wh=resolution_wh)
np.testing.assert_array_equal(result, expected_result)
@pytest.mark.parametrize(
"masks, expected_result, exception",
[
(
np.array(
[
[
[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0],
]
]
),
np.array([[0, 0]]),
DoesNotRaise(),
), # single mask with all zeros
(
np.array(
[
[
[1, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 1],
]
]
),
np.array([[2, 2]]),
DoesNotRaise(),
), # single mask with all ones
(
np.array(
[
[
[0, 1, 1, 0],
[1, 1, 1, 1],
[1, 1, 1, 1],
[0, 1, 1, 0],
]
]
),
np.array([[2, 2]]),
DoesNotRaise(),
), # single mask with symmetric ones
(
np.array(
[
[
[0, 0, 0, 0],
[0, 0, 1, 1],
[0, 0, 1, 1],
[0, 0, 0, 0],
]
]
),
np.array([[3, 2]]),
DoesNotRaise(),
), # single mask with asymmetric ones
(
np.array(
[
[
[0, 1, 1, 0],
[1, 1, 1, 1],
[1, 1, 1, 1],
[0, 1, 1, 0],
],
[
[0, 0, 0, 0],
[0, 0, 1, 1],
[0, 0, 1, 1],
[0, 0, 0, 0],
],
]
),
np.array([[2, 2], [3, 2]]),
DoesNotRaise(),
), # two masks
],
)
def test_calculate_masks_centroids(
masks: np.ndarray,
expected_result: np.ndarray,
exception: Exception,
) -> None:
with exception:
result = calculate_masks_centroids(masks=masks)
assert np.array_equal(result, expected_result)
@pytest.mark.parametrize(
"mask, expected_result, exception",
[
(
np.array([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 0, 0], [0, 1, 1, 0]]).astype(
bool
),
False,
DoesNotRaise(),
), # foreground object in one continuous piece
(
np.array([[1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0], [0, 1, 1, 0]]).astype(
bool
),
False,
DoesNotRaise(),
), # foreground object in 2 separate elements
(
np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]).astype(
bool
),
False,
DoesNotRaise(),
), # no foreground pixels in mask
(
np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]).astype(
bool
),
False,
DoesNotRaise(),
), # only foreground pixels in mask
(
np.array([[1, 1, 1, 0], [1, 0, 1, 0], [1, 1, 1, 0], [0, 0, 0, 0]]).astype(
bool
),
True,
DoesNotRaise(),
), # foreground object has 1 hole
(
np.array([[1, 1, 1, 0], [1, 0, 1, 1], [1, 1, 0, 1], [0, 1, 1, 1]]).astype(
bool
),
True,
DoesNotRaise(),
), # foreground object has 2 holes
],
)
def test_contains_holes(
mask: npt.NDArray[np.bool_], expected_result: bool, exception: Exception
) -> None:
with exception:
result = contains_holes(mask)
assert result == expected_result
@pytest.mark.parametrize(
"mask, connectivity, expected_result, exception",
[
(
np.array([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 0, 0], [0, 1, 1, 0]]).astype(
bool
),
4,
False,
DoesNotRaise(),
), # foreground object in one continuous piece
(
np.array([[1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0], [0, 1, 1, 0]]).astype(
bool
),
4,
True,
DoesNotRaise(),
), # foreground object in 2 separate elements
(
np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]).astype(
bool
),
4,
False,
DoesNotRaise(),
), # no foreground pixels in mask
(
np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]).astype(
bool
),
4,
False,
DoesNotRaise(),
), # only foreground pixels in mask
(
np.array([[1, 1, 1, 0], [1, 0, 1, 1], [1, 1, 0, 1], [0, 1, 1, 1]]).astype(
bool
),
4,
False,
DoesNotRaise(),
), # foreground object has 2 holes, but is in single piece
(
np.array([[1, 1, 0, 0], [1, 1, 0, 1], [1, 0, 1, 1], [0, 0, 1, 1]]).astype(
bool
),
4,
True,
DoesNotRaise(),
), # foreground object in 2 elements with respect to 4-way connectivity
(
np.array([[1, 1, 0, 0], [1, 1, 0, 1], [1, 0, 1, 1], [0, 0, 1, 1]]).astype(
bool
),
8,
False,
DoesNotRaise(),
), # foreground object in single piece with respect to 8-way connectivity
(
np.array([[1, 1, 0, 0], [1, 1, 0, 1], [1, 0, 1, 1], [0, 0, 1, 1]]).astype(
bool
),
5,
None,
pytest.raises(ValueError),
), # Incorrect connectivity parameter value, raises ValueError
],
)
def test_contains_multiple_segments(
mask: npt.NDArray[np.bool_],
connectivity: int,
expected_result: bool,
exception: Exception,
) -> None:
with exception:
result = contains_multiple_segments(mask=mask, connectivity=connectivity)
assert result == expected_result
@pytest.mark.parametrize(
"mask, connectivity, mode, absolute_distance, relative_distance, expected_result, exception", # noqa: E501
[
# single component, unchanged
(
np.array(
[
[0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
8,
"edge",
2.0,
None,
np.array(
[
[0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
DoesNotRaise(),
),
# two components, edge distance 2, kept with abs=1
(
np.array(
[
[0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 1],
[0, 1, 1, 1, 0, 1],
[0, 1, 1, 1, 0, 1],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
8,
"edge",
2.0,
None,
np.array(
[
[0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 1],
[0, 1, 1, 1, 0, 1],
[0, 1, 1, 1, 0, 1],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
DoesNotRaise(),
),
# centroid mode, far centroids, dropped with small relative threshold
(
np.array(
[
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 1],
[0, 0, 0, 1, 1, 1],
],
dtype=bool,
),
8,
"centroid",
None,
0.3, # diagonal ~8.49, threshold ~2.55, centroid gap ~4.24
np.array(
[
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
DoesNotRaise(),
),
# centroid mode, larger relative threshold, kept
(
np.array(
[
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 1],
[0, 0, 0, 1, 1, 1],
],
dtype=bool,
),
8,
"centroid",
None,
0.6, # diagonal ~8.49, threshold ~5.09, centroid gap ~4.24
np.array(
[
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 1],
[0, 0, 0, 1, 1, 1],
],
dtype=bool,
),
DoesNotRaise(),
),
# empty mask
(
np.zeros((4, 4), dtype=bool),
4,
"edge",
2.0,
None,
np.zeros((4, 4), dtype=bool),
DoesNotRaise(),
),
# full mask
(
np.ones((4, 4), dtype=bool),
8,
"centroid",
None,
0.2,
np.ones((4, 4), dtype=bool),
DoesNotRaise(),
),
# two components, pixel distance = 2, kept with abs=2
(
np.array(
[
[0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 1, 1, 1],
[0, 1, 1, 1, 0, 1, 1, 1],
[0, 1, 1, 1, 0, 1, 1, 1],
[0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
8,
"edge",
2.0, # was 1.0
None,
np.array(
[
[0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 1, 1, 1],
[0, 1, 1, 1, 0, 1, 1, 1],
[0, 1, 1, 1, 0, 1, 1, 1],
[0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
DoesNotRaise(),
),
# two components, pixel distance = 3, dropped with abs=2
(
np.array(
[
[0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 0, 0, 1, 1],
[0, 1, 1, 1, 0, 0, 0, 1, 1],
[0, 1, 1, 1, 0, 0, 0, 1, 1],
[0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
8,
"edge",
2.0, # keep threshold below 3 so the right blob is removed
None,
np.array(
[
[0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0],
],
dtype=bool,
),
DoesNotRaise(),
),
],
)
def test_filter_segments_by_distance_sweep(
mask: npt.NDArray,
connectivity: int,
mode: str,
absolute_distance: float | None,
relative_distance: float | None,
expected_result: npt.NDArray | None,
exception: Exception,
) -> None:
with exception:
result = filter_segments_by_distance(
mask=mask,
connectivity=connectivity,
mode=mode, # type: ignore[arg-type]
absolute_distance=absolute_distance,
relative_distance=relative_distance,
)
assert np.array_equal(result, expected_result)
Reference in New Issue View Git Blame Copy Permalink