154 lines
4.8 KiB
Python
154 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
from ...config import settings
|
|
|
|
WEIGHT_EXTS = {".pt", ".pth", ".onnx", ".engine"}
|
|
|
|
|
|
def sha256_file(path: Path, chunk_size: int = 1024 * 1024 * 8) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
while True:
|
|
chunk = handle.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def iter_source_weights() -> Iterable[Path]:
|
|
project_root = settings.project_root.resolve()
|
|
for path in settings.source_root.rglob("*"):
|
|
if not path.is_file() or path.suffix.lower() not in WEIGHT_EXTS:
|
|
continue
|
|
try:
|
|
path.resolve().relative_to(project_root)
|
|
continue
|
|
except ValueError:
|
|
yield path
|
|
|
|
|
|
def classify_weight(path: Path) -> dict[str, str]:
|
|
try:
|
|
rel = str(path.resolve().relative_to(settings.source_root))
|
|
except ValueError:
|
|
rel = str(path)
|
|
lower = rel.lower()
|
|
if "yolo" in lower:
|
|
family = "yolo"
|
|
elif "mmseg" in lower or "my_local_model" in lower:
|
|
family = "mmseg"
|
|
elif "segmodel" in lower:
|
|
family = "segmodel"
|
|
else:
|
|
family = "misc"
|
|
if "best.pt" in lower or "best.pth" in lower:
|
|
role = "trained_best"
|
|
elif "last.pt" in lower or "last.pth" in lower:
|
|
role = "trained_last"
|
|
elif "pretrain" in lower or "my_local_model" in lower:
|
|
role = "pretrained"
|
|
else:
|
|
role = "weight"
|
|
return {"family": family, "role": role}
|
|
|
|
|
|
def copy_weight(src: Path, dst: Path, mode: str) -> None:
|
|
dst.parent.mkdir(parents=True, exist_ok=True)
|
|
if mode == "hardlink":
|
|
if dst.exists():
|
|
dst.unlink()
|
|
os.link(src, dst)
|
|
elif mode == "reflink":
|
|
subprocess.run(["cp", "--reflink=auto", "--preserve=timestamps", str(src), str(dst)], check=True)
|
|
else:
|
|
shutil.copy2(src, dst)
|
|
|
|
|
|
def sync_weights(mode: str = "copy", hash_files: bool = True, skip_existing: bool = True) -> dict:
|
|
files_dir = settings.weights_root / "files"
|
|
settings.weights_root.mkdir(parents=True, exist_ok=True)
|
|
entries = []
|
|
total_bytes = 0
|
|
for src in sorted(iter_source_weights()):
|
|
rel = src.resolve().relative_to(settings.source_root)
|
|
dst = files_dir / rel
|
|
stat = src.stat()
|
|
total_bytes += stat.st_size
|
|
copied = False
|
|
if not (skip_existing and dst.exists() and dst.stat().st_size == stat.st_size):
|
|
copy_weight(src, dst, mode)
|
|
copied = True
|
|
meta = classify_weight(src)
|
|
entry = {
|
|
"source_path": str(rel),
|
|
"stored_path": str(dst.resolve().relative_to(settings.project_root)),
|
|
"size": stat.st_size,
|
|
"family": meta["family"],
|
|
"role": meta["role"],
|
|
"copied": copied,
|
|
}
|
|
if hash_files:
|
|
entry["sha256"] = sha256_file(dst)
|
|
entries.append(entry)
|
|
|
|
manifest = {
|
|
"generated_at": datetime.now(timezone.utc).isoformat(),
|
|
"updated_at": datetime.now(timezone.utc).isoformat(),
|
|
"source_root": str(settings.source_root),
|
|
"mode": mode,
|
|
"count": len(entries),
|
|
"total_bytes": total_bytes,
|
|
"files": entries,
|
|
}
|
|
manifest_path = settings.weights_root / "manifest.json"
|
|
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
return manifest
|
|
|
|
|
|
def load_manifest() -> dict:
|
|
manifest_path = settings.weights_root / "manifest.json"
|
|
if not manifest_path.exists():
|
|
return {
|
|
"generated_at": None,
|
|
"source_root": str(settings.source_root),
|
|
"count": 0,
|
|
"total_bytes": 0,
|
|
"files": [],
|
|
}
|
|
return json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def verify_weights() -> dict:
|
|
manifest = load_manifest()
|
|
checked = []
|
|
ok_count = 0
|
|
for entry in manifest.get("files", []):
|
|
path = settings.project_root / entry["stored_path"]
|
|
exists = path.exists()
|
|
size_ok = exists and path.stat().st_size == entry.get("size")
|
|
hash_ok = None
|
|
if exists and "sha256" in entry:
|
|
hash_ok = sha256_file(path) == entry["sha256"]
|
|
ok = bool(exists and size_ok and (hash_ok is not False))
|
|
ok_count += int(ok)
|
|
checked.append(
|
|
{
|
|
"stored_path": entry["stored_path"],
|
|
"exists": exists,
|
|
"size_ok": size_ok,
|
|
"hash_ok": hash_ok,
|
|
"ok": ok,
|
|
}
|
|
)
|
|
return {"count": len(checked), "ok_count": ok_count, "items": checked}
|