cmvr_edge_ai/tests/unit/test_yolo_model.py

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace
from typing import Any

import pytest

from cmvr_edge_ai.contracts import BoundingBox, Detection, ImageFrame
from cmvr_edge_ai.detection.models import (
    MOBILE_PHONE_LABELS,
    PEOPLE_TALKING_LABELS,
    PPE_6CLASS_LABELS,
    register_builtin_detection_models,
)
from cmvr_edge_ai.detection.models import yolo as yolo_module
from cmvr_edge_ai.detection.models.yolo import (
    UltralyticsYoloModel,
    YoloDependencyError,
)
from cmvr_edge_ai.detection.registry import DetectionModelRegistry


class _FakeArray:
    def __init__(self, data: bytes, *, channel_reversed: bool = False) -> None:
        self.data = data
        self.channel_reversed = channel_reversed
        self.shape: tuple[int, ...] | None = None

    def reshape(self, shape: tuple[int, ...]) -> _FakeArray:
        self.shape = shape
        return self

    def __getitem__(self, key: object) -> _FakeArray:
        assert key == (slice(None), slice(None), slice(None, None, -1))
        reversed_array = _FakeArray(self.data, channel_reversed=True)
        reversed_array.shape = self.shape
        return reversed_array


class _FakeNumpy:
    uint8 = object()

    def __init__(self) -> None:
        self.frombuffer_calls: list[tuple[bytes, object]] = []

    def frombuffer(self, data: bytes, *, dtype: object) -> _FakeArray:
        self.frombuffer_calls.append((data, dtype))
        return _FakeArray(data)


class _FakeTensor:
    """Exercise the detach/cpu/tolist path used by real torch tensors."""

    def __init__(self, value: list[Any]) -> None:
        self.value = value
        self.detached = False
        self.on_cpu = False

    def detach(self) -> _FakeTensor:
        self.detached = True
        return self

    def cpu(self) -> _FakeTensor:
        self.on_cpu = True
        return self

    def tolist(self) -> list[Any]:
        return self.value


class _FakeBoxes:
    def __init__(
        self,
        xyxy: list[list[float]],
        confidence: list[float],
        classes: list[float],
    ) -> None:
        self.xyxy = _FakeTensor(xyxy)
        self.conf = _FakeTensor(confidence)
        self.cls = _FakeTensor(classes)


def _weights(tmp_path: Path) -> Path:
    path = tmp_path / "best.pt"
    path.write_bytes(b"test-only-placeholder")
    return path


def _install_fake_runtime(
    monkeypatch: pytest.MonkeyPatch,
    *,
    names: object,
    results: list[object] | None = None,
) -> tuple[_FakeNumpy, list[Any]]:
    numpy = _FakeNumpy()
    created: list[Any] = []

    class FakeYOLO:
        def __init__(self, weights: str, *, task: str) -> None:
            self.weights = weights
            self.task = task
            self.names = names
            self.predict_calls: list[dict[str, Any]] = []
            created.append(self)

        def predict(self, **kwargs: Any) -> list[object]:
            self.predict_calls.append(kwargs)
            return [] if results is None else results

    ultralytics = SimpleNamespace(YOLO=FakeYOLO)

    def fake_import_module(name: str) -> object:
        if name == "numpy":
            return numpy
        if name == "ultralytics":
            return ultralytics
        raise AssertionError(f"unexpected lazy import: {name}")

    monkeypatch.setattr(yolo_module, "import_module", fake_import_module)
    return numpy, created


def _frame(
    *,
    pixel_format: str = "BGR8",
    codec: str = "none",
    data: bytes = bytes(range(12)),
) -> ImageFrame:
    return ImageFrame(
        data=data,
        width=2,
        height=2,
        pixel_format=pixel_format,
        codec=codec,
    )


def test_yolo_dependency_is_loaded_lazily(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    calls: list[str] = []
    fake_numpy = _FakeNumpy()

    def missing_ultralytics(name: str) -> object:
        calls.append(name)
        if name == "numpy":
            return fake_numpy
        raise ImportError("ultralytics is intentionally unavailable")

    monkeypatch.setattr(yolo_module, "import_module", missing_ultralytics)
    model = UltralyticsYoloModel(
        {"weights": str(_weights(tmp_path))},
        expected_labels=("No-Helmet",),
    )

    assert calls == []
    with pytest.raises(YoloDependencyError, match=r"cmvr-edge-ai\[yolo\]"):
        model.load()
    assert calls == ["numpy", "ultralytics"]


def test_yolo_load_requires_checkpoint_labels_to_match_registration_exactly(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    weights = _weights(tmp_path)
    _, created = _install_fake_runtime(
        monkeypatch,
        names={0: "Vest", 1: "No-Helmet"},
    )
    model = UltralyticsYoloModel(
        {"weights": str(weights)},
        expected_labels=("No-Helmet", "Vest"),
    )

    with pytest.raises(ValueError, match="checkpoint labels do not match"):
        model.load()
    assert len(created) == 1
    assert created[0].weights == str(weights)
    assert created[0].task == "detect"

    # A complete, ordered mapping is accepted; comparison is not set-based.
    _, accepted = _install_fake_runtime(
        monkeypatch,
        names={0: "No-Helmet", 1: "Vest"},
    )
    matching = UltralyticsYoloModel(
        {"weights": str(weights)},
        expected_labels=("No-Helmet", "Vest"),
    )
    matching.load()
    assert len(accepted) == 1


def test_ppe_6classes_builtin_registration_preserves_checkpoint_label_order() -> None:
    registry = DetectionModelRegistry()
    register_builtin_detection_models(registry)

    spec = registry.resolve("ppe-6classes-yolov8n@1")

    assert spec.name == "PPE Detection YOLOv8n (6 Classes)"
    assert spec.backend == "ultralytics-yolo"
    assert spec.supported_labels == (
        "Gloves",
        "Vest",
        "goggles",
        "helmet",
        "mask",
        "safety_shoe",
    )
    assert spec.supported_labels == PPE_6CLASS_LABELS


def test_ppe_6classes_factory_enforces_checkpoint_label_order(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    registry = DetectionModelRegistry()
    register_builtin_detection_models(registry)
    spec = registry.resolve("ppe-6classes-yolov8n@1")
    weights = _weights(tmp_path)
    model = spec.factory({"weights": str(weights)})
    _install_fake_runtime(
        monkeypatch,
        names={
            0: "Vest",
            1: "Gloves",
            2: "goggles",
            3: "helmet",
            4: "mask",
            5: "safety_shoe",
        },
    )

    assert isinstance(model, UltralyticsYoloModel)
    with pytest.raises(ValueError, match="checkpoint labels do not match"):
        model.load()

    _, created = _install_fake_runtime(
        monkeypatch,
        names=dict(enumerate(PPE_6CLASS_LABELS)),
    )
    matching = spec.factory({"weights": str(weights)})
    matching.load()

    assert isinstance(matching, UltralyticsYoloModel)
    assert len(created) == 1
    assert created[0].weights == str(weights)
    assert created[0].task == "detect"


def test_people_talking_builtin_registration_preserves_checkpoint_label_order() -> None:
    registry = DetectionModelRegistry()
    register_builtin_detection_models(registry)

    spec = registry.resolve("people-talking-yolov8x@1")

    assert spec.name == "People Talking YOLOv8x"
    assert spec.backend == "ultralytics-yolo"
    assert spec.supported_labels == ("label", "talking on phone")
    assert spec.supported_labels == PEOPLE_TALKING_LABELS


def test_people_talking_factory_enforces_generic_checkpoint_label(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    registry = DetectionModelRegistry()
    register_builtin_detection_models(registry)
    spec = registry.resolve("people-talking-yolov8x@1")
    weights = _weights(tmp_path)
    model = spec.factory({"weights": str(weights)})
    _install_fake_runtime(
        monkeypatch,
        names={0: "talking on phone"},
    )

    assert isinstance(model, UltralyticsYoloModel)
    with pytest.raises(ValueError, match="checkpoint labels do not match"):
        model.load()

    _, created = _install_fake_runtime(
        monkeypatch,
        names=dict(enumerate(PEOPLE_TALKING_LABELS)),
    )
    matching = spec.factory({"weights": str(weights)})
    matching.load()

    assert isinstance(matching, UltralyticsYoloModel)
    assert len(created) == 1
    assert created[0].weights == str(weights)
    assert created[0].task == "detect"


def test_mobile_phone_builtin_registration_preserves_checkpoint_label_order() -> None:
    registry = DetectionModelRegistry()
    register_builtin_detection_models(registry)

    spec = registry.resolve("yolov8n-mobile-phone@1")

    assert spec.name == "YOLOv8n Mobile Phone"
    assert spec.backend == "ultralytics-yolo"
    assert spec.supported_labels == ("mobile_phone",)
    assert spec.supported_labels == MOBILE_PHONE_LABELS


def test_mobile_phone_factory_enforces_checkpoint_label(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    registry = DetectionModelRegistry()
    register_builtin_detection_models(registry)
    spec = registry.resolve("yolov8n-mobile-phone@1")
    weights = _weights(tmp_path)
    model = spec.factory({"weights": str(weights)})
    _install_fake_runtime(monkeypatch, names={0: "cell_phone"})

    assert isinstance(model, UltralyticsYoloModel)
    with pytest.raises(ValueError, match="checkpoint labels do not match"):
        model.load()

    _, created = _install_fake_runtime(
        monkeypatch,
        names=dict(enumerate(MOBILE_PHONE_LABELS)),
    )
    matching = spec.factory({"weights": str(weights)})
    matching.load()

    assert isinstance(matching, UltralyticsYoloModel)
    assert len(created) == 1
    assert created[0].weights == str(weights)
    assert created[0].task == "detect"


@pytest.mark.parametrize(
    ("frame", "message"),
    [
        (_frame(codec="h264"), "decoded ImageFrame"),
        (_frame(pixel_format="GRAY8"), "supports BGR8 or RGB8"),
        (_frame(data=b"too-short"), r"width\*height\*3"),
    ],
)
def test_yolo_rejects_encoded_unsupported_or_malformed_buffers(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
    frame: ImageFrame,
    message: str,
) -> None:
    _install_fake_runtime(monkeypatch, names=("No-Helmet",))
    model = UltralyticsYoloModel(
        {"weights": str(_weights(tmp_path))},
        expected_labels=("No-Helmet",),
    )
    model.load()

    with pytest.raises(ValueError, match=message):
        model.predict(frame, ("No-Helmet",), 0.5)


def test_yolo_converts_rgb_to_bgr_but_keeps_bgr_buffer_order(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    numpy, created = _install_fake_runtime(monkeypatch, names=("No-Helmet",))
    model = UltralyticsYoloModel(
        {"weights": str(_weights(tmp_path))},
        expected_labels=("No-Helmet",),
    )
    model.load()

    model.predict(_frame(pixel_format=" BGR8 "), ("No-Helmet",), 0.5)
    bgr_source = created[0].predict_calls[-1]["source"]
    assert isinstance(bgr_source, _FakeArray)
    assert bgr_source.shape == (2, 2, 3)
    assert bgr_source.channel_reversed is False

    model.predict(_frame(pixel_format="RGB8"), ("No-Helmet",), 0.5)
    rgb_source = created[0].predict_calls[-1]["source"]
    assert isinstance(rgb_source, _FakeArray)
    assert rgb_source.shape == (2, 2, 3)
    assert rgb_source.channel_reversed is True
    assert numpy.frombuffer_calls == [
        (bytes(range(12)), numpy.uint8),
        (bytes(range(12)), numpy.uint8),
    ]


def test_yolo_predict_forwards_options_and_normalizes_boxes(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    boxes = _FakeBoxes(
        xyxy=[
            [1.0, 2.0, 30.0, 40.0],
            [5.5, 6.5, 50.5, 60.5],
            [0.0, 0.0, 1.0, 1.0],
        ],
        confidence=[0.91, 0.82, 0.99],
        classes=[1.0, 2.0, 0.0],
    )
    _, created = _install_fake_runtime(
        monkeypatch,
        names={0: "Worker", 1: "No-Helmet", 2: "No-Vest"},
        results=[SimpleNamespace(boxes=boxes)],
    )
    model = UltralyticsYoloModel(
        {
            "weights": str(_weights(tmp_path)),
            "device": "cuda:0",
            "imgsz": 320,
            "iou": 0.45,
            "half": True,
            "max_det": 42,
            "agnostic_nms": True,
        },
        expected_labels=("Worker", "No-Helmet", "No-Vest"),
    )
    model.load()

    detections = model.predict(
        _frame(),
        ("No-Helmet", "No-Vest"),
        0.73,
    )

    call = created[0].predict_calls[-1]
    assert call["classes"] == [1, 2]
    assert call["conf"] == pytest.approx(0.73)
    assert call["iou"] == pytest.approx(0.45)
    assert call["imgsz"] == 320
    assert call["device"] == "cuda:0"
    assert call["half"] is True
    assert call["max_det"] == 42
    assert call["agnostic_nms"] is True
    assert call["verbose"] is False
    assert detections == (
        Detection(
            label="No-Helmet",
            confidence=0.91,
            box=BoundingBox(1.0, 2.0, 30.0, 40.0),
        ),
        Detection(
            label="No-Vest",
            confidence=0.82,
            box=BoundingBox(5.5, 6.5, 50.5, 60.5),
        ),
    )
    assert boxes.xyxy.detached and boxes.xyxy.on_cpu
    assert boxes.conf.detached and boxes.conf.on_cpu
    assert boxes.cls.detached and boxes.cls.on_cpu


def test_yolo_close_is_idempotent_and_requires_reload_before_predict(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    _install_fake_runtime(monkeypatch, names=("No-Helmet",))
    model = UltralyticsYoloModel(
        {"weights": str(_weights(tmp_path))},
        expected_labels=("No-Helmet",),
    )
    model.load()
    model.close()
    model.close()

    with pytest.raises(RuntimeError, match="has not been loaded"):
        model.predict(_frame(), ("No-Helmet",), 0.5)