diff --git a/src/supervision/detection/tools/polygon_zone.py b/src/supervision/detection/tools/polygon_zone.py index 336169dd..58e979f9 100644 --- a/src/supervision/detection/tools/polygon_zone.py +++ b/src/supervision/detection/tools/polygon_zone.py @@ -1,7 +1,6 @@ from __future__ import annotations from collections.abc import Iterable -from dataclasses import replace from typing import Any, cast import cv2 @@ -9,7 +8,6 @@ import numpy as np import numpy.typing as npt from supervision import Detections -from supervision.detection.utils.boxes import clip_boxes from supervision.detection.utils.converters import polygon_to_mask from supervision.draw.color import Color from supervision.draw.utils import draw_filled_polygon, draw_polygon, draw_text @@ -73,7 +71,6 @@ class PolygonZone: self.current_count = 0 x_max, y_max = np.max(polygon, axis=0) - self.frame_resolution_wh = (x_max + 1, y_max + 1) self.mask = polygon_to_mask( polygon=polygon, resolution_wh=(x_max + 2, y_max + 2) ) @@ -82,33 +79,36 @@ class PolygonZone: """ Determines if the detections are within the polygon zone. + Anchor points are calculated from original (unclipped) detection boxes to + avoid per-zone clipping shifting anchor positions. This prevents a single + detection from being counted in multiple non-overlapping zones due to + clipping artifacts, although overlapping zones may still legitimately + contain the same detection. + Args: - detections: The detections - to be checked against the polygon zone + detections: The detections to be checked against the polygon zone Returns: A boolean numpy array indicating if each detection is within the polygon zone """ + if len(detections) == 0: + self.current_count = 0 + return np.array([], dtype=bool) - clipped_xyxy = clip_boxes( - xyxy=detections.xyxy, resolution_wh=self.frame_resolution_wh - ) - clipped_detections = replace(detections, xyxy=clipped_xyxy) - all_clipped_anchors = np.array( + all_anchors = np.array( [ - np.ceil(clipped_detections.get_anchors_coordinates(anchor)).astype(int) - for anchor in self.triggering_anchors + np.ceil(detections.get_anchors_coordinates(anchors)).astype(int) + for anchors in self.triggering_anchors ] ) - is_in_zone: npt.NDArray[np.bool_] = ( - self.mask[all_clipped_anchors[:, :, 1], all_clipped_anchors[:, :, 0]] - .transpose() - .astype(bool) - ) - - is_in_zone = np.all(is_in_zone, axis=1) + mask_h, mask_w = self.mask.shape + x, y = all_anchors[:, :, 0], all_anchors[:, :, 1] + in_bounds = (x >= 0) & (y >= 0) & (x < mask_w) & (y < mask_h) + x_safe = np.clip(x, 0, mask_w - 1) + y_safe = np.clip(y, 0, mask_h - 1) + is_in_zone = np.all(in_bounds & self.mask[y_safe, x_safe], axis=0) self.current_count = int(np.sum(is_in_zone)) return is_in_zone.astype(bool) diff --git a/tests/detection/test_polygonzone.py b/tests/detection/test_polygonzone.py index 3a416d2b..6e952ec7 100644 --- a/tests/detection/test_polygonzone.py +++ b/tests/detection/test_polygonzone.py @@ -28,78 +28,139 @@ DETECTIONS = _create_detections( POLYGON = np.array([[100, 100], [200, 100], [200, 200], [100, 200]]) -@pytest.mark.parametrize( - ("detections", "polygon_zone", "expected_results", "exception"), - [ - ( - DETECTIONS, - sv.PolygonZone( +class TestPolygonZoneInit: + @pytest.mark.parametrize( + ("polygon", "triggering_anchors", "exception"), + [ + (POLYGON, [sv.Position.CENTER], DoesNotRaise()), + ( POLYGON, - triggering_anchors=( - sv.Position.TOP_LEFT, - sv.Position.TOP_RIGHT, - sv.Position.BOTTOM_LEFT, - sv.Position.BOTTOM_RIGHT, + [], + pytest.raises(ValueError, match="Triggering anchors cannot be empty"), + ), + ], + ) + def test_empty_anchors_raises(self, polygon, triggering_anchors, exception): + with exception: + sv.PolygonZone(polygon, triggering_anchors=triggering_anchors) + + +class TestPolygonZoneTrigger: + @pytest.mark.parametrize( + ("detections", "polygon_zone", "expected_results", "exception"), + [ + ( + DETECTIONS, + sv.PolygonZone( + POLYGON, + triggering_anchors=( + sv.Position.TOP_LEFT, + sv.Position.TOP_RIGHT, + sv.Position.BOTTOM_LEFT, + sv.Position.BOTTOM_RIGHT, + ), ), - ), - np.array( - [False, False, False, True, True, True, False, False, False], dtype=bool - ), - DoesNotRaise(), - ), # Test all four corners - ( - DETECTIONS, - sv.PolygonZone( - POLYGON, - ), - np.array( - [False, False, True, True, True, True, False, False, False], dtype=bool - ), - DoesNotRaise(), - ), # Test default behaviour when no anchors are provided - ( - DETECTIONS, - sv.PolygonZone( - POLYGON, - triggering_anchors=[sv.Position.BOTTOM_CENTER], - ), - np.array( - [False, False, True, True, True, True, False, False, False], dtype=bool - ), - DoesNotRaise(), - ), # Test default behaviour with deprecated api. - ( - sv.Detections.empty(), - sv.PolygonZone( - POLYGON, - ), - np.array([], dtype=bool), - DoesNotRaise(), - ), # Test empty detections - ], -) -def test_polygon_zone_trigger( - detections: sv.Detections, - polygon_zone: sv.PolygonZone, - expected_results: np.ndarray, - exception: Exception, -) -> None: - with exception: - in_zone = polygon_zone.trigger(detections) - assert np.all(in_zone == expected_results) + np.array( + [False, False, False, True, True, True, False, False, False], + dtype=bool, + ), + DoesNotRaise(), + ), # all four corners must be inside + ( + DETECTIONS, + sv.PolygonZone(POLYGON), + np.array( + [False, False, True, True, True, True, False, False, False], + dtype=bool, + ), + DoesNotRaise(), + ), # default anchor (BOTTOM_CENTER) + ( + DETECTIONS, + sv.PolygonZone( + POLYGON, + triggering_anchors=[sv.Position.BOTTOM_CENTER], + ), + np.array( + [False, False, True, True, True, True, False, False, False], + dtype=bool, + ), + DoesNotRaise(), + ), # explicit BOTTOM_CENTER matches default + ( + sv.Detections.empty(), + sv.PolygonZone(POLYGON), + np.array([], dtype=bool), + DoesNotRaise(), + ), # empty detections return empty array + ], + ) + def test_anchor_configurations( + self, + detections: sv.Detections, + polygon_zone: sv.PolygonZone, + expected_results: np.ndarray, + exception: Exception, + ) -> None: + with exception: + in_zone = polygon_zone.trigger(detections) + assert np.all(in_zone == expected_results) + assert polygon_zone.current_count == int(np.sum(expected_results)) + def test_straddling_detection_assigned_to_one_zone(self) -> None: + """Detection straddling two adjacent zones is counted in exactly one zone. -@pytest.mark.parametrize( - ("polygon", "triggering_anchors", "exception"), - [ - (POLYGON, [sv.Position.CENTER], DoesNotRaise()), - ( - POLYGON, - [], - pytest.raises(ValueError, match="Triggering anchors cannot be empty"), - ), - ], -) -def test_polygon_zone_initialization(polygon, triggering_anchors, exception): - with exception: - sv.PolygonZone(polygon, triggering_anchors=triggering_anchors) + The old implementation clipped each detection box to the zone's bounding box + before computing anchors, so a box straddling two adjacent zones got a + different anchor per zone and was double-counted. The fix computes anchors + from the original unclipped box, giving a single consistent position. + + Setup: zone_left covers x 0-99, zone_right covers x 100-200. + Detection [60, 80, 140, 120] has BOTTOM_CENTER = ceil((60+140)/2), ceil(120) + = (100, 120), which falls in zone_right only. + """ + zone_left = sv.PolygonZone( + np.array([[0, 0], [99, 0], [99, 200], [0, 200]], dtype=np.int32) + ) + zone_right = sv.PolygonZone( + np.array([[100, 0], [200, 0], [200, 200], [100, 200]], dtype=np.int32) + ) + detections = _create_detections( + xyxy=[[60.0, 80.0, 140.0, 120.0]], + class_id=[0], + ) + + results = np.array( + [zone.trigger(detections)[0] for zone in [zone_left, zone_right]], + dtype=bool, + ) + assert np.sum(results) == 1, ( + "Detection should appear in exactly one zone, not zero or multiple" + ) + assert results[1], "BOTTOM_CENTER (100, 120) should be in zone_right" + + def test_out_of_bounds_anchor_excluded(self) -> None: + """Anchors with negative coordinates are excluded, not wrapped or clamped.""" + zone = sv.PolygonZone( + np.array([[0, 0], [100, 0], [100, 100], [0, 100]]), + triggering_anchors=[sv.Position.CENTER], + ) + # CENTER = (ceil((-50+0)/2), ceil((25+75)/2)) = (-25, 50) — x < 0. + detections = _create_detections( + xyxy=[[-50.0, 25.0, 0.0, 75.0]], + class_id=[0], + ) + result = zone.trigger(detections) + assert not result[0] + + def test_anchor_on_polygon_boundary_included(self) -> None: + """An anchor landing exactly on a polygon corner is counted as inside.""" + polygon = np.array([[0, 0], [100, 0], [100, 100], [0, 100]]) + zone = sv.PolygonZone(polygon, triggering_anchors=[sv.Position.BOTTOM_RIGHT]) + # BOTTOM_RIGHT = (ceil(x2), ceil(y2)) = (100, 100) — the polygon corner. + detections = _create_detections( + xyxy=[[50.0, 50.0, 100.0, 100.0]], + class_id=[0], + ) + result = zone.trigger(detections) + assert result[0]