217 lines
6.5 KiB
Python
217 lines
6.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import importlib.util
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from types import ModuleType, SimpleNamespace
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
|
||
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
SCRIPT = PROJECT_ROOT / "scripts" / "export_detection_onnx.py"
|
||
|
|
|
||
|
|
|
||
|
|
def _load_script() -> ModuleType:
|
||
|
|
spec = importlib.util.spec_from_file_location("cmvr_export_detection_onnx", SCRIPT)
|
||
|
|
assert spec is not None and spec.loader is not None
|
||
|
|
module = importlib.util.module_from_spec(spec)
|
||
|
|
# dataclasses resolves postponed annotations through sys.modules.
|
||
|
|
import sys
|
||
|
|
|
||
|
|
sys.modules[spec.name] = module
|
||
|
|
spec.loader.exec_module(module)
|
||
|
|
return module
|
||
|
|
|
||
|
|
|
||
|
|
exporter = _load_script()
|
||
|
|
|
||
|
|
|
||
|
|
def test_export_script_disables_ultralytics_autoinstall_and_has_four_models() -> None:
|
||
|
|
assert os.environ["YOLO_AUTOINSTALL"] == "false"
|
||
|
|
assert set(exporter.EXPORT_DEFINITIONS) == {
|
||
|
|
"construction-ppe-yolov8@2",
|
||
|
|
"ppe-6classes-yolov8n@2",
|
||
|
|
"people-talking-yolov8x@2",
|
||
|
|
"yolov8n-mobile-phone@2",
|
||
|
|
}
|
||
|
|
assert all(
|
||
|
|
definition.output.as_posix().endswith("/v2/model.onnx")
|
||
|
|
for definition in exporter.EXPORT_DEFINITIONS.values()
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_export_refuses_existing_destination_without_force(tmp_path: Path) -> None:
|
||
|
|
source = tmp_path / "source.pt"
|
||
|
|
output = tmp_path / "model.onnx"
|
||
|
|
manifest = tmp_path / "manifest.json"
|
||
|
|
source.write_bytes(b"trusted checkpoint")
|
||
|
|
output.write_bytes(b"existing")
|
||
|
|
|
||
|
|
with pytest.raises(FileExistsError, match="without --force"):
|
||
|
|
exporter._validate_paths(
|
||
|
|
source=source,
|
||
|
|
output=output,
|
||
|
|
manifest=manifest,
|
||
|
|
force=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_export_rejects_an_lfs_pointer_source(tmp_path: Path) -> None:
|
||
|
|
source = tmp_path / "source.pt"
|
||
|
|
source.write_text(
|
||
|
|
"version https://git-lfs.github.com/spec/v1\n"
|
||
|
|
"oid sha256:abc\nsize 123\n",
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="Git LFS pointer"):
|
||
|
|
exporter._validate_paths(
|
||
|
|
source=source,
|
||
|
|
output=tmp_path / "model.onnx",
|
||
|
|
manifest=tmp_path / "manifest.json",
|
||
|
|
force=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fake_export_is_staged_validated_and_published_with_digest(
|
||
|
|
tmp_path: Path,
|
||
|
|
monkeypatch: pytest.MonkeyPatch,
|
||
|
|
) -> None:
|
||
|
|
definition = exporter.EXPORT_DEFINITIONS["yolov8n-mobile-phone@2"]
|
||
|
|
source = tmp_path / "trusted.pt"
|
||
|
|
output = tmp_path / "v2" / "model.onnx"
|
||
|
|
manifest = output.parent / "manifest.json"
|
||
|
|
source.write_bytes(b"trusted-checkpoint")
|
||
|
|
export_calls: list[dict[str, Any]] = []
|
||
|
|
runtime_paths: list[Path] = []
|
||
|
|
|
||
|
|
class FakeYOLO:
|
||
|
|
task = "detect"
|
||
|
|
names = dict(enumerate(definition.labels))
|
||
|
|
|
||
|
|
def __init__(self, path: str, *, task: str) -> None:
|
||
|
|
assert task == "detect"
|
||
|
|
self.path = Path(path)
|
||
|
|
|
||
|
|
def export(self, **kwargs: Any) -> str:
|
||
|
|
export_calls.append(kwargs)
|
||
|
|
exported = self.path.with_suffix(".onnx")
|
||
|
|
exported.write_bytes(b"validated-onnx")
|
||
|
|
return str(exported)
|
||
|
|
|
||
|
|
class FakeRuntime:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
params: dict[str, Any],
|
||
|
|
*,
|
||
|
|
expected_labels: tuple[str, ...],
|
||
|
|
) -> None:
|
||
|
|
assert expected_labels == definition.labels
|
||
|
|
runtime_paths.append(Path(params["weights"]))
|
||
|
|
|
||
|
|
def load(self) -> None:
|
||
|
|
assert runtime_paths[-1].is_file()
|
||
|
|
|
||
|
|
def close(self) -> None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
fake_model = SimpleNamespace(
|
||
|
|
metadata_props=[
|
||
|
|
SimpleNamespace(
|
||
|
|
key="description",
|
||
|
|
value="trained on /private/training/data.yaml",
|
||
|
|
),
|
||
|
|
SimpleNamespace(key="date", value="non-deterministic"),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
sanitized_metadata: dict[str, str] = {}
|
||
|
|
|
||
|
|
def set_model_props(model: Any, metadata: dict[str, str]) -> None:
|
||
|
|
assert model is fake_model
|
||
|
|
sanitized_metadata.update(metadata)
|
||
|
|
|
||
|
|
fake_onnx = SimpleNamespace(
|
||
|
|
__version__="test-onnx",
|
||
|
|
load=lambda path: fake_model,
|
||
|
|
helper=SimpleNamespace(set_model_props=set_model_props),
|
||
|
|
save_model=lambda model, path: None,
|
||
|
|
checker=SimpleNamespace(check_model=lambda model: None),
|
||
|
|
)
|
||
|
|
dependencies = {
|
||
|
|
"YOLO": FakeYOLO,
|
||
|
|
"onnx": fake_onnx,
|
||
|
|
"onnxruntime": SimpleNamespace(__version__="test-ort"),
|
||
|
|
"onnxslim": None,
|
||
|
|
"torch": SimpleNamespace(__version__="test-torch"),
|
||
|
|
"ultralytics": SimpleNamespace(__version__="test-ultralytics"),
|
||
|
|
}
|
||
|
|
monkeypatch.setattr(
|
||
|
|
exporter,
|
||
|
|
"_require_export_dependencies",
|
||
|
|
lambda *, simplify: dependencies,
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(exporter, "OnnxYoloModel", FakeRuntime)
|
||
|
|
args = exporter.parse_args(
|
||
|
|
[
|
||
|
|
"--model-id",
|
||
|
|
definition.target_model_id,
|
||
|
|
"--source",
|
||
|
|
str(source),
|
||
|
|
"--output",
|
||
|
|
str(output),
|
||
|
|
"--manifest",
|
||
|
|
str(manifest),
|
||
|
|
"--imgsz",
|
||
|
|
"32",
|
||
|
|
"--opset",
|
||
|
|
"17",
|
||
|
|
"--no-simplify",
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
published_output, published_manifest = exporter.export_model(args)
|
||
|
|
|
||
|
|
assert published_output == output
|
||
|
|
assert published_manifest == manifest
|
||
|
|
assert output.read_bytes() == b"validated-onnx"
|
||
|
|
payload = json.loads(manifest.read_text(encoding="utf-8"))
|
||
|
|
assert payload["model_id"] == definition.target_model_id
|
||
|
|
assert payload["source_model_id"] == definition.source_model_id
|
||
|
|
assert payload["artifact_sha256"] == hashlib.sha256(
|
||
|
|
b"validated-onnx"
|
||
|
|
).hexdigest()
|
||
|
|
assert payload["export"] == {
|
||
|
|
"batch": 1,
|
||
|
|
"dynamic": False,
|
||
|
|
"format": "onnx",
|
||
|
|
"half": False,
|
||
|
|
"nms": False,
|
||
|
|
"opset": 17,
|
||
|
|
"simplify": False,
|
||
|
|
}
|
||
|
|
assert export_calls == [
|
||
|
|
{
|
||
|
|
"format": "onnx",
|
||
|
|
"imgsz": 32,
|
||
|
|
"batch": 1,
|
||
|
|
"dynamic": False,
|
||
|
|
"simplify": False,
|
||
|
|
"opset": 17,
|
||
|
|
"nms": False,
|
||
|
|
"half": False,
|
||
|
|
"device": "cpu",
|
||
|
|
}
|
||
|
|
]
|
||
|
|
assert runtime_paths and output.parent in runtime_paths[0].parents
|
||
|
|
assert not list(output.parent.glob(".onnx-export-*"))
|
||
|
|
assert "date" not in sanitized_metadata
|
||
|
|
assert "/private/" not in sanitized_metadata["description"]
|
||
|
|
assert sanitized_metadata["cmvr_model_id"] == definition.target_model_id
|
||
|
|
assert sanitized_metadata["cmvr_source_sha256"] == hashlib.sha256(
|
||
|
|
b"trusted-checkpoint"
|
||
|
|
).hexdigest()
|