From ce7e3a1de4eba124e67aec68a983cdc3ea748303 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 31 Oct 2023 18:04:13 +0100 Subject: [PATCH] Refactor clip_boxes function and enhance detections processing Refactored the `clip_boxes` function for clarity by renaming arguments from `boxes_xyxy` and `frame_resolution_wh` to `xyxy` and `resolution_wh`, respectively. These change makes the function arguments more intuitive and improves code readability. The processing of detections in `supervision/annotators/core.py` has been updated to include clipping of detection boxes to the image bounds before processing. This prevents errors and ensures detections beyond the image dimensions are handled correctly. Adjustments were also made in the test cases and in `polygon_zone.py` to match the updated `clip_boxes` function. --- supervision/annotators/core.py | 10 +++++++--- supervision/detection/tools/polygon_zone.py | 2 +- supervision/detection/utils.py | 12 +++++------- test/detection/test_utils.py | 8 ++++---- 4 files changed, 17 insertions(+), 15 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index b6378bc0..8da3a659 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -7,6 +7,7 @@ import numpy as np from supervision.annotators.base import BaseAnnotator from supervision.annotators.utils import ColorLookup, Trace, resolve_color from supervision.detection.core import Detections +from supervision.detection.utils import clip_boxes from supervision.draw.color import Color, ColorPalette from supervision.geometry.core import Position @@ -891,10 +892,13 @@ class BlurAnnotator(BaseAnnotator): ![blur-annotator-example](https://media.roboflow.com/ supervision-annotator-examples/blur-annotator-example-purple.png) """ - for detection_idx in range(len(detections)): - x1, y1, x2, y2 = np.maximum(detections.xyxy[detection_idx].astype(int), 0) - roi = scene[y1:y2, x1:x2] + image_height, image_width = scene.shape[:2] + clipped_xyxy = clip_boxes( + xyxy=detections.xyxy, + resolution_wh=(image_width, image_height)).astype(int) + for x1, y1, x2, y2 in clipped_xyxy: + roi = scene[y1:y2, x1:x2] roi = cv2.blur(roi, (self.kernel_size, self.kernel_size)) scene[y1:y2, x1:x2] = roi diff --git a/supervision/detection/tools/polygon_zone.py b/supervision/detection/tools/polygon_zone.py index f1dba839..15ee1aa2 100644 --- a/supervision/detection/tools/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -56,7 +56,7 @@ class PolygonZone: """ clipped_xyxy = clip_boxes( - boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh + xyxy=detections.xyxy, resolution_wh=self.frame_resolution_wh ) clipped_detections = replace(detections, xyxy=clipped_xyxy) clipped_anchors = np.ceil( diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 7a5eb546..7aa2c846 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -110,17 +110,15 @@ def non_max_suppression( return keep[sort_index.argsort()] -def clip_boxes( - boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int] -) -> np.ndarray: +def clip_boxes(xyxy: np.ndarray, 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 + 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)` + resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the frame. Returns: @@ -128,8 +126,8 @@ def clip_boxes( 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 = np.copy(xyxy) + width, height = 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 09691adf..6b3249b4 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -122,7 +122,7 @@ def test_non_max_suppression( @pytest.mark.parametrize( - "boxes_xyxy, frame_resolution_wh, expected_result", + "xyxy, resolution_wh, expected_result", [ ( np.empty(shape=(0, 4)), @@ -157,11 +157,11 @@ def test_non_max_suppression( ], ) def test_clip_boxes( - boxes_xyxy: np.ndarray, - frame_resolution_wh: Tuple[int, int], + xyxy: np.ndarray, + resolution_wh: Tuple[int, int], expected_result: np.ndarray, ) -> None: - result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh) + result = clip_boxes(xyxy=xyxy, resolution_wh=resolution_wh) assert np.array_equal(result, expected_result)