File size: 2,332 Bytes
ebeed67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from __future__ import annotations
import pytest
from camstack_train.errors import TrainerError
from camstack_train.export import artefact_names, calibration_frames, export_all, remove_export
from .fakes import FakeToolchain
CID = "custom-x"
ALL = ["onnx", "openvino-fp16", "openvino-int8", "coreml-fp16"]
def run(tmp_path, toolchain, exports=ALL, calib=True):
pt = tmp_path / "best.pt"
pt.write_bytes(b"w")
images = tmp_path / "train"
images.mkdir()
for i in range(5):
(images / f"{i:06d}.jpg").write_bytes(b"j")
frames = calibration_frames(images, 3) if calib else []
return export_all(exports, pt, CID, 640, frames, tmp_path / "stage", toolchain)
def test_every_format_lands_under_its_flat_name(tmp_path):
tc = FakeToolchain()
out = run(tmp_path, tc)
assert out.produced == ALL and out.dropped == []
stage = tmp_path / "stage"
for name in artefact_names(CID).values():
assert (stage / name).exists(), name
assert (stage / f"{CID}-int8.bin").is_file()
assert [p.name for p in tc.calib] == ["000000.jpg", "000001.jpg", "000002.jpg"]
def test_failing_coreml_drops_only_coreml(tmp_path):
out = run(tmp_path, FakeToolchain(fail=("coreml",)))
assert "coreml-fp16" not in out.produced
assert out.dropped == [{"export": "coreml-fp16", "reason": "coreml exploded"}]
assert "openvino-int8" in out.produced
def test_onnx_failure_fails_the_job(tmp_path):
with pytest.raises(TrainerError) as err:
run(tmp_path, FakeToolchain(fail=("onnx",)))
assert err.value.code == "export-onnx" and err.value.exit_code == 70
def test_int8_without_calibration_frames_is_dropped(tmp_path):
out = run(tmp_path, FakeToolchain(), calib=False)
assert {"export": "openvino-int8", "reason": "no calibration frames in the train split"} in out.dropped
def test_only_requested_exports_run(tmp_path):
tc = FakeToolchain()
run(tmp_path, tc, exports=["onnx"])
assert tc.calls == ["onnx"]
def test_remove_export_deletes_xml_and_bin(tmp_path):
run(tmp_path, FakeToolchain())
remove_export(tmp_path / "stage", CID, "openvino-int8")
assert not (tmp_path / "stage" / f"{CID}-int8.xml").exists()
assert not (tmp_path / "stage" / f"{CID}-int8.bin").exists()
assert (tmp_path / "stage" / f"{CID}-fp16.bin").exists()
|