from __future__ import annotations import asyncio from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from typing import Any from cmvr_edge_ai.application import ( create_default_model_registry, create_default_registry, ) from cmvr_edge_ai.compiler import compile_pipeline from cmvr_edge_ai.config import load_config_data from cmvr_edge_ai.contracts import ( BoundingBox, Detection, DetectionAlert, DetectionResult, ImageFrame, ) from cmvr_edge_ai.core import Envelope, Sink from cmvr_edge_ai.detection import ( DetectionModel, DetectionModelRegistry, DetectionModelSpec, ) from cmvr_edge_ai.plugins import PluginKind, PluginSpec SIX_CLASS_MODEL_ID = "ppe-6classes-yolov8n@1" SIX_CLASS_LABELS = ( "Gloves", "Vest", "goggles", "helmet", "mask", "safety_shoe", ) class _FixedDetectionModel(DetectionModel): def __init__(self, label: str) -> None: self.label = label self.load_calls = 0 self.close_calls = 0 self.frames: list[ImageFrame] = [] self.selected_labels: list[tuple[str, ...]] = [] def load(self) -> None: self.load_calls += 1 def predict( self, frame: ImageFrame, labels: tuple[str, ...], confidence: float, ) -> Sequence[Detection]: del confidence self.frames.append(frame) self.selected_labels.append(labels) return ( Detection( label=self.label, confidence=0.95, box=BoundingBox(1.0, 2.0, 20.0, 30.0), ), ) def close(self) -> None: self.close_calls += 1 class _CollectSink(Sink): def __init__(self) -> None: self.envelopes: list[Envelope[Any]] = [] self.input_ports: list[str] = [] async def consume( self, envelope: Envelope[Any], input_port: str = "input", ) -> None: self.envelopes.append(envelope) self.input_ports.append(input_port) async def _keep_event_loop_responsive() -> None: while True: await asyncio.sleep(0.01) def test_default_registry_contains_six_class_ppe_model_contract() -> None: registry = create_default_model_registry(discover_entry_points=False) spec = registry.resolve(SIX_CLASS_MODEL_ID) assert spec.model_id == SIX_CLASS_MODEL_ID assert spec.supported_labels == SIX_CLASS_LABELS def test_decoder_fans_out_to_alert_and_raw_detection_branches() -> None: ppe_model = _FixedDetectionModel("No-Helmet") six_class_model = _FixedDetectionModel("helmet") models = DetectionModelRegistry() models.register( DetectionModelSpec( model_id="construction-ppe-yolov8@1", name="Test Construction PPE", supported_labels=("No-Helmet",), factory=lambda options: ppe_model, backend="fake", ) ) models.register( DetectionModelSpec( model_id=SIX_CLASS_MODEL_ID, name="Test Six-Class PPE", supported_labels=SIX_CLASS_LABELS, factory=lambda options: six_class_model, backend="fake", ) ) registry = create_default_registry( discover_entry_points=False, model_registry=models, ) alert_sink = _CollectSink() raw_sink = _CollectSink() registry.register( PluginSpec( plugin_id="test.dual_alert_sink@1", kind=PluginKind.SINK, factory=lambda node_id, params: alert_sink, inputs={"alerts": "DetectionAlert/v1"}, ) ) registry.register( PluginSpec( plugin_id="test.dual_raw_detection_sink@1", kind=PluginKind.SINK, factory=lambda node_id, params: raw_sink, inputs={"detections": "DetectionResult/v1"}, ) ) width, height = 4, 3 frames = ( ImageFrame(bytes((0, 0, 255)) * (width * height), width, height, "BGR8"), ImageFrame(bytes((0, 255, 0)) * (width * height), width, height, "BGR8"), ) config = load_config_data( { "api_version": "cmvr.edge.ai/v1", "pipelines": { "dual_detection": { "nodes": { "source": { "uses": "core.sequence_source@1", "with": { "items": list(frames), "schema_name": "ImageFrame", }, }, "decoder": {"uses": "media.video_decoder.pyav@1"}, "ppe_detector": { "uses": "detection.model@1", "with": { "model": "construction-ppe-yolov8@1", "detect_labels": ["No-Helmet"], }, }, "ppe_gate": { "uses": "detection.repeat_gate@1", "with": { "time_source": "received", "rules": [ { "id": "no-helmet", "labels": ["No-Helmet"], "min_hits": 2, "window_ms": 60_000, "cooldown_ms": 60_000, "scope": "source", } ], }, }, "alert_sink": {"uses": "test.dual_alert_sink@1"}, "six_class_detector": { "uses": "detection.model@1", "with": { "model": SIX_CLASS_MODEL_ID, "detect_labels": list(SIX_CLASS_LABELS), "attach_frame": False, }, }, "raw_sink": { "uses": "test.dual_raw_detection_sink@1" }, }, "edges": [ { "from": "source.output", "to": "decoder.frames", "qos": { "profile": "video_contiguous", "capacity": 4, "overflow": "block", }, }, { "from": "decoder.frames", "to": "ppe_detector.frames", "qos": { "profile": "telemetry", "capacity": 4, "overflow": "block", }, }, { "from": "decoder.frames", "to": "six_class_detector.frames", "qos": { "profile": "telemetry", "capacity": 4, "overflow": "block", }, }, { "from": "ppe_detector.detections", "to": "ppe_gate.detections", "qos": { "profile": "telemetry", "capacity": 4, "overflow": "block", }, }, { "from": "ppe_gate.alerts", "to": "alert_sink.alerts", "qos": { "profile": "telemetry", "capacity": 4, "overflow": "block", }, }, { "from": "six_class_detector.detections", "to": "raw_sink.detections", "qos": { "profile": "telemetry", "capacity": 4, "overflow": "block", }, }, ], } }, } ) async def scenario(): # type: ignore[no-untyped-def] loop = asyncio.get_running_loop() executor = ThreadPoolExecutor(max_workers=2) loop.set_default_executor(executor) ticker = asyncio.create_task(_keep_event_loop_responsive()) try: compiled = compile_pipeline( config, "dual_detection", registry, metadata={"thread_executor": executor}, ) await compiled.runtime.run() return compiled finally: ticker.cancel() await asyncio.gather(ticker, return_exceptions=True) executor.shutdown(wait=True, cancel_futures=True) loop._default_executor = None # type: ignore[attr-defined] compiled = asyncio.run(scenario()) assert compiled.plugin_specs["ppe_gate"].outputs == { "alerts": "DetectionAlert/v1" } assert compiled.plugin_specs["six_class_detector"].outputs == { "detections": "DetectionResult/v1" } assert ppe_model.load_calls == ppe_model.close_calls == 1 assert six_class_model.load_calls == six_class_model.close_calls == 1 assert len(ppe_model.frames) == len(six_class_model.frames) == len(frames) assert all( ppe_frame is six_class_frame for ppe_frame, six_class_frame in zip( ppe_model.frames, six_class_model.frames, ) ) assert ppe_model.selected_labels == [("No-Helmet",)] * len(frames) assert six_class_model.selected_labels == [SIX_CLASS_LABELS] * len(frames) assert len(alert_sink.envelopes) == 1 alert_envelope = alert_sink.envelopes[0] assert alert_sink.input_ports == ["alerts"] assert alert_envelope.schema == "DetectionAlert/v1" assert alert_envelope.source_id == "source" assert alert_envelope.sequence == 1 assert isinstance(alert_envelope.payload, DetectionAlert) assert alert_envelope.payload.model_id == "construction-ppe-yolov8@1" assert alert_envelope.payload.rule_id == "no-helmet" assert raw_sink.input_ports == ["detections", "detections"] assert [envelope.schema for envelope in raw_sink.envelopes] == [ "DetectionResult/v1", "DetectionResult/v1", ] assert [envelope.source_id for envelope in raw_sink.envelopes] == [ "source", "source", ] assert [envelope.sequence for envelope in raw_sink.envelopes] == [0, 1] assert all( isinstance(envelope.payload, DetectionResult) for envelope in raw_sink.envelopes ) assert all( envelope.payload.model_id == SIX_CLASS_MODEL_ID for envelope in raw_sink.envelopes ) assert all( envelope.payload.source_frame is None for envelope in raw_sink.envelopes ) assert alert_envelope.trace_id == raw_sink.envelopes[1].trace_id stats = compiled.runtime.edge_stats() ppe_fanout = next( value for name, value in stats.items() if name.endswith("decoder.frames->ppe_detector.frames") ) six_class_fanout = next( value for name, value in stats.items() if name.endswith("decoder.frames->six_class_detector.frames") ) assert ppe_fanout.enqueued == len(frames) assert six_class_fanout.enqueued == len(frames)