From 7cbf16be8a3841af2c090c6edcddabc640ad1f94 Mon Sep 17 00:00:00 2001
From: SkalskiP <piotr.skalski92@gmail.com>
Date: Wed, 6 Dec 2023 12:16:22 +0100
Subject: [PATCH] Add `scale_boxes` function to detection utils

Added a new function scale_boxes to the detection utilities in order to scale the dimensions of bounding boxes. This function accepts bounding boxes and a scaling factor as input and returns the scaled bounding boxes. It can enlarge or shrink the input bounding boxes based on the value of the passed factor. The function was also documented in the utils.md file and necessary tests were added to assure its correctness.
---
 docs/detection/utils.md        |  9 +++++++
 supervision/detection/utils.py | 21 ++++++++++++++-
 test/detection/test_utils.py   | 49 +++++++++++++++++++++++++++++++++-
 3 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/docs/detection/utils.md b/docs/detection/utils.md
index 9a23fcb4..7b221231 100644
--- a/docs/detection/utils.md
+++ b/docs/detection/utils.md
@@ -25,3 +25,12 @@
 ## filter_polygons_by_area
 
 :::supervision.detection.utils.filter_polygons_by_area
+
+## move_boxes
+
+:::supervision.detection.utils.move_boxes
+
+## scale_boxes
+
+:::supervision.detection.utils.scale_boxes
+
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index 1e5fdc35..33f49131 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -383,7 +383,7 @@ def process_roboflow_result(
 
 def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
     """
-    Args:
+    Parameters:
         xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes
             coordinates in format `[x1, y1, x2, y2]`
         offset (np.array): An array of shape `(2,)` containing offset values in format
@@ -395,6 +395,25 @@ def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
     return xyxy + np.hstack([offset, offset])
 
 
+def scale_boxes(xyxy: np.ndarray, factor: float) -> np.ndarray:
+    """
+    Scale the dimensions of bounding boxes.
+
+    Parameters:
+        xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes
+            coordinates in format `[x1, y1, x2, y2]`
+        factor (float): A float value representing the factor by which the box
+            dimensions are scaled. A factor greater than 1 enlarges the boxes, while a
+            factor less than 1 shrinks them.
+
+    Returns:
+        (np.ndarray) scaled bounding boxes
+    """
+    centers = (xyxy[:, :2] + xyxy[:, 2:]) / 2
+    new_sizes = (xyxy[:, 2:] - xyxy[:, :2]) * factor
+    return np.concatenate((centers - new_sizes / 2, centers + new_sizes / 2), axis=1)
+
+
 def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray:
     """
     Calculate the centroids of binary masks in a tensor.
diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py
index dc8e9e16..66c84ae6 100644
--- a/test/detection/test_utils.py
+++ b/test/detection/test_utils.py
@@ -10,7 +10,7 @@ from supervision.detection.utils import (
     filter_polygons_by_area,
     move_boxes,
     non_max_suppression,
-    process_roboflow_result,
+    process_roboflow_result, scale_boxes,
 )
 
 TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool)
@@ -504,6 +504,53 @@ def test_move_boxes(
         assert np.array_equal(result, expected_result)
 
 
+@pytest.mark.parametrize(
+    "xyxy, factor, expected_result, exception",
+    [
+        (
+            np.empty(shape=(0, 4)),
+            2.0,
+            np.empty(shape=(0, 4)),
+            DoesNotRaise(),
+        ),  # empty xyxy array
+        (
+            np.array([[0, 0, 10, 10]]),
+            1.0,
+            np.array([[0, 0, 10, 10]]),
+            DoesNotRaise(),
+        ),  # single box with factor equal to 1.0
+        (
+            np.array([[0, 0, 10, 10]]),
+            2.0,
+            np.array([[-5, -5, 15, 15]]),
+            DoesNotRaise(),
+        ),  # single box with factor equal to 2.0
+        (
+            np.array([[0, 0, 10, 10]]),
+            0.5,
+            np.array([[2.5, 2.5, 7.5, 7.5]]),
+            DoesNotRaise(),
+        ),  # single box with factor equal to 0.5
+        (
+            np.array([[0, 0, 10, 10], [10, 10, 30, 30]]),
+            2.0,
+            np.array([[-5, -5, 15, 15], [0, 0, 40, 40]]),
+            DoesNotRaise(),
+        ),  # two boxes with factor equal to 2.0
+    ]
+)
+def test_scale_boxes(
+    xyxy: np.ndarray,
+    factor: float,
+    expected_result: np.ndarray,
+    exception: Exception,
+) -> None:
+    with exception:
+        result = scale_boxes(xyxy=xyxy, factor=factor)
+        assert np.array_equal(result, expected_result)
+
+
+
 @pytest.mark.parametrize(
     "masks, expected_result, exception",
     [