152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from cmvr_edge_ai.config import load_config
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
CONSTRUCTION_WEIGHTS = "models/detection/construction-ppe-yolov8/v1/best.pt"
|
|
PEOPLE_TALKING_WEIGHTS = "models/detection/people-talking-yolov8x/v1/best.pt"
|
|
MOBILE_PHONE_WEIGHTS = (
|
|
"models/detection/yolov8n-mobile-phone/yolov8n-mobile-phone.pt"
|
|
)
|
|
MODEL_ARTIFACTS = (
|
|
(
|
|
"construction-ppe-yolov8",
|
|
22_537_898,
|
|
"31ef3ca04a17cf545f3fcfc64c4af8993a41d52ccc460e82aff01d5354603533",
|
|
),
|
|
(
|
|
"ppe-6classes-yolov8n",
|
|
5_625_014,
|
|
"07172ef3ae9e256c40a1fb0ce3eefe5547d90170645aa73dded0fffc382cdb31",
|
|
),
|
|
(
|
|
"people-talking-yolov8x",
|
|
136_694_313,
|
|
"86cd63926de1c69f70dd7f8755e4de9e382bbdeb0668821b7c3846b826471526",
|
|
),
|
|
)
|
|
|
|
|
|
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_name", "expected_size", "expected_sha256"),
|
|
MODEL_ARTIFACTS,
|
|
)
|
|
def test_model_artifact_is_complete(
|
|
model_name: str,
|
|
expected_size: int,
|
|
expected_sha256: str,
|
|
) -> None:
|
|
model_dir = PROJECT_ROOT / "models" / "detection" / model_name / "v1"
|
|
weights = model_dir / "best.pt"
|
|
|
|
assert (model_dir / "README.md").is_file()
|
|
assert weights.is_file()
|
|
assert weights.stat().st_size == expected_size
|
|
assert _sha256(weights) == expected_sha256
|
|
|
|
|
|
def test_mobile_phone_model_artifact_is_complete() -> None:
|
|
weights = PROJECT_ROOT / MOBILE_PHONE_WEIGHTS
|
|
model_card = weights.parent / "README.md"
|
|
|
|
assert model_card.is_file()
|
|
assert weights.is_file()
|
|
assert weights.stat().st_size == 6_234_666
|
|
assert (
|
|
_sha256(weights)
|
|
== "9230e4bfa7cba7134215c4c7f228b5e58760b02138788ccaff0149258c2d2e19"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("config_path", "pipeline_id"),
|
|
(
|
|
(PROJECT_ROOT / "configs" / "edge_ai.yaml", "detection"),
|
|
(
|
|
PROJECT_ROOT / "configs" / "debug" / "detection_viewer.yaml",
|
|
"detection_show",
|
|
),
|
|
),
|
|
)
|
|
def test_detection_configs_use_repository_construction_weights(
|
|
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@1"
|
|
assert detector.params["model_options"]["weights"] == CONSTRUCTION_WEIGHTS
|
|
assert (PROJECT_ROOT / CONSTRUCTION_WEIGHTS).is_file()
|
|
|
|
|
|
def test_production_pipeline_uses_repository_people_talking_weights() -> None:
|
|
config = load_config(PROJECT_ROOT / "configs" / "edge_ai.yaml")
|
|
pipeline = config.pipelines["detection"]
|
|
detector = pipeline.nodes["phone_detector"]
|
|
gate = pipeline.nodes["phone_repeat_gate"]
|
|
|
|
assert detector.params["model"] == "people-talking-yolov8x@1"
|
|
assert detector.params["detect_labels"] == ["talking on phone"]
|
|
assert detector.params["model_options"]["weights"] == PEOPLE_TALKING_WEIGHTS
|
|
assert (PROJECT_ROOT / PEOPLE_TALKING_WEIGHTS).is_file()
|
|
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
|
|
)
|
|
|
|
|
|
def test_debug_viewer_uses_repository_mobile_phone_weights() -> None:
|
|
config = load_config(PROJECT_ROOT / "configs" / "debug" / "detection_viewer.yaml")
|
|
pipeline = config.pipelines["detection_show"]
|
|
ppe_detector = pipeline.nodes["detector"]
|
|
detector = pipeline.nodes["phone_detector"]
|
|
viewer = pipeline.nodes["viewer"]
|
|
|
|
assert detector.params["model"] == "yolov8n-mobile-phone@1"
|
|
assert detector.params["detect_labels"] == ["mobile_phone"]
|
|
assert detector.params["model_options"]["weights"] == MOBILE_PHONE_WEIGHTS
|
|
assert ppe_detector.params["attach_frame"] is True
|
|
assert detector.params["attach_frame"] is True
|
|
assert viewer.params["expected_inputs"] == ["ppe", "phone"]
|
|
assert (PROJECT_ROOT / MOBILE_PHONE_WEIGHTS).is_file()
|
|
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
|
|
)
|