roboflow--supervision/tests/detection/test_compact_mask_integration.py

"""Integration tests: CompactMask <-> Detections, annotators, merge."""

from contextlib import ExitStack as DoesNotRaise

import numpy as np
import pytest

import supervision as sv
from supervision.detection.compact_mask import CompactMask
from supervision.detection.core import Detections


def _full_xyxy(n: int, h: int, w: int) -> np.ndarray:
    """N boxes covering the whole image (ensures crop == full mask)."""
    return np.tile(np.array([0, 0, w, h], dtype=np.float32), (n, 1))


def _make_compact_detections(
    n: int, h: int = 40, w: int = 40
) -> tuple[Detections, np.ndarray]:
    """Detections with a CompactMask backed by full-image bounding boxes.

    Using full-image xyxy means all True pixels are within the crop region,
    so from_dense -> to_dense is lossless.
    """
    rng = np.random.default_rng(42)
    masks = rng.integers(0, 2, size=(n, h, w)).astype(bool)
    xyxy = _full_xyxy(n, h, w)
    cm = CompactMask.from_dense(masks, xyxy, image_shape=(h, w))
    det = Detections(
        xyxy=xyxy,
        mask=cm,
        confidence=np.ones(n, dtype=np.float32) * 0.9,
        class_id=np.arange(n),
    )
    return det, masks


class TestConstruction:
    """Tests for building Detections with a CompactMask.

    Verifies that a CompactMask is accepted as a valid mask argument and that
    the validator raises ValueError when the mask length does not match the
    number of bounding boxes.
    """

    def test_detections_construction_with_compact_mask(self) -> None:
        with DoesNotRaise():
            det, _ = _make_compact_detections(3)
        assert isinstance(det.mask, CompactMask)
        assert len(det) == 3

    def test_detections_compact_mask_validation_mismatch(self) -> None:
        n, h, w = 3, 20, 20
        xyxy = _full_xyxy(n, h, w)
        masks_wrong_n = np.zeros((n + 1, h, w), dtype=bool)
        cm = CompactMask.from_dense(masks_wrong_n, _full_xyxy(n + 1, h, w), (h, w))
        with pytest.raises(ValueError, match="mask must contain"):
            Detections(xyxy=xyxy, mask=cm)


class TestFiltering:
    """Tests for Detections.__getitem__ with a CompactMask.

    Verifies that integer, slice, and boolean-array indexing all preserve the
    CompactMask type and return the correct subset of masks.
    """

    def test_int_wraps_to_compact_mask(self) -> None:
        det, _ = _make_compact_detections(3)
        # Detections converts int to [int] internally -> subset has 1 element
        subset = det[1]
        assert isinstance(subset.mask, CompactMask)
        assert len(subset) == 1

    def test_slice_preserves_compact_mask(self) -> None:
        det, masks = _make_compact_detections(4)
        subset = det[1:3]
        assert isinstance(subset.mask, CompactMask)
        assert len(subset) == 2
        np.testing.assert_array_equal(subset.mask.to_dense(), masks[1:3])

    def test_bool_array_preserves_compact_mask(self) -> None:
        det, masks = _make_compact_detections(4)
        selector = np.array([True, False, True, False])
        subset = det[selector]
        assert isinstance(subset.mask, CompactMask)
        assert len(subset) == 2
        np.testing.assert_array_equal(subset.mask.to_dense(), masks[[0, 2]])


class TestIteration:
    """Tests for iterating over Detections with a CompactMask.

    Verifies that each iteration step yields a 2-D boolean (H, W) array
    identical to the corresponding dense mask, so downstream code that
    iterates over detections needs no changes.
    """

    def test_iter_yields_2d_dense(self) -> None:
        h, w = 20, 20
        det, masks = _make_compact_detections(3, h, w)
        for i, (_, mask_2d, *_) in enumerate(det):
            assert mask_2d is not None
            assert isinstance(mask_2d, np.ndarray)
            assert mask_2d.shape == (h, w)
            assert mask_2d.dtype == bool
            np.testing.assert_array_equal(mask_2d, masks[i])


class TestEquality:
    """Tests for Detections.__eq__ mixing CompactMask and dense arrays.

    Verifies that a Detections object backed by a CompactMask compares equal
    to an otherwise identical Detections object backed by a dense ndarray.
    """

    def test_compact_vs_dense(self) -> None:
        h, w = 20, 20
        det_compact, masks = _make_compact_detections(2, h, w)
        xyxy = det_compact.xyxy.copy()
        det_dense = Detections(
            xyxy=xyxy,
            mask=masks,
            confidence=np.ones(2, dtype=np.float32) * 0.9,
            class_id=np.arange(2),
        )
        assert det_compact == det_dense


class TestArea:
    """Tests for the Detections.area property with a CompactMask.

    Verifies that the fast CompactMask path in Detections.area returns the
    same per-detection pixel counts as summing the equivalent dense array.
    """

    def test_compact_matches_dense(self) -> None:
        det_compact, masks = _make_compact_detections(3)
        expected_area = np.array([m.sum() for m in masks])
        np.testing.assert_array_equal(det_compact.area, expected_area)


class TestMerge:
    """Tests for merging Detections objects that contain CompactMask instances.

    Covers three scenarios:
    - All-compact merge: result is a CompactMask.
    - Mixed compact + dense: dense inputs are converted; result is a CompactMask.
    - Inner pair merge (merge_inner_detection_object_pair): used during NMS-like
      operations, each input must contain exactly one detection.
    """

    def test_all_compact(self) -> None:
        h, w = 30, 30
        det1, masks1 = _make_compact_detections(2, h, w)

        rng = np.random.default_rng(7)
        masks2 = rng.integers(0, 2, size=(3, h, w)).astype(bool)
        xyxy2 = _full_xyxy(3, h, w)
        cm2 = CompactMask.from_dense(masks2, xyxy2, (h, w))
        det2 = Detections(
            xyxy=xyxy2,
            mask=cm2,
            confidence=np.ones(3, dtype=np.float32) * 0.8,
            class_id=np.arange(3),
        )

        merged = Detections.merge([det1, det2])
        assert isinstance(merged.mask, CompactMask)
        assert len(merged) == 5
        expected = np.concatenate([masks1, masks2], axis=0)
        np.testing.assert_array_equal(merged.mask.to_dense(), expected)

    def test_mixed_compact_and_dense(self) -> None:
        """Merging a CompactMask with a dense ndarray returns a CompactMask."""
        h, w = 20, 20
        det_compact, masks_compact = _make_compact_detections(2, h, w)
        masks_dense = np.zeros((1, h, w), dtype=bool)
        masks_dense[0, 3:8, 3:8] = True
        xyxy_dense = _full_xyxy(1, h, w)
        det_dense = Detections(
            xyxy=xyxy_dense,
            mask=masks_dense,
            confidence=np.array([0.5], dtype=np.float32),
            class_id=np.array([0]),
        )

        merged = Detections.merge([det_compact, det_dense])
        assert isinstance(merged.mask, CompactMask)
        assert len(merged) == 3
        expected = np.concatenate([masks_compact, masks_dense], axis=0)
        np.testing.assert_array_equal(merged.mask.to_dense(), expected)
        assert merged.mask.image_shape == (h, w)

    def test_inner_pair_with_compact(self) -> None:
        from supervision.detection.core import merge_inner_detection_object_pair

        h, w = 20, 20
        masks_a = np.zeros((1, h, w), dtype=bool)
        masks_a[0, 0:5, 0:5] = True
        xyxy_a = _full_xyxy(1, h, w)
        cm_a = CompactMask.from_dense(masks_a, xyxy_a, (h, w))
        det_a = Detections(
            xyxy=xyxy_a,
            mask=cm_a,
            confidence=np.array([0.9], dtype=np.float32),
            class_id=np.array([1]),
        )

        masks_b = np.zeros((1, h, w), dtype=bool)
        masks_b[0, 5:10, 5:10] = True
        xyxy_b = _full_xyxy(1, h, w)
        cm_b = CompactMask.from_dense(masks_b, xyxy_b, (h, w))
        det_b = Detections(
            xyxy=xyxy_b,
            mask=cm_b,
            confidence=np.array([0.7], dtype=np.float32),
            class_id=np.array([1]),
        )

        with DoesNotRaise():
            result = merge_inner_detection_object_pair(det_a, det_b)
        assert len(result) == 1


class TestAnnotators:
    """Tests for annotators that consume CompactMask via Detections.

    Verifies that MaskAnnotator and PolygonAnnotator produce pixel-identical
    output when given Detections backed by a CompactMask versus the equivalent
    dense ndarray, confirming that the annotators are transparent to the mask
    representation.
    """

    def test_mask_annotator(self) -> None:
        h, w = 40, 40
        det_compact, masks = _make_compact_detections(2, h, w)
        det_dense = Detections(
            xyxy=det_compact.xyxy.copy(),
            mask=masks,
            confidence=det_compact.confidence.copy(),
            class_id=det_compact.class_id.copy(),
        )

        image = np.zeros((h, w, 3), dtype=np.uint8)
        annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)

        annotated_compact = annotator.annotate(image.copy(), det_compact)
        annotated_dense = annotator.annotate(image.copy(), det_dense)

        np.testing.assert_array_equal(
            annotated_compact,
            annotated_dense,
            err_msg="MaskAnnotator output differs between CompactMask and dense mask",
        )

    def test_polygon_annotator(self) -> None:
        h, w = 40, 40
        # Use solid rectangular masks for stable polygon results.
        masks = np.zeros((2, h, w), dtype=bool)
        masks[0, 5:15, 5:15] = True
        masks[1, 20:30, 20:30] = True
        xyxy = _full_xyxy(2, h, w)
        cm = CompactMask.from_dense(masks, xyxy, (h, w))

        det_compact = Detections(xyxy=xyxy, mask=cm, class_id=np.array([0, 1]))
        det_dense = Detections(xyxy=xyxy, mask=masks, class_id=np.array([0, 1]))

        image = np.zeros((h, w, 3), dtype=np.uint8)
        annotator = sv.PolygonAnnotator(color_lookup=sv.ColorLookup.INDEX)

        annotated_compact = annotator.annotate(image.copy(), det_compact)
        annotated_dense = annotator.annotate(image.copy(), det_dense)

        np.testing.assert_array_equal(annotated_compact, annotated_dense)