File size: 1,407 Bytes
afef12f d31e060 afef12f | 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 | #!/usr/bin/env python3
"""Check whether a checkpoint can be exported at the selected input resolution."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from .converter_io import JsonObject, RejectedInput, atomic_write, canonical_bytes, sha256_bytes
from .export_yolo26n import ExportedModel, export_model
_VALIDATION_VERSION = 1
def _validation_result(exported: ExportedModel) -> JsonObject:
"""Describe the successfully exported artifact without retaining model data."""
return {
"version": _VALIDATION_VERSION,
"status": "accepted",
"artifact": {
"bytes": len(exported.payload),
"sha256": sha256_bytes(exported.payload),
},
}
def _main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--resolution", type=int, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
result = _validation_result(export_model(args.checkpoint, args.resolution))
atomic_write(args.output, canonical_bytes(result))
print(json.dumps(result, sort_keys=True))
if __name__ == "__main__":
try:
_main()
except (OSError, RejectedInput) as error:
raise SystemExit(f"YOLO26n checkpoint rejected: {error}") from None
|