85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import tarfile
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
|
|
def digest(path: Path) -> str:
|
|
value = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
while chunk := handle.read(1024 * 1024):
|
|
value.update(chunk)
|
|
return value.hexdigest()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Build a versioned ImageFind model bundle")
|
|
parser.add_argument("source", type=Path)
|
|
parser.add_argument("output", type=Path)
|
|
parser.add_argument("--version", default="1")
|
|
args = parser.parse_args()
|
|
source = args.source.resolve()
|
|
output = args.output.resolve()
|
|
for required in ("visual/image", "visual/text"):
|
|
if not (source / required).is_dir():
|
|
parser.error(f"missing {required}")
|
|
optional_components = {
|
|
"ocr": ("det.onnx", "rec.onnx", "cls.onnx"),
|
|
"faces": ("detector.xml", "detector.bin", "reidentification.xml", "reidentification.bin"),
|
|
}
|
|
for component, required_files in optional_components.items():
|
|
root = source / component
|
|
if root.exists():
|
|
missing = [name for name in required_files if not (root / name).is_file()]
|
|
if missing:
|
|
parser.error(f"incomplete {component}: missing {', '.join(missing)}")
|
|
audio_root = source / "audio"
|
|
if audio_root.exists() and (not (audio_root / "config.json").is_file() or not list(audio_root.glob("*.xml"))):
|
|
parser.error("incomplete audio: missing config.json or OpenVINO XML files")
|
|
try:
|
|
output.relative_to(source)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
parser.error("output must be outside the source model directory")
|
|
|
|
files = sorted(
|
|
path for path in source.rglob("*") if path.is_file() and path.relative_to(source).as_posix() != "manifest.json"
|
|
)
|
|
created_at = datetime.now(UTC)
|
|
manifest = {
|
|
"format_version": 2,
|
|
"version": args.version,
|
|
"created_at": created_at.isoformat(),
|
|
"source": "imagefind-build-model-bundle",
|
|
"components": {
|
|
"visual": {"version": args.version},
|
|
**({"ocr": {"version": args.version}} if (source / "ocr").is_dir() else {}),
|
|
**({"faces": {"version": args.version}} if (source / "faces").is_dir() else {}),
|
|
**({"audio": {"version": args.version}} if audio_root.is_dir() else {}),
|
|
},
|
|
"files": {path.relative_to(source).as_posix(): digest(path) for path in files},
|
|
}
|
|
manifest_bytes = (json.dumps(manifest, ensure_ascii=False, indent=2) + "\n").encode()
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tarfile.open(output, "w:gz") as archive:
|
|
for path in sorted(source.rglob("*")):
|
|
if path.relative_to(source).as_posix() == "manifest.json":
|
|
continue
|
|
archive.add(path, path.relative_to(source).as_posix(), recursive=False)
|
|
info = tarfile.TarInfo("manifest.json")
|
|
info.size = len(manifest_bytes)
|
|
info.mode = 0o644
|
|
info.mtime = int(created_at.timestamp())
|
|
archive.addfile(info, io.BytesIO(manifest_bytes))
|
|
print(f"bundle={output} sha256={digest(output)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|