Initial Seg Data Server Net platform

backend/app/modules/weights/service.py

@@ -0,0 +1,153 @@
+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}