feat: speed up box iou batch using 2d in-place ops instead of 3d

without creating intermediary (N, M, 2) arrays.
This commit is contained in:
AnonymDevOSS 2025-10-28 13:51:28 +01:00
parent 171687f3ab
commit 35ff87b08d
3 changed files with 142 additions and 0 deletions

View File

@ -231,6 +231,96 @@ def box_iou_batch(
return ious
def box_iou_batch_alt(
boxes_true: np.ndarray,
boxes_detection: np.ndarray,
overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> np.ndarray:
"""
Compute Intersection over Union (IoU) of two sets of bounding boxes -
`boxes_true` and `boxes_detection`. Both sets
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
Note:
Use `box_iou` when computing IoU between two individual boxes.
For comparing multiple boxes (arrays of boxes), use `box_iou_batch` for better
performance.
Args:
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes.
`shape = (N, 4)` where `N` is number of true objects.
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes.
`shape = (M, 4)` where `M` is number of detected objects.
overlap_metric (OverlapMetric): Metric used to compute the degree of overlap
between pairs of boxes (e.g., IoU, IoS).
Returns:
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
`shape = (N, M)` where `N` is number of true objects and
`M` is number of detected objects.
Examples:
```python
import numpy as np
import supervision as sv
boxes_true = np.array([
[100, 100, 200, 200],
[300, 300, 400, 400]
])
boxes_detection = np.array([
[150, 150, 250, 250],
[320, 320, 420, 420]
])
sv.box_iou_batch(boxes_true=boxes_true, boxes_detection=boxes_detection)
# array([
# [0.14285714, 0. ],
# [0. , 0.47058824]
# ])
```
"""
tx1, ty1, tx2, ty2 = boxes_true.T
dx1, dy1, dx2, dy2 = boxes_detection.T
N, M = boxes_true.shape[0], boxes_detection.shape[0]
top_left_x = np.empty((N, M), dtype=np.float32)
bottom_right_x = np.empty_like(top_left_x)
top_left_y = np.empty_like(top_left_x)
bottom_right_y = np.empty_like(top_left_x)
np.maximum(tx1[:, None], dx1[None, :], out=top_left_x)
np.minimum(tx2[:, None], dx2[None, :], out=bottom_right_x)
np.maximum(ty1[:, None], dy1[None, :], out=top_left_y)
np.minimum(ty2[:, None], dy2[None, :], out=bottom_right_y)
np.subtract(bottom_right_x, top_left_x, out=bottom_right_x) # W
np.subtract(bottom_right_y, top_left_y, out=bottom_right_y) # H
np.clip(bottom_right_x, 0.0, None, out=bottom_right_x)
np.clip(bottom_right_y, 0.0, None, out=bottom_right_y)
area_inter = bottom_right_x * bottom_right_y
area_true = (tx2 - tx1) * (ty2 - ty1)
area_detection = (dx2 - dx1) * (dy2 - dy1)
if overlap_metric == OverlapMetric.IOU:
denom = area_true[:, None] + area_detection[None, :] - area_inter
elif overlap_metric == OverlapMetric.IOS:
denom = np.minimum(area_true[:, None], area_detection[None, :])
else:
raise ValueError(
f"overlap_metric {overlap_metric} is not supported, "
"only 'IOU' and 'IOS' are supported"
)
out = np.zeros_like(area_inter, dtype=np.float32)
np.divide(area_inter, denom, out=out, where=denom > 0)
return out
def _jaccard(box_a: list[float], box_b: list[float], is_crowd: bool) -> float:
"""
Calculate the Jaccard index (intersection over union) between two bounding boxes.

View File

@ -0,0 +1,38 @@
import random
import numpy as np
def generate_boxes(
n: int,
W: int = 1920,
H: int = 1080,
min_size: int = 20,
max_size: int = 200,
seed: int | None = 1,
):
"""
Generate N valid bounding boxes of format [x_min, y_min, x_max, y_max].
Args:
n (int): Number of boexs to generate
W (int): Image width
H (int): Image height
min_size (int): Minimum box size (width/height)
max_size (int): Maximum box size (width/height)
seed (int | None): Random seed for reproducibility
Returns:
list[list[float]] | np.ndarray: List of boxes
"""
random.seed(seed)
boxes = []
for _ in range(n):
w = random.uniform(min_size, max_size)
h = random.uniform(min_size, max_size)
x1 = random.uniform(0, W - w)
y1 = random.uniform(0, H - h)
x2 = x1 + w
y2 = y1 + h
boxes.append([x1, y1, x2, y2])
return np.array(boxes, dtype=np.float32)

View File

@ -7,10 +7,13 @@ import pytest
from supervision.detection.utils.iou_and_nms import (
_group_overlapping_boxes,
box_iou_batch,
box_iou_batch_alt,
box_non_max_suppression,
mask_non_max_merge,
mask_non_max_suppression,
)
from test.detection.utils.functions import generate_boxes
@pytest.mark.parametrize(
@ -631,3 +634,14 @@ def test_mask_non_max_merge(
sorted_result = sorted([sorted(group) for group in result])
sorted_expected_result = sorted([sorted(group) for group in expected_result])
assert sorted_result == sorted_expected_result
def test_box_iou_batch_and_alt_equivalence():
boxes_true = generate_boxes(20, seed=1)
boxes_detection = generate_boxes(30, seed=2)
iou_a = box_iou_batch(boxes_true, boxes_detection)
iou_b = box_iou_batch_alt(boxes_true, boxes_detection)
assert iou_a.shape == iou_b.shape
assert np.allclose(iou_a, iou_b, rtol=1e-6, atol=1e-6)