350 lines
12 KiB
Python
350 lines
12 KiB
Python
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)
|