roboflow--supervision/tests/detection/tools/test_inference_slicer.py

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,
    move_detections,
)
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


class TestInferenceSlicerBatch:
    """Tests for InferenceSlicer batch_size > 1 path."""

    @pytest.mark.parametrize(
        "batch_size",
        [
            pytest.param(0, id="zero"),
            pytest.param(-1, id="negative"),
        ],
    )
    def test_raises_on_invalid_batch_size(self, batch_size: int) -> None:
        """ValueError raised for batch_size < 1."""
        with pytest.raises(ValueError, match="batch_size"):
            InferenceSlicer(
                callback=lambda x: Detections.empty(), batch_size=batch_size
            )

    def test_batch_size_one_callback_receives_ndarray(self) -> None:
        """batch_size=1 delivers np.ndarray to callback, not list."""
        received_types: list[type] = []

        def callback(tile: np.ndarray) -> Detections:
            received_types.append(type(tile))
            return Detections.empty()

        image = np.zeros((128, 128, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback, slice_wh=64, overlap_wh=0, batch_size=1
        )
        slicer(image)

        assert all(t is np.ndarray for t in received_types)

    def test_batch_callback_receives_list_of_ndarrays(self) -> None:
        """batch_size > 1 delivers list[np.ndarray] to callback."""
        received: list[list] = []

        def callback(tiles: list) -> list:
            received.append(tiles)
            return [Detections.empty() for _ in tiles]

        # 128x128, slice_wh=64, overlap=0 → 4 slices → 2 batches of 2
        image = np.zeros((128, 128, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback, slice_wh=64, overlap_wh=0, batch_size=2
        )
        slicer(image)

        assert len(received) == 2
        assert all(isinstance(batch, list) for batch in received)
        assert all(isinstance(tile, np.ndarray) for batch in received for tile in batch)

    @pytest.mark.parametrize(
        ("image_wh", "batch_size", "expected_batch_sizes"),
        [
            pytest.param((320, 64), 3, [3, 2], id="5-slices-batch-3"),
            pytest.param((256, 64), 4, [4], id="4-slices-batch-4"),
            pytest.param((384, 64), 2, [2, 2, 2], id="6-slices-batch-2"),
        ],
    )
    def test_last_batch_shorter_when_not_divisible(
        self,
        image_wh: tuple[int, int],
        batch_size: int,
        expected_batch_sizes: list[int],
    ) -> None:
        """Last batch has remaining slices when total not divisible by batch_size."""
        received_batch_sizes: list[int] = []

        def callback(tiles: list) -> list:
            received_batch_sizes.append(len(tiles))
            return [Detections.empty() for _ in tiles]

        image = np.zeros((image_wh[1], image_wh[0], 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback,
            slice_wh=64,
            overlap_wh=0,
            batch_size=batch_size,
            thread_workers=1,
        )
        slicer(image)

        assert received_batch_sizes == expected_batch_sizes

    def test_batch_wrong_return_type_raises(self) -> None:
        """ValueError raised when batch callback returns Detections instead of list."""

        def callback(tiles: list) -> Detections:  # type: ignore[return]
            return Detections.empty()

        image = np.zeros((128, 128, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback, slice_wh=64, overlap_wh=0, batch_size=2
        )
        with pytest.raises(ValueError, match="list\\[Detections\\]"):
            slicer(image)

    def test_batch_length_mismatch_raises(self) -> None:
        """ValueError raised when batch callback returns list of wrong length."""

        def callback(tiles: list) -> list:
            return [Detections.empty()]  # always 1, regardless of batch size

        # 128x128, batch_size=4 → one batch of 4 slices; callback returns 1
        image = np.zeros((128, 128, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback, slice_wh=64, overlap_wh=0, batch_size=4
        )
        with pytest.raises(ValueError, match="Lengths must match"):
            slicer(image)

    def test_batch_with_thread_workers_merges_correctly(self) -> None:
        """batch_size + thread_workers > 1 yields correct merged detection count."""
        call_count = 0

        def callback(tiles: list) -> list:
            nonlocal call_count
            call_count += 1
            return [
                Detections(xyxy=np.array([[0, 0, 10, 10]], dtype=float)) for _ in tiles
            ]

        # 128x128, slice_wh=64, overlap=0 → 4 slices → 2 batches of 2
        image = np.zeros((128, 128, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback,
            slice_wh=64,
            overlap_wh=0,
            batch_size=2,
            thread_workers=4,
            overlap_filter=OverlapFilter.NONE,
        )
        detections = slicer(image)

        assert call_count == 2
        assert len(detections) == 4

    def test_batch_obb_forces_sequential_and_warns(self) -> None:
        """OBB in first batch forces sequential execution and emits one warning."""
        active_calls = 0
        max_active_calls = 0
        lock = threading.Lock()

        def callback(tiles: list) -> list:
            nonlocal active_calls, max_active_calls
            with lock:
                active_calls += 1
                max_active_calls = max(max_active_calls, active_calls)
            with lock:
                active_calls -= 1
            return [
                Detections(
                    xyxy=np.array([[0, 0, 10, 10]], dtype=float),
                    data={
                        ORIENTED_BOX_COORDINATES: np.array(
                            [[[0, 0], [10, 0], [10, 10], [0, 10]]], dtype=float
                        )
                    },
                )
                for _ in tiles
            ]

        # 192x192, batch_size=2 → 9 slices → 5 batches; first sync, 4 remaining
        image = np.zeros((192, 192, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback,
            slice_wh=64,
            overlap_wh=0,
            batch_size=2,
            thread_workers=4,
            overlap_filter=OverlapFilter.NONE,
        )

        with pytest.warns(SupervisionWarnings, match="oriented bounding boxes"):
            detections = slicer(image)

        assert max_active_calls == 1
        assert len(detections) == 9

    def test_batch_warns_out_of_bounds_once(self) -> None:
        """Out-of-slice-bounds warning fires exactly once in batch path."""

        def callback(tiles: list) -> list:
            return [
                Detections(xyxy=np.array([[0, 0, 512, 512]], dtype=float))
                for _ in tiles
            ]

        image = np.zeros((128, 128, 3), dtype=np.uint8)
        slicer = InferenceSlicer(
            callback=callback,
            slice_wh=64,
            overlap_wh=0,
            batch_size=2,
            overlap_filter=OverlapFilter.NONE,
        )
        with pytest.warns(SupervisionWarnings, match="outside the slice bounds"):
            slicer(image)

    def test_move_detections_returns_a_copy(self) -> None:
        """move_detections must not mutate the caller's Detections object."""
        detections = Detections(
            xyxy=np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32),
            class_id=np.array([0]),
        )
        original_xyxy = detections.xyxy.copy()

        moved = move_detections(
            detections=detections,
            offset=np.array([10, 20]),
            resolution_wh=(100, 100),
        )

        np.testing.assert_array_equal(detections.xyxy, original_xyxy)
        np.testing.assert_array_equal(moved.xyxy, np.array([[11.0, 22.0, 13.0, 24.0]]))