106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from cmvr_edge_ai.application import create_default_registry
|
||
|
|
from cmvr_edge_ai.compiler import PipelineCompileError, compile_pipeline
|
||
|
|
from cmvr_edge_ai.config import load_config_data
|
||
|
|
|
||
|
|
|
||
|
|
def _config(
|
||
|
|
rule_label: str,
|
||
|
|
*,
|
||
|
|
attach_frame: bool = False,
|
||
|
|
alert_image: bool = False,
|
||
|
|
): # type: ignore[no-untyped-def]
|
||
|
|
return load_config_data(
|
||
|
|
{
|
||
|
|
"api_version": "cmvr.edge.ai/v1",
|
||
|
|
"pipelines": {
|
||
|
|
"detection": {
|
||
|
|
"nodes": {
|
||
|
|
"source": {
|
||
|
|
"uses": "core.sequence_source@1",
|
||
|
|
"with": {
|
||
|
|
"items": [],
|
||
|
|
"schema_name": "ImageFrame",
|
||
|
|
},
|
||
|
|
},
|
||
|
|
"detector": {
|
||
|
|
"uses": "detection.model@1",
|
||
|
|
"with": {
|
||
|
|
"model": "construction-ppe-yolov8@1",
|
||
|
|
"detect_labels": ["No-Helmet"],
|
||
|
|
"attach_frame": attach_frame,
|
||
|
|
"model_options": {"weights": "/not-loaded.pt"},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
"gate": {
|
||
|
|
"uses": "detection.repeat_gate@1",
|
||
|
|
"with": {
|
||
|
|
"alert_image": {"enabled": alert_image},
|
||
|
|
"rules": [
|
||
|
|
{
|
||
|
|
"id": "ppe-rule",
|
||
|
|
"labels": [rule_label],
|
||
|
|
"min_hits": 2,
|
||
|
|
"window_ms": 1000,
|
||
|
|
}
|
||
|
|
]
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
"edges": [
|
||
|
|
{"from": "source.output", "to": "detector.frames"},
|
||
|
|
{
|
||
|
|
"from": "detector.detections",
|
||
|
|
"to": "gate.detections",
|
||
|
|
},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
},
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_compile_accepts_rule_label_selected_by_detector() -> None:
|
||
|
|
compiled = compile_pipeline(
|
||
|
|
_config("No-Helmet"),
|
||
|
|
"detection",
|
||
|
|
create_default_registry(discover_entry_points=False),
|
||
|
|
)
|
||
|
|
|
||
|
|
assert set(compiled.plugin_specs) == {"source", "detector", "gate"}
|
||
|
|
|
||
|
|
|
||
|
|
def test_compile_rejects_rule_label_ignored_by_detector() -> None:
|
||
|
|
with pytest.raises(PipelineCompileError, match="No-Vest"):
|
||
|
|
compile_pipeline(
|
||
|
|
_config("No-Vest"),
|
||
|
|
"detection",
|
||
|
|
create_default_registry(discover_entry_points=False),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_compile_accepts_alert_images_when_detector_attaches_frames() -> None:
|
||
|
|
compiled = compile_pipeline(
|
||
|
|
_config("No-Helmet", attach_frame=True, alert_image=True),
|
||
|
|
"detection",
|
||
|
|
create_default_registry(discover_entry_points=False),
|
||
|
|
)
|
||
|
|
|
||
|
|
assert compiled.runtime.nodes["detector"].attach_frame is True
|
||
|
|
assert compiled.runtime.nodes["gate"].alert_image_enabled is True
|
||
|
|
|
||
|
|
|
||
|
|
def test_compile_rejects_alert_images_without_attached_frames() -> None:
|
||
|
|
with pytest.raises(
|
||
|
|
PipelineCompileError,
|
||
|
|
match="enables alert_image.*attach_frame=true",
|
||
|
|
):
|
||
|
|
compile_pipeline(
|
||
|
|
_config("No-Helmet", attach_frame=False, alert_image=True),
|
||
|
|
"detection",
|
||
|
|
create_default_registry(discover_entry_points=False),
|
||
|
|
)
|