From 38f52a04ae40339533f50de819633132907ed286 Mon Sep 17 00:00:00 2001
From: SkalskiP <piotr.skalski92@gmail.com>
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/utils.py | 21 ++++++++++
 test/detection/test_utils.py   | 70 +++++++++++++++++++++++++++++++---
 2 files changed, 86 insertions(+), 5 deletions(-)

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)