443 lines
15 KiB
Python
443 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import threading
|
|
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Any, TypeVar
|
|
|
|
import pytest
|
|
|
|
from cmvr_edge_ai.contracts import BoundingBox, Detection, DetectionResult, ImageFrame
|
|
from cmvr_edge_ai.core.component import ComponentContext, Emission
|
|
from cmvr_edge_ai.core.envelope import Envelope
|
|
from cmvr_edge_ai.detection.base import DetectionModel
|
|
from cmvr_edge_ai.detection.operator import DetectionOperator
|
|
from cmvr_edge_ai.detection.registry import (
|
|
DetectionModelRegistry,
|
|
DetectionModelSpec,
|
|
UnknownDetectionModelError,
|
|
)
|
|
|
|
|
|
ResultT = TypeVar("ResultT")
|
|
|
|
|
|
class RecordingDetectionModel(DetectionModel):
|
|
def __init__(self, detections: Sequence[Detection] = ()) -> None:
|
|
self.detections = tuple(detections)
|
|
self.load_calls = 0
|
|
self.predict_calls = 0
|
|
self.close_calls = 0
|
|
self.predict_arguments: list[tuple[ImageFrame, tuple[str, ...], float]] = []
|
|
self.call_threads: dict[str, list[int]] = {
|
|
"load": [],
|
|
"predict": [],
|
|
"close": [],
|
|
}
|
|
|
|
def load(self) -> None:
|
|
self.load_calls += 1
|
|
self.call_threads["load"].append(threading.get_ident())
|
|
|
|
def predict(
|
|
self,
|
|
frame: ImageFrame,
|
|
labels: tuple[str, ...],
|
|
confidence: float,
|
|
) -> Sequence[Detection]:
|
|
self.predict_calls += 1
|
|
self.predict_arguments.append((frame, labels, confidence))
|
|
self.call_threads["predict"].append(threading.get_ident())
|
|
return self.detections
|
|
|
|
def close(self) -> None:
|
|
self.close_calls += 1
|
|
self.call_threads["close"].append(threading.get_ident())
|
|
|
|
|
|
def _detection(label: str, confidence: float) -> Detection:
|
|
return Detection(
|
|
label=label,
|
|
confidence=confidence,
|
|
box=BoundingBox(x_min=1.0, y_min=2.0, x_max=11.0, y_max=22.0),
|
|
)
|
|
|
|
|
|
def _frame(*, codec: str = "none") -> ImageFrame:
|
|
return ImageFrame(
|
|
data=bytes(range(12)),
|
|
width=2,
|
|
height=2,
|
|
pixel_format="BGR8",
|
|
codec=codec,
|
|
is_key_frame=codec != "none",
|
|
)
|
|
|
|
|
|
def _envelope(frame: ImageFrame | None = None) -> Envelope[ImageFrame]:
|
|
return Envelope(
|
|
payload=_frame() if frame is None else frame,
|
|
schema_name="ImageFrame",
|
|
schema_version=1,
|
|
source_id="camera-right",
|
|
sequence=42,
|
|
captured_at_ns=1_234_000_000,
|
|
received_at_ns=1_235_000_000,
|
|
deadline_ns=1_500_000_000,
|
|
trace_id="trace-detection-42",
|
|
session_id="camera-session",
|
|
attributes={"site": "park-a"},
|
|
)
|
|
|
|
|
|
def _registry(
|
|
model: RecordingDetectionModel,
|
|
*,
|
|
labels: tuple[str, ...] = ("Worker", "No-Helmet", "No-Vest"),
|
|
factory_options: list[dict[str, Any]] | None = None,
|
|
) -> DetectionModelRegistry:
|
|
def factory(options: Mapping[str, Any]) -> DetectionModel:
|
|
if factory_options is not None:
|
|
factory_options.append(dict(options))
|
|
return model
|
|
|
|
registry = DetectionModelRegistry()
|
|
registry.register(
|
|
DetectionModelSpec(
|
|
model_id="construction-ppe@1",
|
|
name="Construction PPE",
|
|
supported_labels=labels,
|
|
factory=factory,
|
|
backend="fake",
|
|
)
|
|
)
|
|
return registry
|
|
|
|
|
|
def _context(executor: ThreadPoolExecutor) -> ComponentContext:
|
|
return ComponentContext(
|
|
pipeline_id="detection-test",
|
|
node_id="detector",
|
|
shutdown_event=asyncio.Event(),
|
|
metadata={"thread_executor": executor},
|
|
)
|
|
|
|
|
|
async def _keep_restricted_event_loop_responsive() -> None:
|
|
"""Provide wakeups where sandboxed self-pipe notifications are unavailable."""
|
|
|
|
while True:
|
|
await asyncio.sleep(0.01)
|
|
|
|
|
|
def _run_with_thread_executor(
|
|
exercise: Callable[[ThreadPoolExecutor], Awaitable[ResultT]],
|
|
) -> ResultT:
|
|
async def runner() -> ResultT:
|
|
loop = asyncio.get_running_loop()
|
|
executor = ThreadPoolExecutor(max_workers=1)
|
|
loop.set_default_executor(executor)
|
|
ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
|
|
try:
|
|
return await exercise(executor)
|
|
finally:
|
|
ticker.cancel()
|
|
await asyncio.gather(ticker, return_exceptions=True)
|
|
executor.shutdown(wait=True, cancel_futures=True)
|
|
# asyncio.run() otherwise tries to shut down the same executor again.
|
|
loop._default_executor = None # type: ignore[attr-defined]
|
|
|
|
return asyncio.run(runner())
|
|
|
|
|
|
def test_operator_loads_once_predicts_and_closes_idempotently() -> None:
|
|
model = RecordingDetectionModel([_detection("Worker", 0.9)])
|
|
options_seen: list[dict[str, Any]] = []
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{
|
|
"model": "construction-ppe@1",
|
|
"model_options": {"weights": "/models/ppe.pt", "device": "cpu"},
|
|
},
|
|
model_registry=_registry(model, factory_options=options_seen),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
result = await operator.process(_envelope())
|
|
await operator.stop()
|
|
await operator.stop()
|
|
|
|
assert isinstance(result, Emission)
|
|
assert model.load_calls == 1
|
|
assert model.predict_calls == 1
|
|
assert model.close_calls == 1
|
|
assert options_seen == [{"weights": "/models/ppe.pt", "device": "cpu"}]
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_operator_filters_selected_labels_and_per_label_confidence() -> None:
|
|
selected_worker = _detection("Worker", 0.50)
|
|
selected_violation = _detection("No-Helmet", 0.80)
|
|
model = RecordingDetectionModel(
|
|
[
|
|
_detection("Worker", 0.49),
|
|
selected_worker,
|
|
_detection("No-Helmet", 0.79),
|
|
selected_violation,
|
|
_detection("No-Vest", 0.99),
|
|
]
|
|
)
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{
|
|
"model": "construction-ppe@1",
|
|
"detect_labels": ["Worker", "No-Helmet"],
|
|
"confidence": 0.50,
|
|
"label_confidence": {"No-Helmet": 0.80},
|
|
},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
emission = await operator.process(_envelope())
|
|
await operator.stop()
|
|
|
|
assert isinstance(emission, Emission)
|
|
payload = emission.envelope.payload
|
|
assert isinstance(payload, DetectionResult)
|
|
assert payload.detections == (selected_worker, selected_violation)
|
|
assert operator.selected_labels == ("Worker", "No-Helmet")
|
|
assert model.predict_arguments[0][1:] == (("Worker", "No-Helmet"), 0.50)
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_operator_rejects_unknown_model_during_construction() -> None:
|
|
with pytest.raises(UnknownDetectionModelError, match="missing@1"):
|
|
DetectionOperator(
|
|
"detector",
|
|
{"model": "missing@1"},
|
|
model_registry=DetectionModelRegistry(),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"params",
|
|
[
|
|
{
|
|
"model": "construction-ppe@1",
|
|
"detect_labels": ["Unknown-Label"],
|
|
},
|
|
{
|
|
"model": "construction-ppe@1",
|
|
"label_confidence": {"Unknown-Label": 0.9},
|
|
},
|
|
],
|
|
)
|
|
def test_operator_rejects_unknown_labels_during_construction(
|
|
params: dict[str, Any],
|
|
) -> None:
|
|
with pytest.raises(ValueError, match="Unknown-Label"):
|
|
DetectionOperator(
|
|
"detector",
|
|
params,
|
|
model_registry=_registry(RecordingDetectionModel()),
|
|
)
|
|
|
|
|
|
def test_operator_rejects_encoded_frame_without_calling_model() -> None:
|
|
model = RecordingDetectionModel()
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{"model": "construction-ppe@1"},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
with pytest.raises(ValueError, match="decoded ImageFrame"):
|
|
await operator.process(_envelope(_frame(codec="h264")))
|
|
await operator.stop()
|
|
|
|
assert model.predict_calls == 0
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_operator_emits_result_metadata_and_preserves_envelope_correlation() -> None:
|
|
source = _envelope()
|
|
expected_detection = _detection("No-Helmet", 0.93)
|
|
model = RecordingDetectionModel([expected_detection])
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{"model": "construction-ppe@1"},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
emission = await operator.process(source)
|
|
await operator.stop()
|
|
|
|
assert isinstance(emission, Emission)
|
|
assert emission.port == "detections"
|
|
output = emission.envelope
|
|
assert output.schema == "DetectionResult/v1"
|
|
assert isinstance(output.payload, DetectionResult)
|
|
assert output.payload.detections == (expected_detection,)
|
|
assert output.payload.model_id == "construction-ppe@1"
|
|
assert output.payload.model_name == "Construction PPE"
|
|
assert output.payload.inference_ms >= 0.0
|
|
assert output.payload.source_frame is None
|
|
assert operator.attach_frame is False
|
|
assert output.source_id == source.source_id
|
|
assert output.sequence == source.sequence
|
|
assert output.captured_at_ns == source.captured_at_ns
|
|
assert output.received_at_ns == source.received_at_ns
|
|
assert output.deadline_ns == source.deadline_ns
|
|
assert output.trace_id == source.trace_id
|
|
assert output.session_id == source.session_id
|
|
assert output.attributes == {
|
|
"site": "park-a",
|
|
"detection_model_id": "construction-ppe@1",
|
|
"detection_model_name": "Construction PPE",
|
|
}
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_operator_attach_frame_keeps_the_decoded_frame_by_reference() -> None:
|
|
source = _envelope()
|
|
model = RecordingDetectionModel([_detection("No-Helmet", 0.93)])
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{"model": "construction-ppe@1", "attach_frame": True},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
emission = await operator.process(source)
|
|
await operator.stop()
|
|
|
|
assert isinstance(emission, Emission)
|
|
assert isinstance(emission.envelope.payload, DetectionResult)
|
|
assert emission.envelope.payload.source_frame is source.payload
|
|
assert operator.attach_frame is True
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_blocking_model_methods_run_on_the_supplied_thread_pool() -> None:
|
|
event_loop_thread = threading.get_ident()
|
|
model = RecordingDetectionModel()
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{"model": "construction-ppe@1"},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
await operator.process(_envelope())
|
|
await operator.stop()
|
|
|
|
assert all(len(thread_ids) == 1 for thread_ids in model.call_threads.values())
|
|
worker_threads = {
|
|
thread_id
|
|
for thread_ids in model.call_threads.values()
|
|
for thread_id in thread_ids
|
|
}
|
|
assert len(worker_threads) == 1
|
|
assert event_loop_thread not in worker_threads
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_max_fps_skips_frames_without_invoking_model_again() -> None:
|
|
model = RecordingDetectionModel()
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{"model": "construction-ppe@1", "max_fps": 0.001},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
first = await operator.process(_envelope())
|
|
second = await operator.process(_envelope())
|
|
await operator.stop()
|
|
|
|
assert isinstance(first, Emission)
|
|
assert second is None
|
|
assert model.predict_calls == 1
|
|
|
|
_run_with_thread_executor(scenario)
|
|
|
|
|
|
def test_operator_logs_model_load_and_configured_inference_heartbeat(
|
|
caplog: pytest.LogCaptureFixture,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
timestamps = iter((0, 1_000_000_000, 6_000_000_000))
|
|
monkeypatch.setattr(
|
|
"cmvr_edge_ai.detection.operator.monotonic_ns",
|
|
lambda: next(timestamps),
|
|
)
|
|
model = RecordingDetectionModel([_detection("No-Helmet", 0.93)])
|
|
operator = DetectionOperator(
|
|
"detector",
|
|
{
|
|
"model": "construction-ppe@1",
|
|
"detect_labels": ["No-Helmet"],
|
|
"inference_log_interval_s": 5,
|
|
},
|
|
model_registry=_registry(model),
|
|
)
|
|
|
|
async def scenario(executor: ThreadPoolExecutor) -> None:
|
|
await operator.setup(_context(executor))
|
|
await operator.process(_envelope())
|
|
await operator.process(_envelope())
|
|
await operator.process(_envelope())
|
|
await operator.stop()
|
|
|
|
with caplog.at_level(logging.INFO, logger="cmvr_edge_ai.detection.operator"):
|
|
_run_with_thread_executor(scenario)
|
|
|
|
messages = [record.getMessage() for record in caplog.records]
|
|
assert any(
|
|
"detection model loaded node=detector model=construction-ppe@1" in message
|
|
for message in messages
|
|
)
|
|
inference_messages = [
|
|
message for message in messages if message.startswith("detection inference ")
|
|
]
|
|
assert len(inference_messages) == 2
|
|
assert "total_frames=1" in inference_messages[0]
|
|
assert "window_frames=1" in inference_messages[0]
|
|
assert "window_detections=1" in inference_messages[0]
|
|
assert "hit_labels=No-Helmet:1" in inference_messages[0]
|
|
assert "total_frames=3" in inference_messages[1]
|
|
assert "window_frames=2" in inference_messages[1]
|
|
assert "window_detections=2" in inference_messages[1]
|
|
assert "hit_labels=No-Helmet:2" in inference_messages[1]
|
|
|
|
|
|
@pytest.mark.parametrize("value", [0, -1, 3600.1, True])
|
|
def test_operator_rejects_invalid_inference_log_interval(value: Any) -> None:
|
|
with pytest.raises(ValueError, match="inference_log_interval_s"):
|
|
DetectionOperator(
|
|
"detector",
|
|
{
|
|
"model": "construction-ppe@1",
|
|
"inference_log_interval_s": value,
|
|
},
|
|
model_registry=_registry(RecordingDetectionModel()),
|
|
)
|