73 lines
2.0 KiB
Python
73 lines
2.0 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"
|
|
MODEL_ARTIFACTS = (
|
|
(
|
|
"construction-ppe-yolov8",
|
|
22_537_898,
|
|
"31ef3ca04a17cf545f3fcfc64c4af8993a41d52ccc460e82aff01d5354603533",
|
|
),
|
|
(
|
|
"ppe-6classes-yolov8n",
|
|
5_625_014,
|
|
"07172ef3ae9e256c40a1fb0ce3eefe5547d90170645aa73dded0fffc382cdb31",
|
|
),
|
|
)
|
|
|
|
|
|
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
|
|
|
|
|
|
@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()
|