roboflow--supervision/tests/annotators/test_utils.py

from contextlib import ExitStack as DoesNotRaise
from typing import cast

import numpy as np
import pytest

from supervision.annotators.utils import (
    ColorLookup,
    Trace,
    get_labels_text,
    hex_to_rgba,
    is_valid_hex,
    resolve_color,
    resolve_color_idx,
    resolve_text_background_xyxy,
    rgba_to_hex,
    wrap_text,
)
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
from tests.helpers import _create_detections


@pytest.mark.parametrize(
    ("detections", "detection_idx", "color_lookup", "expected_result", "exception"),
    [
        (
            _create_detections(
                xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
                class_id=[5, 3],
                tracker_id=[2, 6],
            ),
            0,
            ColorLookup.INDEX,
            0,
            DoesNotRaise(),
        ),  # multiple detections; index lookup
        (
            _create_detections(
                xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
                class_id=[5, 3],
                tracker_id=[2, 6],
            ),
            0,
            ColorLookup.CLASS,
            5,
            DoesNotRaise(),
        ),  # multiple detections; class lookup
        (
            _create_detections(
                xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
                class_id=[5, 3],
                tracker_id=[2, 6],
            ),
            0,
            ColorLookup.TRACK,
            2,
            DoesNotRaise(),
        ),  # multiple detections; track lookup
        (
            Detections.empty(),
            0,
            ColorLookup.INDEX,
            None,
            pytest.raises(ValueError, match="out of bounds for detections of length 0"),
        ),  # no detections; index lookup; out of bounds
        (
            _create_detections(
                xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
                class_id=[5, 3],
                tracker_id=[2, 6],
            ),
            2,
            ColorLookup.INDEX,
            None,
            pytest.raises(ValueError, match="Detection index 2"),
        ),  # multiple detections; index lookup; out of bounds
        (
            _create_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
            0,
            ColorLookup.CLASS,
            None,
            pytest.raises(ValueError, match="resolve color by class"),
        ),  # multiple detections; class lookup; no class_id
        (
            _create_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
            0,
            ColorLookup.TRACK,
            None,
            pytest.raises(ValueError, match="resolve color by track"),
        ),  # multiple detections; class lookup; no track_id
        (
            _create_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
            0,
            np.array([1, 0]),
            1,
            DoesNotRaise(),
        ),  # multiple detections; custom lookup; correct length
        (
            _create_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
            0,
            np.array([1]),
            None,
            pytest.raises(ValueError, match="Length of color lookup 1"),
        ),  # multiple detections; custom lookup; wrong length
    ],
)
def test_resolve_color_idx(
    detections: Detections,
    detection_idx: int,
    color_lookup: ColorLookup | np.ndarray,
    expected_result: int | None,
    exception: Exception,
) -> None:
    with exception:
        result = resolve_color_idx(
            detections=detections,
            detection_idx=detection_idx,
            color_lookup=color_lookup,
        )
        assert result == expected_result


def test_resolve_color_accepts_negative_class_id() -> None:
    """Negative class ids map deterministically instead of aborting annotation."""
    detections = _create_detections(xyxy=[[0, 0, 10, 10]], class_id=[-1])
    palette = ColorPalette.from_hex(["#ff0000", "#00ff00", "#0000ff"])

    result = resolve_color(
        color=palette,
        detections=detections,
        detection_idx=0,
        color_lookup=ColorLookup.CLASS,
    )

    assert result == Color.from_hex("#0000ff")


def test_resolve_color_accepts_palette_from_src_namespace() -> None:
    """Palette instances imported through src.supervision still resolve by index."""
    from src.supervision.draw.color import ColorPalette as SrcColorPalette

    detections = _create_detections(xyxy=[[0, 0, 10, 10]], class_id=[0])

    result = resolve_color(
        color=SrcColorPalette.DEFAULT,
        detections=detections,
        detection_idx=0,
        color_lookup=ColorLookup.CLASS,
    )

    assert result.as_bgr() == (251, 81, 163)


@pytest.mark.parametrize(
    ("text", "max_line_length", "expected_result", "exception"),
    [
        (None, None, [""], DoesNotRaise()),  # text is None
        ("", None, [""], DoesNotRaise()),  # empty string
        ("   \t  ", 3, [""], DoesNotRaise()),  # whitespace-only (spaces + tab)
        (12345, None, ["12345"], DoesNotRaise()),  # plain integer
        (-6789, None, ["-6789"], DoesNotRaise()),  # negative integer
        (np.int64(1000), None, ["1000"], DoesNotRaise()),  # NumPy int64
        ([1, 2, 3], None, ["[1, 2, 3]"], DoesNotRaise()),  # list to string
        (
            "When you play the game of thrones, you win or you die.\nFear cuts deeper than swords.\nA mind needs books as a sword needs a whetstone.",  # noqa: E501
            None,
            [
                "When you play the game of thrones, you win or you die.",
                "Fear cuts deeper than swords.",
                "A mind needs books as a sword needs a whetstone.",
            ],
            DoesNotRaise(),
        ),  # Game-of-Thrones quotes, multiline
        ("\n", None, [""], DoesNotRaise()),  # single newline
        (
            "valarmorghulisvalardoharis",
            6,
            ["valarm", "orghul", "isvala", "rdohar", "is"],
            DoesNotRaise(),
        ),  # long Valyrian phrase, wrapped
        (
            "Winter is coming\nFire and blood",
            10,
            [
                "Winter is",
                "coming",
                "Fire and",
                "blood",
            ],
            DoesNotRaise(),
        ),  # mix of short/long with newline
        (
            "What is dead may never die",
            0,
            None,
            pytest.raises(ValueError, match="max_line_length must be"),
        ),  # width 0 - invalid
        (
            "A Lannister always pays his debts",
            -1,
            None,
            pytest.raises(ValueError, match="positive integer"),
        ),  # width -1 - invalid
        (None, 10, [""], DoesNotRaise()),  # text None, width set
    ],
)
def test_wrap_text(
    text: object,
    max_line_length: int | None,
    expected_result: list[str],
    exception: Exception,
) -> None:
    with exception:
        result = wrap_text(text=text, max_line_length=max_line_length)
        assert result == expected_result


def test_resolve_text_background_xyxy_rejects_unknown_position() -> None:
    """Unsupported positions must raise instead of returning an implicit None."""
    with pytest.raises(ValueError, match="Unsupported position"):
        resolve_text_background_xyxy(
            center_coordinates=(10, 10),
            text_wh=(20, 10),
            position=cast(Position, "invalid"),
        )


def test_trace_put_requires_tracker_id_before_mutation() -> None:
    """Trace.put must not mutate internal history before validating tracker ids."""
    trace = Trace()
    detections = _create_detections(xyxy=[[0, 0, 1, 1]], class_id=[0])

    before_frame_id = trace.current_frame_id
    before_history = trace.frame_id.copy()
    before_xy = trace.xy.copy()
    before_tracker_id = trace.tracker_id.copy()

    with pytest.raises(ValueError, match="tracker_id"):
        trace.put(detections)

    assert trace.current_frame_id == before_frame_id
    assert np.array_equal(trace.frame_id, before_history)
    assert np.array_equal(trace.xy, before_xy)
    assert np.array_equal(trace.tracker_id, before_tracker_id)


@pytest.mark.parametrize(
    ("hex_color", "expected_rgba"),
    [
        ("#FF00FF", (255, 0, 255, 255)),
        ("FF00FF", (255, 0, 255, 255)),
        ("#FF00FF80", (255, 0, 255, 128)),
        ("00FF0080", (0, 255, 0, 128)),
        ("  #ff00ff80  ", (255, 0, 255, 128)),
        ("abcdef", (171, 205, 239, 255)),
    ],
)
def test_hex_to_rgba_valid(
    hex_color: str, expected_rgba: tuple[int, int, int, int]
) -> None:
    assert hex_to_rgba(hex_color) == expected_rgba


@pytest.mark.parametrize(
    "hex_color",
    [
        pytest.param("#FF00F", id="five-digits"),
        pytest.param("#GGHHII", id="non-hex-digits"),
        pytest.param("#FFF", id="short-form"),
        pytest.param("1234567", id="seven-digits"),
        pytest.param("##000000", id="multiple-prefixes"),
    ],
)
def test_hex_to_rgba_invalid(hex_color: str) -> None:
    """Invalid formats raise instead of being normalized into valid colors."""
    with pytest.raises(ValueError, match="Invalid hex"):
        hex_to_rgba(hex_color)


@pytest.mark.parametrize(
    ("rgba", "expected_hex"),
    [
        ((0, 0, 0, 0), "#00000000"),
        ((255, 0, 255, 255), "#FF00FFFF"),
        ((0, 255, 0, 128), "#00FF0080"),
        ((255, 255, 255, 255), "#FFFFFFFF"),
    ],
)
def test_rgba_to_hex(rgba: tuple[int, int, int, int], expected_hex: str) -> None:
    assert rgba_to_hex(rgba) == expected_hex


@pytest.mark.parametrize(
    "rgba",
    [
        (255, 0, 0),
        (256, 0, 0, 255),
        (-1, 0, 0, 255),
        (255, 0, 0, -1),
    ],
)
def test_rgba_to_hex_invalid(rgba: tuple[int, ...]) -> None:
    with pytest.raises(ValueError, match="RGBA must be a 4-tuple"):
        rgba_to_hex(rgba)  # type: ignore[arg-type]


@pytest.mark.parametrize(
    ("hex_color", "expected_result"),
    [
        ("#FF00FF", True),
        ("ff00ff", True),
        ("00FF0080", True),
        (" 00ff0080 ", True),
        ("#XYZ123", False),
        ("FF00F", False),
        ("#FFF", False),
        ("#1234567", False),
    ],
)
def test_is_valid_hex(hex_color: str, expected_result: bool) -> None:
    assert is_valid_hex(hex_color) is expected_result


class TestGetLabelsText:
    """Tests for the get_labels_text helper."""

    def test_returns_custom_labels_when_provided(self) -> None:
        """Custom labels are returned as-is without inspecting detections."""
        detections = _create_detections(xyxy=[[0, 0, 1, 1]], class_id=[0])
        custom = ["cat"]
        assert get_labels_text(detections, custom) is custom

    def test_uses_class_name_data_field(self) -> None:
        """Labels come from CLASS_NAME_DATA_FIELD when present in data."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2]],
            class_id=[0, 1],
            data={"class_name": np.array(["cat", "dog"])},
        )
        assert get_labels_text(detections, None) == ["cat", "dog"]

    def test_falls_back_to_class_id(self) -> None:
        """Labels fall back to class_id strings when no class_name data."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2]],
            class_id=[5, 3],
        )
        assert get_labels_text(detections, None) == ["5", "3"]

    def test_falls_back_to_class_id_single_detection(self) -> None:
        """class_id fallback holds for a single detection (N=1 edge case)."""
        detections = _create_detections(xyxy=[[0, 0, 1, 1]], class_id=[7])
        assert get_labels_text(detections, None) == ["7"]

    def test_falls_back_to_index(self) -> None:
        """Labels fall back to detection indices when no class_id."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2], [2, 2, 3, 3]],
        )
        assert get_labels_text(detections, None) == ["0", "1", "2"]

    def test_falls_back_to_index_single_detection(self) -> None:
        """Index fallback holds for a single detection (N=1 edge case)."""
        detections = _create_detections(xyxy=[[0, 0, 1, 1]])
        assert get_labels_text(detections, None) == ["0"]

    def test_custom_labels_take_precedence_over_class_name(self) -> None:
        """Custom labels win over the class_name branch specifically."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2]],
            class_id=[0, 1],
            data={"class_name": np.array(["cat", "dog"])},
        )
        custom = ["a", "b"]
        assert get_labels_text(detections, custom) is custom

    def test_uses_class_name_data_field_with_non_str_dtype(self) -> None:
        """str() coercion is load-bearing when class_name holds a non-str dtype."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2]],
            class_id=[0, 1],
            data={"class_name": np.array([1, 2])},
        )
        assert get_labels_text(detections, None) == ["1", "2"]

    def test_empty_detections(self) -> None:
        """Empty detections produce an empty label list, via the class_id branch."""
        detections = Detections.empty()
        assert detections.class_id is not None
        assert get_labels_text(detections, None) == []

    def test_raises_on_class_name_length_mismatch(self) -> None:
        """A class_name array shorter than detections raises instead of truncating."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2]],
            class_id=[0, 1],
            data={"class_name": np.array(["cat", "dog"])},
        )
        detections.data["class_name"] = np.array(["cat"])  # bypass alignment checks

        with pytest.raises(ValueError, match="class_name"):
            get_labels_text(detections, None)

    def test_raises_on_class_id_length_mismatch(self) -> None:
        """A class_id array shorter than detections raises instead of truncating."""
        detections = _create_detections(
            xyxy=[[0, 0, 1, 1], [1, 1, 2, 2]],
            class_id=[0, 1],
        )
        detections.class_id = np.array([0])  # bypass alignment checks

        with pytest.raises(ValueError, match="class_id"):
            get_labels_text(detections, None)