500 lines
16 KiB
Python
500 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
import warnings
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from supervision.config import ORIENTED_BOX_COORDINATES
|
|
from supervision.detection.core import Detections
|
|
from supervision.detection.tools.inference_slicer import InferenceSlicer
|
|
from supervision.detection.utils.iou_and_nms import OverlapFilter
|
|
from supervision.utils.internal import SupervisionWarnings
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_callback():
|
|
"""Mock callback function for testing."""
|
|
|
|
def callback(_: np.ndarray) -> Detections:
|
|
return Detections(xyxy=np.array([[0, 0, 10, 10]]))
|
|
|
|
return callback
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("resolution_wh", "slice_wh", "overlap_wh", "expected_offsets"),
|
|
[
|
|
# Case 1: Square image, square slices, no overlap
|
|
(
|
|
(256, 256),
|
|
(128, 128),
|
|
(0, 0),
|
|
np.array(
|
|
[
|
|
[0, 0, 128, 128],
|
|
[128, 0, 256, 128],
|
|
[0, 128, 128, 256],
|
|
[128, 128, 256, 256],
|
|
]
|
|
),
|
|
),
|
|
# Case 2: Square image, square slices, non-zero overlap
|
|
(
|
|
(256, 256),
|
|
(128, 128),
|
|
(64, 64),
|
|
np.array(
|
|
[
|
|
[0, 0, 128, 128],
|
|
[64, 0, 192, 128],
|
|
[128, 0, 256, 128],
|
|
[0, 64, 128, 192],
|
|
[64, 64, 192, 192],
|
|
[128, 64, 256, 192],
|
|
[0, 128, 128, 256],
|
|
[64, 128, 192, 256],
|
|
[128, 128, 256, 256],
|
|
]
|
|
),
|
|
),
|
|
# Case 3: Rectangle image (horizontal), square slices, no overlap
|
|
(
|
|
(192, 128),
|
|
(64, 64),
|
|
(0, 0),
|
|
np.array(
|
|
[
|
|
[0, 0, 64, 64],
|
|
[64, 0, 128, 64],
|
|
[128, 0, 192, 64],
|
|
[0, 64, 64, 128],
|
|
[64, 64, 128, 128],
|
|
[128, 64, 192, 128],
|
|
]
|
|
),
|
|
),
|
|
# Case 4: Rectangle image (horizontal), square slices, non-zero overlap
|
|
(
|
|
(192, 128),
|
|
(64, 64),
|
|
(32, 32),
|
|
np.array(
|
|
[
|
|
[0, 0, 64, 64],
|
|
[32, 0, 96, 64],
|
|
[64, 0, 128, 64],
|
|
[96, 0, 160, 64],
|
|
[128, 0, 192, 64],
|
|
[0, 32, 64, 96],
|
|
[32, 32, 96, 96],
|
|
[64, 32, 128, 96],
|
|
[96, 32, 160, 96],
|
|
[128, 32, 192, 96],
|
|
[0, 64, 64, 128],
|
|
[32, 64, 96, 128],
|
|
[64, 64, 128, 128],
|
|
[96, 64, 160, 128],
|
|
[128, 64, 192, 128],
|
|
]
|
|
),
|
|
),
|
|
# Case 5: Rectangle image (vertical), square slices, no overlap
|
|
(
|
|
(128, 192),
|
|
(64, 64),
|
|
(0, 0),
|
|
np.array(
|
|
[
|
|
[0, 0, 64, 64],
|
|
[64, 0, 128, 64],
|
|
[0, 64, 64, 128],
|
|
[64, 64, 128, 128],
|
|
[0, 128, 64, 192],
|
|
[64, 128, 128, 192],
|
|
]
|
|
),
|
|
),
|
|
# Case 6: Rectangle image (vertical), square slices, non-zero overlap
|
|
(
|
|
(128, 192),
|
|
(64, 64),
|
|
(32, 32),
|
|
np.array(
|
|
[
|
|
[0, 0, 64, 64],
|
|
[32, 0, 96, 64],
|
|
[64, 0, 128, 64],
|
|
[0, 32, 64, 96],
|
|
[32, 32, 96, 96],
|
|
[64, 32, 128, 96],
|
|
[0, 64, 64, 128],
|
|
[32, 64, 96, 128],
|
|
[64, 64, 128, 128],
|
|
[0, 96, 64, 160],
|
|
[32, 96, 96, 160],
|
|
[64, 96, 128, 160],
|
|
[0, 128, 64, 192],
|
|
[32, 128, 96, 192],
|
|
[64, 128, 128, 192],
|
|
]
|
|
),
|
|
),
|
|
# Case 7: Square image, rectangular slices (horizontal), no overlap
|
|
(
|
|
(160, 160),
|
|
(80, 40),
|
|
(0, 0),
|
|
np.array(
|
|
[
|
|
[0, 0, 80, 40],
|
|
[80, 0, 160, 40],
|
|
[0, 40, 80, 80],
|
|
[80, 40, 160, 80],
|
|
[0, 80, 80, 120],
|
|
[80, 80, 160, 120],
|
|
[0, 120, 80, 160],
|
|
[80, 120, 160, 160],
|
|
]
|
|
),
|
|
),
|
|
# Case 8: Square image, rectangular slices (vertical), non-zero overlap
|
|
(
|
|
(160, 160),
|
|
(40, 80),
|
|
(10, 20),
|
|
np.array(
|
|
[
|
|
[0, 0, 40, 80],
|
|
[30, 0, 70, 80],
|
|
[60, 0, 100, 80],
|
|
[90, 0, 130, 80],
|
|
[120, 0, 160, 80],
|
|
[0, 60, 40, 140],
|
|
[30, 60, 70, 140],
|
|
[60, 60, 100, 140],
|
|
[90, 60, 130, 140],
|
|
[120, 60, 160, 140],
|
|
[0, 80, 40, 160],
|
|
[30, 80, 70, 160],
|
|
[60, 80, 100, 160],
|
|
[90, 80, 130, 160],
|
|
[120, 80, 160, 160],
|
|
]
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_generate_offset(
|
|
resolution_wh: tuple[int, int],
|
|
slice_wh: tuple[int, int],
|
|
overlap_wh: tuple[int, int],
|
|
expected_offsets: np.ndarray,
|
|
) -> None:
|
|
offsets = InferenceSlicer._generate_offset(
|
|
resolution_wh=resolution_wh,
|
|
slice_wh=slice_wh,
|
|
overlap_wh=overlap_wh,
|
|
)
|
|
|
|
assert np.array_equal(offsets, expected_offsets), (
|
|
f"Expected {expected_offsets}, got {offsets}"
|
|
)
|
|
|
|
|
|
def test_run_callback_warns_when_detections_outside_slice_bounds() -> None:
|
|
"""Test that a warning is emitted when callback returns detections with
|
|
coordinates outside the slice bounds."""
|
|
|
|
def out_of_bounds_callback(_: np.ndarray) -> Detections:
|
|
# Return detections with coordinates exceeding the 64x64 slice size
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 128, 128]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
image = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(callback=out_of_bounds_callback, slice_wh=64, overlap_wh=0)
|
|
|
|
with pytest.warns(SupervisionWarnings, match="outside the slice bounds"):
|
|
slicer(image)
|
|
|
|
|
|
def test_run_callback_warns_only_once_for_out_of_bounds_detections() -> None:
|
|
"""Test that the out-of-bounds warning is only emitted once even across
|
|
multiple slices."""
|
|
|
|
def out_of_bounds_callback(_: np.ndarray) -> Detections:
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 128, 128]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
image = np.zeros((256, 256, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(callback=out_of_bounds_callback, slice_wh=64, overlap_wh=0)
|
|
|
|
with warnings.catch_warnings(record=True) as recorded_warnings:
|
|
warnings.simplefilter("always")
|
|
slicer(image)
|
|
|
|
out_of_bounds_warnings = [
|
|
w
|
|
for w in recorded_warnings
|
|
if issubclass(w.category, SupervisionWarnings)
|
|
and "outside the slice bounds" in str(w.message)
|
|
]
|
|
assert len(out_of_bounds_warnings) == 1
|
|
|
|
|
|
def test_run_callback_no_warning_when_detections_inside_slice_bounds() -> None:
|
|
"""Test that no warning is emitted when callback returns detections within
|
|
the slice bounds."""
|
|
|
|
def in_bounds_callback(_: np.ndarray) -> Detections:
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 10, 10]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
image = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(callback=in_bounds_callback, slice_wh=64, overlap_wh=0)
|
|
|
|
with warnings.catch_warnings(record=True) as recorded_warnings:
|
|
warnings.simplefilter("always")
|
|
slicer(image)
|
|
|
|
out_of_bounds_warnings = [
|
|
w
|
|
for w in recorded_warnings
|
|
if issubclass(w.category, SupervisionWarnings)
|
|
and "outside the slice bounds" in str(w.message)
|
|
]
|
|
assert len(out_of_bounds_warnings) == 0
|
|
|
|
|
|
def test_run_callback_warns_when_detections_have_negative_coordinates() -> None:
|
|
"""Test that a warning is emitted when callback returns detections with
|
|
negative coordinates, indicating wrong reference frame."""
|
|
|
|
def negative_coords_callback(_: np.ndarray) -> Detections:
|
|
# Return detections with negative coordinates (e.g., returned in full-image
|
|
# coordinates that are to the left/top of this slice's origin)
|
|
return Detections(
|
|
xyxy=np.array([[-10, -10, 10, 10]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
image = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(
|
|
callback=negative_coords_callback, slice_wh=64, overlap_wh=0
|
|
)
|
|
|
|
with pytest.warns(SupervisionWarnings, match="outside the slice bounds"):
|
|
slicer(image)
|
|
|
|
|
|
def test_run_callback_warns_only_once_with_multiple_threads() -> None:
|
|
"""Test that exactly one warning fires even with thread_workers > 1, validating
|
|
that the threading.Lock makes the check-and-set atomic."""
|
|
|
|
def out_of_bounds_callback(_: np.ndarray) -> Detections:
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 128, 128]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
# 512x512 / 64 slice -> 64 slices; all 4 threads will see out-of-bounds detections
|
|
image = np.zeros((512, 512, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(
|
|
callback=out_of_bounds_callback,
|
|
slice_wh=64,
|
|
overlap_wh=0,
|
|
thread_workers=4,
|
|
)
|
|
|
|
with warnings.catch_warnings(record=True) as recorded_warnings:
|
|
warnings.simplefilter("always")
|
|
slicer(image)
|
|
|
|
out_of_bounds_warnings = [
|
|
w
|
|
for w in recorded_warnings
|
|
if issubclass(w.category, SupervisionWarnings)
|
|
and "outside the slice bounds" in str(w.message)
|
|
]
|
|
assert len(out_of_bounds_warnings) == 1
|
|
|
|
|
|
def test_run_callback_no_warning_for_detection_exactly_at_slice_boundary() -> None:
|
|
"""Test that a detection whose coordinates exactly equal the slice dimensions
|
|
does not trigger the warning (boundary is exclusive: > not >=)."""
|
|
|
|
def at_boundary_callback(_: np.ndarray) -> Detections:
|
|
# x2=64, y2=64 on a 64x64 slice — touching the edge but not exceeding it
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 64, 64]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
image = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(callback=at_boundary_callback, slice_wh=64, overlap_wh=0)
|
|
|
|
with warnings.catch_warnings(record=True) as recorded_warnings:
|
|
warnings.simplefilter("always")
|
|
slicer(image)
|
|
|
|
out_of_bounds_warnings = [
|
|
w
|
|
for w in recorded_warnings
|
|
if issubclass(w.category, SupervisionWarnings)
|
|
and "outside the slice bounds" in str(w.message)
|
|
]
|
|
assert len(out_of_bounds_warnings) == 0
|
|
|
|
|
|
def test_run_callback_does_not_rewarn_on_second_call() -> None:
|
|
"""Test that a second call to the same slicer instance does not re-emit
|
|
the out-of-bounds warning even when detections are still out of bounds."""
|
|
|
|
def out_of_bounds_callback(_: np.ndarray) -> Detections:
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 128, 128]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
)
|
|
|
|
image = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(callback=out_of_bounds_callback, slice_wh=64, overlap_wh=0)
|
|
|
|
with warnings.catch_warnings(record=True) as recorded_warnings:
|
|
warnings.simplefilter("always")
|
|
slicer(image) # first call — warning fires
|
|
slicer(image) # second call — must not re-warn
|
|
|
|
out_of_bounds_warnings = [
|
|
w
|
|
for w in recorded_warnings
|
|
if issubclass(w.category, SupervisionWarnings)
|
|
and "outside the slice bounds" in str(w.message)
|
|
]
|
|
assert len(out_of_bounds_warnings) == 1
|
|
|
|
|
|
def test_obb_callbacks_run_sequentially_even_with_multiple_workers() -> None:
|
|
"""Test that OBB callbacks are serialized even when thread_workers > 1."""
|
|
|
|
active_calls = 0
|
|
max_active_calls = 0
|
|
concurrent_callbacks = 0
|
|
callback_lock = threading.Lock()
|
|
|
|
def obb_callback(_: np.ndarray) -> Detections:
|
|
nonlocal active_calls, max_active_calls, concurrent_callbacks
|
|
|
|
with callback_lock:
|
|
active_calls += 1
|
|
max_active_calls = max(max_active_calls, active_calls)
|
|
if active_calls > 1:
|
|
concurrent_callbacks += 1
|
|
|
|
with callback_lock:
|
|
active_calls -= 1
|
|
|
|
return Detections(
|
|
xyxy=np.array([[0, 0, 10, 10]], dtype=float),
|
|
confidence=np.array([0.9]),
|
|
class_id=np.array([0]),
|
|
data={
|
|
ORIENTED_BOX_COORDINATES: np.array(
|
|
[[[0, 0], [10, 0], [10, 10], [0, 10]]], dtype=float
|
|
)
|
|
},
|
|
)
|
|
|
|
image = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(
|
|
callback=obb_callback,
|
|
slice_wh=64,
|
|
overlap_wh=0,
|
|
thread_workers=4,
|
|
)
|
|
|
|
with pytest.warns(SupervisionWarnings, match="oriented bounding boxes"):
|
|
detections = slicer(image)
|
|
|
|
assert max_active_calls == 1
|
|
assert concurrent_callbacks == 0
|
|
assert len(detections) == 4
|
|
|
|
|
|
def _rotated_rect(
|
|
cx: float, cy: float, w: float, h: float, angle_deg: float
|
|
) -> np.ndarray:
|
|
angle = np.deg2rad(angle_deg)
|
|
cos, sin = np.cos(angle), np.sin(angle)
|
|
rot = np.array([[cos, -sin], [sin, cos]])
|
|
corners = np.array(
|
|
[[-w / 2, -h / 2], [w / 2, -h / 2], [w / 2, h / 2], [-w / 2, h / 2]]
|
|
)
|
|
return (corners @ rot.T + [cx, cy]).astype(np.float32)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"overlap_filter",
|
|
[OverlapFilter.NON_MAX_SUPPRESSION, OverlapFilter.NON_MAX_MERGE],
|
|
)
|
|
def test_inference_slicer_keeps_crossed_obb_detections(
|
|
overlap_filter: OverlapFilter,
|
|
) -> None:
|
|
"""Regression for issue #1679: the SAHI workflow with OBB detections
|
|
dropped valid detections at the merge step because `with_nms`/`with_nmm`
|
|
historically used axis-aligned IoU. For crossed thin rectangles the AABBs
|
|
are nearly identical (IoU ≈ 1.0) while the OBBs barely overlap (IoU ≈ 0.06)
|
|
— so AABB-NMS suppressed one of them.
|
|
|
|
Both crossed OBBs must survive end-to-end through `InferenceSlicer`.
|
|
"""
|
|
quad_a = _rotated_rect(50, 50, 80, 8, +45)
|
|
quad_b = _rotated_rect(50, 50, 80, 8, -45)
|
|
aabb_a = [
|
|
quad_a[:, 0].min(),
|
|
quad_a[:, 1].min(),
|
|
quad_a[:, 0].max(),
|
|
quad_a[:, 1].max(),
|
|
]
|
|
aabb_b = [
|
|
quad_b[:, 0].min(),
|
|
quad_b[:, 1].min(),
|
|
quad_b[:, 0].max(),
|
|
quad_b[:, 1].max(),
|
|
]
|
|
|
|
def callback(_: np.ndarray) -> Detections:
|
|
return Detections(
|
|
xyxy=np.array([aabb_a, aabb_b], dtype=np.float32),
|
|
confidence=np.array([0.9, 0.85], dtype=np.float32),
|
|
class_id=np.array([0, 0], dtype=int),
|
|
data={ORIENTED_BOX_COORDINATES: np.stack([quad_a, quad_b])},
|
|
)
|
|
|
|
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
slicer = InferenceSlicer(
|
|
callback=callback,
|
|
slice_wh=100,
|
|
overlap_wh=0,
|
|
thread_workers=1,
|
|
overlap_filter=overlap_filter,
|
|
iou_threshold=0.5,
|
|
)
|
|
|
|
detections = slicer(image)
|
|
|
|
assert len(detections) == 2
|