Merge pull request #46 from roboflow/fix/041_index_error_on_polygon_z…

…ones

fix/041_index_error_on_polygon_zones

supervision/detection/polygon_zone.py

@@ -1,10 +1,11 @@
+from dataclasses import replace
 from typing import Optional, Tuple
 
 import cv2
 import numpy as np
 
 from supervision import Detections
-from supervision.detection.utils import generate_2d_mask
+from supervision.detection.utils import clip_boxes, generate_2d_mask
 from supervision.draw.color import Color
 from supervision.draw.utils import draw_polygon, draw_text
 from supervision.geometry.core import Position
@@ -21,17 +22,22 @@ def __init__(
         self.polygon = polygon
         self.frame_resolution_wh = frame_resolution_wh
         self.triggering_position = triggering_position
-        self.mask = generate_2d_mask(polygon=polygon, resolution_wh=frame_resolution_wh)
         self.current_count = 0
 
+        width, height = frame_resolution_wh
+        self.mask = generate_2d_mask(
+            polygon=polygon, resolution_wh=(width + 1, height + 1)
+        )
+
     def trigger(self, detections: Detections) -> np.ndarray:
-        anchors = (
-            np.ceil(
-                detections.get_anchor_coordinates(anchor=self.triggering_position)
-            ).astype(int)
-            - 1
+        clipped_xyxy = clip_boxes(
+            boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh
         )
-        is_in_zone = self.mask[anchors[:, 1], anchors[:, 0]]
+        clipped_detections = replace(detections, xyxy=clipped_xyxy)
+        clipped_anchors = np.ceil(
+            clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
+        ).astype(int)
+        is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
         self.current_count = np.sum(is_in_zone)
         return is_in_zone.astype(bool)
 

supervision/detection/utils.py

@@ -92,3 +92,26 @@ def non_max_suppression(
         keep = keep & ~condition
 
     return keep[sort_index.argsort()]
+
+
+def clip_boxes(
+    boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+    """
+    Clips bounding boxes coordinates to fit within the frame resolution.
+
+    Args:
+        boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each row corresponds to a bounding box in
+        the format `(x_min, y_min, x_max, y_max)`.
+        frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the
+        frame.
+
+    Returns:
+        np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a bounding box with coordinates
+        clipped to fit within the frame resolution.
+    """
+    result = np.copy(boxes_xyxy)
+    width, height = frame_resolution_wh
+    result[:, [0, 2]] = result[:, [0, 2]].clip(0, width)
+    result[:, [1, 3]] = result[:, [1, 3]].clip(0, height)
+    return result

test/detection/test_utils.py

@@ -1,11 +1,11 @@
 from contextlib import ExitStack as DoesNotRaise
-from typing import Optional
+from typing import Optional, Tuple
 
 import pytest
 
 import numpy as np
 
-from supervision.detection.utils import non_max_suppression
+from supervision.detection.utils import non_max_suppression, clip_boxes
 
 
 @pytest.mark.parametrize(
@@ -61,14 +61,14 @@
             ]),
             DoesNotRaise()
         ),  # two boxes with different category
-(
+        (
             np.array([
                 [10.0, 10.0, 40.0, 40.0, 0.8, 0],
                 [15.0, 15.0, 40.0, 40.0, 0.9, 0],
             ]),
             0.5,
             np.array([
-                True,
+                False,
                 True
             ]),
             DoesNotRaise()
@@ -125,4 +125,64 @@ def test_non_max_suppression(
 ) -> None:
     with exception:
         result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold)
-        np.array_equal(result, expected_result)
+        assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+    "boxes_xyxy, frame_resolution_wh, expected_result",
+    [
+        (
+            np.empty(shape=(0, 4)),
+            (1280, 720),
+            np.empty(shape=(0, 4)),
+        ),
+        (
+            np.array([
+                [1.0, 1.0, 1279.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 1.0, 1279.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [-1.0, 1.0, 1279.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [0.0, 1.0, 1279.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [1.0, -1.0, 1279.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 0.0, 1279.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [1.0, 1.0, 1281.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 1.0, 1280.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [1.0, 1.0, 1279.0, 721.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 1.0, 1279.0, 720.0]
+            ]),
+        ),
+    ]
+)
+def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
+    result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
+    assert np.array_equal(result, expected_result)