diff --git a/test/detection/test_line_counter.py b/test/detection/test_line_counter.py index 73780414..997da352 100644 --- a/test/detection/test_line_counter.py +++ b/test/detection/test_line_counter.py @@ -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)