🧪 ready for tests

This commit is contained in:
SkalskiP 2023-03-14 11:33:31 +01:00
parent 65090f612c
commit 63fd5a3f8a
3 changed files with 97 additions and 14 deletions

View File

@ -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)

View File

@ -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

View File

@ -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)