From 63fd5a3f8acac483764d55db90f35a968eb6f081 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 14 Mar 2023 11:33:31 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=AA=20ready=20for=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/polygon_zone.py | 20 ++++---- supervision/detection/utils.py | 21 ++++++++ test/detection/test_utils.py | 70 +++++++++++++++++++++++++-- 3 files changed, 97 insertions(+), 14 deletions(-) diff --git a/supervision/detection/polygon_zone.py b/supervision/detection/polygon_zone.py index 54bbbf62..2bb321bb 100644 --- a/supervision/detection/polygon_zone.py +++ b/supervision/detection/polygon_zone.py @@ -1,10 +1,11 @@ from typing import Optional, Tuple +from dataclasses import replace import cv2 import numpy as np from supervision import Detections -from supervision.detection.utils import generate_2d_mask +from supervision.detection.utils import generate_2d_mask, clip_boxes from supervision.draw.color import Color from supervision.draw.utils import draw_polygon, draw_text from supervision.geometry.core import Position @@ -21,17 +22,18 @@ class PolygonZone: 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 - ) - is_in_zone = self.mask[anchors[:, 1], anchors[:, 0]] + clipped_xyxy = clip_boxes(boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh) + 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) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 5d8407f4..f3e1ee08 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -92,3 +92,24 @@ 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 diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 110b682e..9ba3354a 100644 --- a/test/detection/test_utils.py +++ b/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 @@ from supervision.detection.utils import non_max_suppression ]), 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)