Add LineZone unit tests

This commit is contained in:
tc360950 2024-05-17 16:17:31 +02:00
parent 34a5b34c88
commit da904ad2da
1 changed files with 323 additions and 3 deletions

View File

@ -1,10 +1,13 @@
from contextlib import ExitStack as DoesNotRaise
from typing import Optional, Tuple
from itertools import chain, combinations
from test.test_utils import mock_detections
from typing import Optional, Tuple, List
import numpy as np
import pytest
from supervision import LineZone
from supervision.geometry.core import Point, Vector
from supervision import Detections, LineZone
from supervision.geometry.core import Point, Position, Vector
@pytest.mark.parametrize(
@ -70,3 +73,320 @@ def test_calculate_region_of_interest_limits(
with exception:
result = LineZone.calculate_region_of_interest_limits(vector=vector)
assert result == expected_result
@pytest.mark.parametrize(
"vector, bbox_sequence, expected_count_in, expected_count_out",
[
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[100, 50, 120, 70],
[-100, 50, -80, 70],
],
[False, False],
[False, True],
),
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[-100, 50, -80, 70],
[100, 50, 120, 70],
],
[False, True],
[False, False],
),
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[-100, 50, -80, 70],
[-10, 50, 20, 70],
[100, 50, 120, 70],
],
[False, False, True],
[False, False, False],
),
(
Vector(
Point(0, 0),
Point(100, 100),
),
[
[50, 45, 70, 30],
[40, 50, 50, 40],
[0, 50, 10, 40],
],
[False, False, False],
[False, False, True],
),
(
Vector(
Point(0, 0),
Point(100, 0),
),
[
[50, -45, 70, -30],
[40, 50, 50, 40],
],
[False, False],
[False, True],
),
(
Vector(
Point(0, 0),
Point(0, -100),
),
[
[100, -50, 120, -70],
[-100, -50, -80, -70],
],
[False, True],
[False, False],
),
(
Vector(
Point(0, 0),
Point(50, 100),
),
[
[50, 50, 70, 30],
[40, 50, 50, 40],
[0, 50, 10, 40],
],
[False, False, False],
[False, False, True],
),
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[100, 50, 120, 70],
[-100, 50, -80, 70],
[100, 50, 120, 70],
[-100, 50, -80, 70],
[100, 50, 120, 70],
[-100, 50, -80, 70],
[100, 50, 120, 70],
[-100, 50, -80, 70],
],
[False, False, True, False, True, False, True, False],
[False, True, False, True, False, True, False, True],
),
(
Vector(
Point(0, 0),
Point(-100, 0),
),
[
[-50, 70, -40, 50],
[-50, -70, -40, -50],
[-50, 70, -40, 50],
[-50, -70, -40, -50],
[-50, 70, -40, 50],
[-50, -70, -40, -50],
[-50, 70, -40, 50],
[-50, -70, -40, -50],
],
[False, False, True, False, True, False, True, False],
[False, True, False, True, False, True, False, True],
),
],
)
def test_line_zone_single_detection(
vector, bbox_sequence, expected_count_in: List[bool], expected_count_out: List[bool]
) -> None:
line_zone = LineZone(start=vector.start, end=vector.end)
for i, bbox in enumerate(bbox_sequence):
detections = mock_detections(
xyxy=[bbox],
tracker_id=[i for i in range(0, 1)],
)
count_in, count_out = line_zone.trigger(detections)
assert count_in[0] == expected_count_in[i]
assert count_out[0] == expected_count_out[i]
assert line_zone.in_count == sum(expected_count_in[: (i + 1)])
assert line_zone.out_count == sum(expected_count_out[: (i + 1)])
@pytest.mark.parametrize(
"vector, bbox_sequence, expected_count_in, expected_count_out, crossing_anchors",
[
(
Vector(
Point(0, 0),
Point(100, 100),
),
[
[50, 30, 60, 20],
[20, 50, 40, 30],
],
[False, False],
[False, True],
[Position.TOP_LEFT, Position.TOP_RIGHT, Position.BOTTOM_LEFT],
),
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[-100, 50, -80, 70],
[-100, 50, 120, 70],
],
[False, True],
[False, False],
[Position.TOP_RIGHT, Position.BOTTOM_RIGHT],
),
],
)
def test_line_zone_single_detection_on_subset_of_anchors(
vector,
bbox_sequence,
expected_count_in: List[bool],
expected_count_out: List[bool],
crossing_anchors,
) -> None:
def powerset(s):
return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
for anchors in powerset(
[
Position.TOP_LEFT,
Position.TOP_RIGHT,
Position.BOTTOM_LEFT,
Position.BOTTOM_RIGHT,
]
):
if not anchors:
continue
line_zone = LineZone(
start=vector.start, end=vector.end, triggering_anchors=anchors
)
for i, bbox in enumerate(bbox_sequence):
detections = mock_detections(
xyxy=[bbox],
tracker_id=[i for i in range(0, 1)],
)
count_in, count_out = line_zone.trigger(detections)
if all(anchor in crossing_anchors for anchor in anchors):
assert count_in == expected_count_in[i]
assert count_out == expected_count_out[i]
else:
assert np.all(not count_in)
assert np.all(not count_out)
@pytest.mark.parametrize(
"vector, bbox_sequence, expected_count_in, expected_count_out",
[
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[[100, 50, 120, 70], [100, 50, 120, 70]],
[[-100, 50, -80, 70], [100, 50, 120, 70]],
[[100, 50, 120, 70], [100, 50, 120, 70]],
],
[[False, False], [False, False], [True, False]],
[[False, False], [True, False], [False, False]],
),
(
Vector(
Point(0, 0),
Point(-100, 0),
),
[
[[-50, 70, -40, 50], [-80, -50, -70, -40]],
[[-50, -70, -40, -50], [-80, 50, -70, 40]],
[[-50, 70, -40, 50], [-80, 50, -70, 40]],
[[-50, -70, -40, -50], [-80, 50, -70, 40]],
[[-50, 70, -40, 50], [-80, 50, -70, 40]],
[[-50, -70, -40, -50], [-80, 50, -70, 40]],
[[-50, 70, -40, 50], [-80, 50, -70, 40]],
[[-50, -70, -40, -50], [-80, -50, -70, -40]],
],
[
(False, False),
(False, True),
(True, False),
(False, False),
(True, False),
(False, False),
(True, False),
(False, False),
],
[
(False, False),
(True, False),
(False, False),
(True, False),
(False, False),
(True, False),
(False, False),
(True, True),
],
),
],
)
def test_line_zone_multiple_detections(
vector, bbox_sequence, expected_count_in: List[bool], expected_count_out: List[bool]
) -> None:
line_zone = LineZone(start=vector.start, end=vector.end)
for i, bboxes in enumerate(bbox_sequence):
detections = mock_detections(
xyxy=bboxes,
tracker_id=[i for i in range(0, len(bboxes))],
)
count_in, count_out = line_zone.trigger(detections)
assert np.all(count_in == expected_count_in[i])
assert np.all(count_out == expected_count_out[i])
@pytest.mark.parametrize(
"vector, bbox_sequence",
[
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[100, 50, 120, 70],
[-100, 50, -80, 70],
],
),
(
Vector(
Point(0, 0),
Point(0, 100),
),
[
[-100, 50, -80, 70],
[100, 50, 120, 70],
],
),
],
)
def test_line_zone_does_not_count_detections_without_tracker_id(vector, bbox_sequence):
line_zone = LineZone(start=vector.start, end=vector.end)
for bbox in bbox_sequence:
detections = Detections(
xyxy=np.array([bbox]).reshape((-1, 4)),
tracker_id=np.array([None for _ in range(0, 1)]),
)
count_in, count_out = line_zone.trigger(detections)
assert np.all(not count_in)
assert np.all(not count_out)