362 lines
14 KiB
Python
362 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import pytest
|
|
|
|
from supervision.config import ORIENTED_BOX_COORDINATES
|
|
from supervision.detection._geometry_dispatch import detection_area, detection_iou
|
|
from supervision.detection.compact_mask import CompactMask
|
|
from supervision.detection.core import Detections
|
|
from supervision.detection.utils.boxes import xyxyxyxy_to_xyxy
|
|
from supervision.detection.utils.iou_and_nms import (
|
|
OverlapMetric,
|
|
box_iou_batch,
|
|
mask_iou_batch,
|
|
oriented_box_iou_batch,
|
|
)
|
|
from supervision.detection.utils.masks import count_mask_pixels
|
|
|
|
|
|
def _rotated_rect(
|
|
center_x: float, center_y: float, width: float, height: float, angle_deg: float
|
|
) -> npt.NDArray[np.float32]:
|
|
"""Return four corners of a rotated rectangle in clockwise order."""
|
|
half_w = width / 2.0
|
|
half_h = height / 2.0
|
|
corners = np.array(
|
|
[
|
|
[-half_w, -half_h],
|
|
[half_w, -half_h],
|
|
[half_w, half_h],
|
|
[-half_w, half_h],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
angle = np.deg2rad(angle_deg)
|
|
rotation = np.array(
|
|
[[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]],
|
|
dtype=np.float32,
|
|
)
|
|
center = np.array([center_x, center_y], dtype=np.float32)
|
|
return (corners @ rotation.T + center).astype(np.float32)
|
|
|
|
|
|
def _detections_from_quads(
|
|
quads: list[npt.NDArray[np.float32]],
|
|
xyxy: npt.NDArray[np.float32] | None = None,
|
|
) -> Detections:
|
|
"""Build Detections carrying OBB coordinates."""
|
|
corners = np.stack(quads).astype(np.float32)
|
|
if xyxy is None:
|
|
xyxy = xyxyxyxy_to_xyxy(corners).astype(np.float32)
|
|
return Detections(
|
|
xyxy=xyxy,
|
|
data={ORIENTED_BOX_COORDINATES: corners},
|
|
)
|
|
|
|
|
|
def _full_image_xyxy(
|
|
count: int, image_shape: tuple[int, int]
|
|
) -> npt.NDArray[np.float32]:
|
|
"""Return full-image xyxy boxes for CompactMask construction."""
|
|
image_height, image_width = image_shape
|
|
xyxy = np.array([0, 0, image_width - 1, image_height - 1], dtype=np.float32)
|
|
return np.tile(xyxy, (count, 1))
|
|
|
|
|
|
class TestDetectionArea:
|
|
"""Tests for geometry-aware area dispatch."""
|
|
|
|
def test_returns_mask_pixel_area(self) -> None:
|
|
"""Masks take precedence over OBBs and AABB envelopes."""
|
|
mask = np.zeros((1, 20, 20), dtype=bool)
|
|
mask[0, 2:5, 3:8] = True
|
|
quad = _rotated_rect(10, 10, 12, 6, 30)
|
|
detections = Detections(
|
|
xyxy=np.array([[0, 0, 20, 20]], dtype=np.float32),
|
|
mask=mask,
|
|
data={ORIENTED_BOX_COORDINATES: quad[np.newaxis]},
|
|
)
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_array_equal(area, np.array([15], dtype=np.int64))
|
|
|
|
def test_dense_mask_branch_reuses_count_mask_pixels(self) -> None:
|
|
"""Dense-mask area delegates to count_mask_pixels for the shared fast path."""
|
|
mask = np.zeros((3, 10, 10), dtype=bool)
|
|
mask[0, :2, :2] = True
|
|
mask[1, :3, :3] = True
|
|
detections = Detections(
|
|
xyxy=np.tile(np.array([[0, 0, 9, 9]], dtype=np.float32), (3, 1)),
|
|
mask=mask,
|
|
)
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_array_equal(area, count_mask_pixels(mask))
|
|
assert area.dtype == np.int64
|
|
|
|
def test_returns_compact_mask_area(self) -> None:
|
|
"""CompactMask inputs use CompactMask.area without dense materialisation."""
|
|
masks = np.zeros((2, 12, 12), dtype=bool)
|
|
masks[0, 1:4, 2:7] = True
|
|
masks[1, 4:9, 4:10] = True
|
|
compact_mask = CompactMask.from_dense(
|
|
masks=masks,
|
|
xyxy=_full_image_xyxy(len(masks), masks.shape[1:]),
|
|
image_shape=masks.shape[1:],
|
|
)
|
|
detections = Detections(
|
|
xyxy=_full_image_xyxy(len(masks), masks.shape[1:]),
|
|
mask=compact_mask,
|
|
)
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_array_equal(area, compact_mask.area)
|
|
|
|
def test_returns_oriented_box_area_when_present(self) -> None:
|
|
"""OBB area is used instead of the larger rotated AABB envelope."""
|
|
quad = _rotated_rect(50, 50, 20, 10, 45)
|
|
detections = _detections_from_quads([quad])
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_allclose(area, np.array([200.0]))
|
|
assert detections.box_area[0] > area[0]
|
|
|
|
def test_returns_box_area_when_no_richer_geometry_is_present(self) -> None:
|
|
"""AABB area is the fallback when masks and OBB corners are absent."""
|
|
detections = Detections(
|
|
xyxy=np.array([[0, 0, 20, 10], [3, 4, 8, 12]], dtype=np.float32)
|
|
)
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_array_equal(area, detections.box_area)
|
|
|
|
@pytest.mark.parametrize(
|
|
"detections",
|
|
[
|
|
pytest.param(
|
|
Detections(xyxy=np.empty((0, 4), dtype=np.float32)),
|
|
id="empty-aabb",
|
|
),
|
|
pytest.param(
|
|
Detections(
|
|
xyxy=np.empty((0, 4), dtype=np.float32),
|
|
mask=np.empty((0, 8, 8), dtype=bool),
|
|
),
|
|
id="empty-mask",
|
|
),
|
|
pytest.param(
|
|
Detections(
|
|
xyxy=np.empty((0, 4), dtype=np.float32),
|
|
data={
|
|
ORIENTED_BOX_COORDINATES: np.empty((0, 4, 2), dtype=np.float32)
|
|
},
|
|
),
|
|
id="empty-obb",
|
|
),
|
|
pytest.param(
|
|
Detections(
|
|
xyxy=np.empty((0, 4), dtype=np.float32),
|
|
mask=CompactMask.from_dense(
|
|
masks=np.empty((0, 8, 8), dtype=bool),
|
|
xyxy=np.empty((0, 4), dtype=np.float32),
|
|
image_shape=(8, 8),
|
|
),
|
|
),
|
|
id="empty-compact-mask",
|
|
),
|
|
],
|
|
)
|
|
def test_returns_empty_array_for_empty_detections(
|
|
self, detections: Detections
|
|
) -> None:
|
|
"""Empty Detections return an empty area array for every geometry branch."""
|
|
area = detection_area(detections)
|
|
|
|
assert area.shape == (0,)
|
|
|
|
@pytest.mark.parametrize(
|
|
"detections",
|
|
[
|
|
pytest.param(
|
|
Detections(
|
|
xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
|
|
mask=np.ones((1, 3, 4), dtype=bool),
|
|
),
|
|
id="mask",
|
|
),
|
|
pytest.param(
|
|
_detections_from_quads([_rotated_rect(50, 50, 20, 10, 30)]),
|
|
id="obb",
|
|
),
|
|
pytest.param(
|
|
Detections(xyxy=np.array([[0, 0, 10, 5]], dtype=np.float32)),
|
|
id="aabb",
|
|
),
|
|
],
|
|
)
|
|
def test_matches_detections_area_property(self, detections: Detections) -> None:
|
|
"""Detections.area delegates to the shared geometry dispatch helper."""
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_array_equal(area, detections.area)
|
|
|
|
def test_keeps_oriented_area_invariant_when_envelope_changes(self) -> None:
|
|
"""Rotating an OBB preserves exact OBB area while changing envelope area."""
|
|
axis_aligned = _detections_from_quads([_rotated_rect(30, 30, 20, 10, 0)])
|
|
rotated = _detections_from_quads([_rotated_rect(30, 30, 20, 10, 45)])
|
|
|
|
areas = np.concatenate([detection_area(axis_aligned), detection_area(rotated)])
|
|
|
|
np.testing.assert_allclose(areas, np.array([200.0, 200.0]))
|
|
assert axis_aligned.box_area[0] != pytest.approx(rotated.box_area[0])
|
|
|
|
def test_mask_area_handles_rotated_arrays(self) -> None:
|
|
"""Mask area counts true pixels consistently after array rotation."""
|
|
mask = np.zeros((80, 80), dtype=bool)
|
|
mask[30:50, 25:55] = 1
|
|
rotated_mask = np.rot90(mask)
|
|
detections = Detections(
|
|
xyxy=np.array([[0, 0, 79, 79], [0, 0, 79, 79]], dtype=np.float32),
|
|
mask=np.stack([mask, rotated_mask]),
|
|
)
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_array_equal(area, np.array([600, 600], dtype=np.int64))
|
|
|
|
def test_degenerate_collinear_obb_has_zero_area(self) -> None:
|
|
"""A collinear (zero-area) OBB reports 0 rather than a well-formed area."""
|
|
collinear = np.array([[0, 0], [5, 0], [10, 0], [15, 0]], dtype=np.float32)
|
|
detections = _detections_from_quads([collinear])
|
|
|
|
area = detection_area(detections)
|
|
|
|
np.testing.assert_allclose(area, np.array([0.0]))
|
|
|
|
|
|
class TestDetectionIou:
|
|
"""Tests for geometry-aware IoU dispatch."""
|
|
|
|
def test_returns_mask_iou_when_both_operands_have_masks(self) -> None:
|
|
"""Mask IoU is used when both operands carry masks."""
|
|
masks_a = np.zeros((1, 16, 16), dtype=bool)
|
|
masks_b = np.zeros((1, 16, 16), dtype=bool)
|
|
masks_a[0, 2:10, 2:10] = True
|
|
masks_b[0, 6:14, 6:14] = True
|
|
detections_a = Detections(
|
|
xyxy=np.array([[0, 0, 16, 16]], dtype=np.float32),
|
|
mask=masks_a,
|
|
)
|
|
detections_b = Detections(
|
|
xyxy=np.array([[0, 0, 16, 16]], dtype=np.float32),
|
|
mask=masks_b,
|
|
)
|
|
|
|
iou = detection_iou(detections_a, detections_b)
|
|
|
|
np.testing.assert_allclose(iou, mask_iou_batch(masks_a, masks_b))
|
|
|
|
def test_returns_oriented_box_iou_when_both_operands_have_obbs(self) -> None:
|
|
"""OBB IoU is used even when AABB envelopes are identical."""
|
|
square = np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32)
|
|
diamond = np.array([[5, 0], [10, 5], [5, 10], [0, 5]], dtype=np.float32)
|
|
shared_xyxy = np.array([[0, 0, 10, 10]], dtype=np.float32)
|
|
detections_a = _detections_from_quads([square], xyxy=shared_xyxy)
|
|
detections_b = _detections_from_quads([diamond], xyxy=shared_xyxy)
|
|
|
|
iou = detection_iou(detections_a, detections_b)
|
|
|
|
np.testing.assert_allclose(
|
|
iou,
|
|
oriented_box_iou_batch(square[np.newaxis], diamond[np.newaxis]),
|
|
)
|
|
assert iou[0, 0] < box_iou_batch(shared_xyxy, shared_xyxy)[0, 0]
|
|
|
|
def test_degenerate_collinear_obb_yields_zero_iou_without_error(self) -> None:
|
|
"""A zero-area (collinear) OBB denominator yields 0 IoU, not NaN or a crash."""
|
|
collinear = np.array([[0, 0], [5, 0], [10, 0], [15, 0]], dtype=np.float32)
|
|
square = np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32)
|
|
detections_a = _detections_from_quads([collinear])
|
|
detections_b = _detections_from_quads([square])
|
|
|
|
iou = detection_iou(detections_a, detections_b)
|
|
|
|
assert not np.isnan(iou).any()
|
|
np.testing.assert_allclose(iou, np.array([[0.0]]))
|
|
|
|
def test_returns_box_iou_when_no_richer_geometry_is_present(self) -> None:
|
|
"""AABB IoU is the fallback when neither operand carries richer geometry."""
|
|
detections_a = Detections(
|
|
xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]], dtype=np.float32)
|
|
)
|
|
detections_b = Detections(xyxy=np.array([[5, 5, 15, 15]], dtype=np.float32))
|
|
|
|
iou = detection_iou(detections_a, detections_b, OverlapMetric.IOS)
|
|
|
|
np.testing.assert_allclose(
|
|
iou,
|
|
box_iou_batch(detections_a.xyxy, detections_b.xyxy, OverlapMetric.IOS),
|
|
)
|
|
|
|
def test_returns_empty_matrix_for_empty_detections(self) -> None:
|
|
"""Empty Detections return an empty pairwise IoU matrix."""
|
|
empty = Detections(xyxy=np.empty((0, 4), dtype=np.float32))
|
|
non_empty = Detections(xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32))
|
|
|
|
iou = detection_iou(empty, non_empty)
|
|
|
|
assert iou.shape == (0, 1)
|
|
|
|
@pytest.mark.parametrize(
|
|
("left_kind", "right_kind", "expected_geometry"),
|
|
[
|
|
pytest.param("mask", "obb", "box", id="mask-vs-obb-falls-back-to-box"),
|
|
pytest.param("mask", "aabb", "box", id="mask-vs-aabb-falls-back-to-box"),
|
|
pytest.param("obb", "aabb", "box", id="obb-vs-aabb-falls-back-to-box"),
|
|
pytest.param("compact", "mask", "mask", id="compact-vs-dense-uses-mask"),
|
|
pytest.param(
|
|
"compact", "compact", "mask", id="compact-vs-compact-uses-mask"
|
|
),
|
|
],
|
|
)
|
|
def test_mixed_geometry_dispatch_uses_shared_geometry_only_when_available(
|
|
self, left_kind: str, right_kind: str, expected_geometry: str
|
|
) -> None:
|
|
"""Mixed geometry uses AABB fallback unless both operands carry masks."""
|
|
image_shape = (16, 16)
|
|
mask = np.zeros((1, *image_shape), dtype=bool)
|
|
mask[0, 2:10, 2:10] = True
|
|
xyxy = _full_image_xyxy(1, image_shape)
|
|
compact_mask = CompactMask.from_dense(
|
|
masks=mask, xyxy=xyxy, image_shape=image_shape
|
|
)
|
|
obb = _detections_from_quads(
|
|
[np.array([[0, 0], [15, 0], [15, 15], [0, 15]], dtype=np.float32)],
|
|
xyxy=xyxy,
|
|
)
|
|
detections_by_kind = {
|
|
"mask": Detections(xyxy=xyxy, mask=mask),
|
|
"compact": Detections(xyxy=xyxy, mask=compact_mask),
|
|
"obb": obb,
|
|
"aabb": Detections(xyxy=xyxy),
|
|
}
|
|
detections_left = detections_by_kind[left_kind]
|
|
detections_right = detections_by_kind[right_kind]
|
|
|
|
result = detection_iou(detections_left, detections_right)
|
|
|
|
if expected_geometry == "mask":
|
|
assert detections_left.mask is not None
|
|
assert detections_right.mask is not None
|
|
expected = mask_iou_batch(detections_left.mask, detections_right.mask)
|
|
else:
|
|
expected = box_iou_batch(detections_left.xyxy, detections_right.xyxy)
|
|
np.testing.assert_allclose(result, expected)
|