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.
This commit is contained in:
SkalskiP 2023-10-31 18:04:13 +01:00
parent 8b88f1f415
commit ce7e3a1de4
4 changed files with 17 additions and 15 deletions

View File

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

View File

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

View File

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

View File

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