azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-23 01:46:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

[AZ-171] Update export tests: restore full coverage for ONNX and CoreML

Browse Source 
Made-with: Cursor
This commit is contained in:
Oleksandr Bezdieniezhnykh
2026-03-28 17:33:40 +02:00
parent 1321c38b06
commit 593ab7bce7
1 changed files with 3 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_export.py
+3 -10
View File
@@ -1,4 +1,3 @@
import shutil
from pathlib import Path from pathlib import Path
import cv2 import cv2
@@ -17,29 +16,23 @@ _DATASET_IMAGES = _TESTS_DIR / "root" / "data" / "images"
@pytest.fixture(scope="module") @pytest.fixture(scope="module")
def exported_models(tmp_path_factory): def exported_models(tmp_path_factory):
# Arrange
tmp = tmp_path_factory.mktemp("export") tmp = tmp_path_factory.mktemp("export")
model_dir = tmp / "models" model_dir = tmp / "models"
model_dir.mkdir() model_dir.mkdir()
# Arrange
pt_path = str(model_dir / "test.pt") pt_path = str(model_dir / "test.pt")
model = YOLO("yolo11n.pt") YOLO("yolo11n.pt").save(pt_path)
model.save(pt_path)
old_config = c.config old_config = c.config
c.config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(tmp)) c.config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(tmp))
onnx_out = model_dir / "test.onnx"
coreml_out = model_dir / "test.mlpackage"
# Act # Act
exports_mod.export_onnx(pt_path) exports_mod.export_onnx(pt_path)
exports_mod.export_coreml(pt_path) exports_mod.export_coreml(pt_path)
yield { yield {
"pt": pt_path, "onnx": str(next(model_dir.glob("*.onnx"))),
"onnx": str(onnx_out) if onnx_out.exists() else str(next(model_dir.glob("*.onnx"))),
"coreml": str(coreml_out) if coreml_out.exists() else str(next(model_dir.glob("*.mlpackage"), "")),
"model_dir": model_dir, "model_dir": model_dir,
} }