314 lines
9.8 KiB
Python
314 lines
9.8 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from cmvr_edge_ai.application import create_default_model_registry
|
|
from cmvr_edge_ai.config import load_config
|
|
from cmvr_edge_ai.contracts import ImageFrame
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
CONSTRUCTION_ONNX = (
|
|
"models/detection/construction-ppe-yolov8/v2/model.onnx"
|
|
)
|
|
PEOPLE_TALKING_ONNX = (
|
|
"models/detection/people-talking-yolov8x/v2/model.onnx"
|
|
)
|
|
MOBILE_PHONE_ONNX = (
|
|
"models/detection/yolov8n-mobile-phone/v2/model.onnx"
|
|
)
|
|
|
|
SOURCE_ARTIFACTS = (
|
|
(
|
|
"construction-ppe-yolov8@1",
|
|
"models/detection/construction-ppe-yolov8/v1/best.pt",
|
|
22_537_898,
|
|
"31ef3ca04a17cf545f3fcfc64c4af8993a41d52ccc460e82aff01d5354603533",
|
|
),
|
|
(
|
|
"ppe-6classes-yolov8n@1",
|
|
"models/detection/ppe-6classes-yolov8n/v1/best.pt",
|
|
5_625_014,
|
|
"07172ef3ae9e256c40a1fb0ce3eefe5547d90170645aa73dded0fffc382cdb31",
|
|
),
|
|
(
|
|
"people-talking-yolov8x@1",
|
|
"models/detection/people-talking-yolov8x/v1/best.pt",
|
|
136_694_313,
|
|
"86cd63926de1c69f70dd7f8755e4de9e382bbdeb0668821b7c3846b826471526",
|
|
),
|
|
(
|
|
"yolov8n-mobile-phone@1",
|
|
"models/detection/yolov8n-mobile-phone/yolov8n-mobile-phone.pt",
|
|
6_234_666,
|
|
"9230e4bfa7cba7134215c4c7f228b5e58760b02138788ccaff0149258c2d2e19",
|
|
),
|
|
)
|
|
|
|
ONNX_ARTIFACTS = (
|
|
(
|
|
"construction-ppe-yolov8@1",
|
|
"construction-ppe-yolov8@2",
|
|
CONSTRUCTION_ONNX,
|
|
44_774_179,
|
|
"b8e5d116a964d0e7091e14b5382ab9301c196ae33b363516df04fddfb0d0b57a",
|
|
),
|
|
(
|
|
"ppe-6classes-yolov8n@1",
|
|
"ppe-6classes-yolov8n@2",
|
|
"models/detection/ppe-6classes-yolov8n/v2/model.onnx",
|
|
10_993_196,
|
|
"6cb0e567b4fbf353aeb61c271ab33c45260d6ddddbe26eb3b64c245c2683c03b",
|
|
),
|
|
(
|
|
"people-talking-yolov8x@1",
|
|
"people-talking-yolov8x@2",
|
|
PEOPLE_TALKING_ONNX,
|
|
272_787_930,
|
|
"b4387307dbdf1083c0229a12e833ccdc4a3ad0063d5bdb7908765072470e840b",
|
|
),
|
|
(
|
|
"yolov8n-mobile-phone@1",
|
|
"yolov8n-mobile-phone@2",
|
|
MOBILE_PHONE_ONNX,
|
|
12_265_264,
|
|
"df170d9bc86984894797da6f2101519617669c3f589d75289acb0ba1937f857f",
|
|
),
|
|
)
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as artifact:
|
|
for chunk in iter(lambda: artifact.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_id", "relative_path", "expected_size", "expected_sha256"),
|
|
SOURCE_ARTIFACTS,
|
|
)
|
|
def test_pt_export_source_is_complete(
|
|
model_id: str,
|
|
relative_path: str,
|
|
expected_size: int,
|
|
expected_sha256: str,
|
|
) -> None:
|
|
weights = PROJECT_ROOT / relative_path
|
|
|
|
assert model_id.endswith("@1")
|
|
assert (weights.parent / "README.md").is_file()
|
|
assert weights.is_file()
|
|
assert weights.stat().st_size == expected_size
|
|
assert _sha256(weights) == expected_sha256
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"source_model_id",
|
|
"model_id",
|
|
"relative_path",
|
|
"expected_size",
|
|
"expected_sha256",
|
|
),
|
|
ONNX_ARTIFACTS,
|
|
)
|
|
def test_onnx_artifact_and_manifest_are_complete(
|
|
source_model_id: str,
|
|
model_id: str,
|
|
relative_path: str,
|
|
expected_size: int,
|
|
expected_sha256: str,
|
|
) -> None:
|
|
weights = PROJECT_ROOT / relative_path
|
|
manifest_path = weights.with_name("manifest.json")
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
registry = create_default_model_registry(discover_entry_points=False)
|
|
spec = registry.resolve(model_id)
|
|
|
|
assert (weights.parent / "README.md").is_file()
|
|
assert weights.is_file()
|
|
assert weights.stat().st_size == expected_size
|
|
assert _sha256(weights) == expected_sha256
|
|
assert manifest["schema_version"] == "cmvr.detection-model-manifest/v1"
|
|
assert manifest["source_model_id"] == source_model_id
|
|
assert manifest["model_id"] == model_id
|
|
assert manifest["backend"] == "onnxruntime-yolov8"
|
|
assert manifest["artifact"] == relative_path
|
|
assert manifest["artifact_sha256"] == expected_sha256
|
|
assert manifest["labels"] == list(spec.supported_labels)
|
|
assert manifest["input"]["shape"] == [1, 3, 640, 640]
|
|
assert manifest["export"] == {
|
|
"batch": 1,
|
|
"dynamic": False,
|
|
"format": "onnx",
|
|
"half": False,
|
|
"nms": False,
|
|
"opset": 17,
|
|
"simplify": True,
|
|
}
|
|
source = PROJECT_ROOT / manifest["source"]
|
|
assert source.is_file()
|
|
assert manifest["source_sha256"] == _sha256(source)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_id", "relative_path"),
|
|
[(item[1], item[2]) for item in ONNX_ARTIFACTS],
|
|
)
|
|
def test_onnx_metadata_has_stable_identity_without_training_host_paths(
|
|
model_id: str,
|
|
relative_path: str,
|
|
) -> None:
|
|
onnx = pytest.importorskip("onnx")
|
|
weights = PROJECT_ROOT / relative_path
|
|
model = onnx.load(str(weights), load_external_data=False)
|
|
metadata = {item.key: item.value for item in model.metadata_props}
|
|
manifest = json.loads(
|
|
weights.with_name("manifest.json").read_text(encoding="utf-8")
|
|
)
|
|
|
|
assert "date" not in metadata
|
|
assert "data.yaml" not in metadata["description"]
|
|
assert metadata["cmvr_model_id"] == model_id
|
|
assert metadata["cmvr_source_sha256"] == manifest["source_sha256"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_id",
|
|
[item[1] for item in ONNX_ARTIFACTS],
|
|
)
|
|
def test_real_onnx_artifact_loads_and_runs_on_cpu(model_id: str) -> None:
|
|
pytest.importorskip("onnxruntime")
|
|
registry = create_default_model_registry(discover_entry_points=False)
|
|
spec = registry.resolve(model_id)
|
|
model = spec.factory(
|
|
{
|
|
"providers": ["CPUExecutionProvider"],
|
|
"intra_op_threads": 1,
|
|
"inter_op_threads": 1,
|
|
}
|
|
)
|
|
frame = ImageFrame(
|
|
data=bytes(32 * 32 * 3),
|
|
width=32,
|
|
height=32,
|
|
pixel_format="BGR8",
|
|
)
|
|
|
|
try:
|
|
model.load()
|
|
detections = tuple(
|
|
model.predict(frame, spec.supported_labels, confidence=0.99)
|
|
)
|
|
finally:
|
|
model.close()
|
|
|
|
for detection in detections:
|
|
assert detection.label in spec.supported_labels
|
|
assert 0.0 <= detection.box.x1 <= detection.box.x2 <= frame.width
|
|
assert 0.0 <= detection.box.y1 <= detection.box.y2 <= frame.height
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("config_path", "pipeline_id"),
|
|
(
|
|
(PROJECT_ROOT / "configs" / "active_detection.yaml", "detection"),
|
|
(
|
|
PROJECT_ROOT / "configs" / "detection_viewer.yaml",
|
|
"detection_show",
|
|
),
|
|
(PROJECT_ROOT / "configs" / "server_detect.yaml", "detect_ppe"),
|
|
),
|
|
)
|
|
def test_detection_configs_use_repository_construction_onnx(
|
|
config_path: Path,
|
|
pipeline_id: str,
|
|
) -> None:
|
|
config = load_config(config_path)
|
|
detector = config.pipelines[pipeline_id].nodes["detector"]
|
|
|
|
assert detector.params["model"] == "construction-ppe-yolov8@2"
|
|
assert detector.params["model_options"]["weights"] == CONSTRUCTION_ONNX
|
|
assert detector.params["model_options"]["providers"] == [
|
|
"CPUExecutionProvider"
|
|
]
|
|
assert "device" not in detector.params["model_options"]
|
|
assert "half" not in detector.params["model_options"]
|
|
|
|
|
|
def test_production_pipeline_uses_repository_people_talking_onnx() -> None:
|
|
config = load_config(PROJECT_ROOT / "configs" / "active_detection.yaml")
|
|
pipeline = config.pipelines["detection"]
|
|
detector = pipeline.nodes["phone_detector"]
|
|
gate = pipeline.nodes["phone_repeat_gate"]
|
|
|
|
assert detector.params["model"] == "people-talking-yolov8x@2"
|
|
assert detector.params["detect_labels"] == ["talking on phone"]
|
|
assert detector.params["model_options"]["weights"] == PEOPLE_TALKING_ONNX
|
|
assert gate.params["rules"][0]["labels"] == ["talking on phone"]
|
|
assert any(
|
|
edge.source == "decoder.frames" and edge.target == "phone_detector.frames"
|
|
for edge in pipeline.edges
|
|
)
|
|
assert any(
|
|
edge.source == "phone_detector.detections"
|
|
and edge.target == "phone_repeat_gate.detections"
|
|
for edge in pipeline.edges
|
|
)
|
|
assert any(
|
|
edge.source == "phone_repeat_gate.alerts"
|
|
and edge.target == "alert_platform.input"
|
|
for edge in pipeline.edges
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("config_name", "pipeline_id"),
|
|
(
|
|
("detection_viewer.yaml", "detection_show"),
|
|
("server_detect.yaml", "detect_mobile_phone"),
|
|
),
|
|
)
|
|
def test_mobile_phone_configs_use_repository_onnx(
|
|
config_name: str,
|
|
pipeline_id: str,
|
|
) -> None:
|
|
config = load_config(PROJECT_ROOT / "configs" / config_name)
|
|
detector = config.pipelines[pipeline_id].nodes["phone_detector" if pipeline_id == "detection_show" else "detector"]
|
|
|
|
assert detector.params["model"] == "yolov8n-mobile-phone@2"
|
|
assert detector.params["detect_labels"] == ["mobile_phone"]
|
|
assert detector.params["model_options"]["weights"] == MOBILE_PHONE_ONNX
|
|
|
|
|
|
def test_debug_viewer_preserves_same_frame_pairing() -> None:
|
|
config = load_config(PROJECT_ROOT / "configs" / "detection_viewer.yaml")
|
|
pipeline = config.pipelines["detection_show"]
|
|
ppe_detector = pipeline.nodes["detector"]
|
|
detector = pipeline.nodes["phone_detector"]
|
|
viewer = pipeline.nodes["viewer"]
|
|
|
|
assert ppe_detector.params["attach_frame"] is True
|
|
assert detector.params["attach_frame"] is True
|
|
assert viewer.params["expected_inputs"] == ["ppe", "phone"]
|
|
assert any(
|
|
edge.source == "frame_fanout.output"
|
|
and edge.target == "phone_detector.frames"
|
|
for edge in pipeline.edges
|
|
)
|
|
assert any(
|
|
edge.source == "phone_detector.detections"
|
|
and edge.target == "viewer.phone"
|
|
for edge in pipeline.edges
|
|
)
|
|
assert any(
|
|
edge.source == "detector.detections" and edge.target == "viewer.ppe"
|
|
for edge in pipeline.edges
|
|
)
|