95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
from contextlib import ExitStack as DoesNotRaise
|
|
from test.utils import mock_detections
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
|
|
from supervision.annotators.utils import ColorMap, resolve_color_idx
|
|
from supervision.detection.core import Detections
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"detections, detection_idx, color_map, expected_result, exception",
|
|
[
|
|
(
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
|
class_id=[5, 3],
|
|
tracker_id=[2, 6],
|
|
),
|
|
0,
|
|
ColorMap.INDEX,
|
|
0,
|
|
DoesNotRaise(),
|
|
), # multiple detections; index mapping
|
|
(
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
|
class_id=[5, 3],
|
|
tracker_id=[2, 6],
|
|
),
|
|
0,
|
|
ColorMap.CLASS,
|
|
5,
|
|
DoesNotRaise(),
|
|
), # multiple detections; class mapping
|
|
(
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
|
class_id=[5, 3],
|
|
tracker_id=[2, 6],
|
|
),
|
|
0,
|
|
ColorMap.TRACK,
|
|
2,
|
|
DoesNotRaise(),
|
|
), # multiple detections; track mapping
|
|
(
|
|
Detections.empty(),
|
|
0,
|
|
ColorMap.INDEX,
|
|
None,
|
|
pytest.raises(ValueError),
|
|
), # no detections; index mapping; out of bounds
|
|
(
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
|
class_id=[5, 3],
|
|
tracker_id=[2, 6],
|
|
),
|
|
2,
|
|
ColorMap.INDEX,
|
|
None,
|
|
pytest.raises(ValueError),
|
|
), # multiple detections; index mapping; out of bounds
|
|
(
|
|
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
|
0,
|
|
ColorMap.CLASS,
|
|
None,
|
|
pytest.raises(ValueError),
|
|
), # multiple detections; class mapping; no class_id
|
|
(
|
|
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
|
0,
|
|
ColorMap.TRACK,
|
|
None,
|
|
pytest.raises(ValueError),
|
|
), # multiple detections; class mapping; no track_id
|
|
],
|
|
)
|
|
def test_resolve_color_idx(
|
|
detections: Detections,
|
|
detection_idx: int,
|
|
color_map: ColorMap,
|
|
expected_result: Optional[int],
|
|
exception: Exception,
|
|
) -> None:
|
|
with exception:
|
|
result = resolve_color_idx(
|
|
detections=detections,
|
|
detection_idx=detection_idx,
|
|
color_map=color_map,
|
|
)
|
|
assert result == expected_result
|