cmvr_edge_ai/tests/unit/test_detection_annotation.py

from __future__ import annotations

from io import BytesIO
from typing import Any

import pytest

from cmvr_edge_ai.contracts import BoundingBox, Detection, EncodedImage, ImageFrame
from cmvr_edge_ai.detection.annotation import render_annotated_jpeg


Image = pytest.importorskip("PIL.Image")
ImageDraw = pytest.importorskip("PIL.ImageDraw")


def _solid_frame(
    *,
    pixel_format: str,
    rgb: tuple[int, int, int],
    width: int = 64,
    height: int = 48,
) -> ImageFrame:
    stored_pixel = rgb if pixel_format == "RGB8" else tuple(reversed(rgb))
    return ImageFrame(
        data=bytes(stored_pixel) * width * height,
        width=width,
        height=height,
        pixel_format=pixel_format,
    )


def _detection(
    label: str,
    box: BoundingBox,
    *,
    confidence: float = 0.93,
) -> Detection:
    return Detection(label=label, confidence=confidence, box=box)


def _record_drawing_calls(
    monkeypatch: pytest.MonkeyPatch,
) -> tuple[list[tuple[float, float, float, float]], list[str]]:
    rectangle_calls: list[tuple[float, float, float, float]] = []
    text_calls: list[str] = []
    original_rectangle = ImageDraw.ImageDraw.rectangle
    original_text = ImageDraw.ImageDraw.text

    def recording_rectangle(
        draw: ImageDraw.ImageDraw,
        xy: Any,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        coordinates = tuple(xy)
        if len(coordinates) == 2 and all(
            isinstance(point, (tuple, list)) for point in coordinates
        ):
            coordinates = tuple(
                coordinate for point in coordinates for coordinate in point
            )
        if len(coordinates) == 4:
            rectangle_calls.append(tuple(float(value) for value in coordinates))
        original_rectangle(draw, xy, *args, **kwargs)

    def recording_text(
        draw: ImageDraw.ImageDraw,
        xy: Any,
        text: str,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        text_calls.append(text)
        original_text(draw, xy, text, *args, **kwargs)

    monkeypatch.setattr(ImageDraw.ImageDraw, "rectangle", recording_rectangle)
    monkeypatch.setattr(ImageDraw.ImageDraw, "text", recording_text)
    return rectangle_calls, text_calls


@pytest.mark.parametrize("pixel_format", ["BGR8", "RGB8"])
def test_rendered_alert_image_preserves_color_and_is_a_decodable_jpeg(
    pixel_format: str,
) -> None:
    expected_rgb = (24, 96, 184)
    frame = _solid_frame(pixel_format=pixel_format, rgb=expected_rgb)
    detection = _detection(
        "No-Helmet",
        BoundingBox(x_min=8.0, y_min=8.0, x_max=32.0, y_max=32.0),
    )

    encoded = render_annotated_jpeg(frame, (detection,), jpeg_quality=90)

    assert isinstance(encoded, EncodedImage)
    assert encoded.media_type == "image/jpeg"
    assert (encoded.width, encoded.height) == (frame.width, frame.height)
    assert encoded.data.startswith(b"\xff\xd8")
    assert encoded.data.endswith(b"\xff\xd9")

    with Image.open(BytesIO(encoded.data)) as decoded:
        assert decoded.format == "JPEG"
        assert decoded.size == (frame.width, frame.height)
        rgb_image = decoded.convert("RGB")

    # Sample well away from the box and its label. JPEG is lossy, so compare
    # with a tolerance while still catching an RGB/BGR channel swap.
    assert rgb_image.getpixel((frame.width - 3, frame.height - 3)) == pytest.approx(
        expected_rgb,
        abs=12,
    )


def test_renderer_draws_each_detection_box_and_label(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    rectangle_calls, text_calls = _record_drawing_calls(monkeypatch)

    frame = _solid_frame(pixel_format="BGR8", rgb=(0, 0, 0))
    detections = (
        _detection(
            "No-Helmet",
            BoundingBox(x_min=5.0, y_min=10.0, x_max=30.0, y_max=35.0),
        ),
        _detection(
            "No-Vest",
            BoundingBox(x_min=34.0, y_min=12.0, x_max=58.0, y_max=40.0),
            confidence=0.81,
        ),
    )

    render_annotated_jpeg(frame, detections, jpeg_quality=85)

    assert (5.0, 10.0, 30.0, 35.0) in rectangle_calls
    assert (34.0, 12.0, 58.0, 40.0) in rectangle_calls
    assert any("No-Helmet" in text for text in text_calls)
    assert any("No-Vest" in text for text in text_calls)


def test_renderer_clips_out_of_bounds_box_and_skips_invalid_boxes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    rectangle_calls, text_calls = _record_drawing_calls(monkeypatch)
    frame = _solid_frame(pixel_format="RGB8", rgb=(0, 0, 0))
    detections = (
        _detection(
            "Clipped",
            BoundingBox(x_min=-5.0, y_min=-10.0, x_max=80.0, y_max=60.0),
        ),
        _detection(
            "Reversed",
            BoundingBox(x_min=30.0, y_min=30.0, x_max=10.0, y_max=10.0),
        ),
        _detection(
            "Outside",
            BoundingBox(x_min=70.0, y_min=5.0, x_max=90.0, y_max=20.0),
        ),
        _detection(
            "Not-Finite",
            BoundingBox(x_min=float("nan"), y_min=1.0, x_max=10.0, y_max=10.0),
        ),
    )

    render_annotated_jpeg(frame, detections, jpeg_quality=85)

    assert (0.0, 0.0, 63.0, 47.0) in rectangle_calls
    assert any("Clipped" in text for text in text_calls)
    assert all("Reversed" not in text for text in text_calls)
    assert all("Outside" not in text for text in text_calls)
    assert all("Not-Finite" not in text for text in text_calls)