379 lines
11 KiB
Python
379 lines
11 KiB
Python
|
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 (
|
||
|
|
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"
|
||
|
|
|
||
|
|
|
||
|
|
@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)
|