training loop: run ledger, studio export importer with an image pool, train_run with gold gate and auto-promotion, SLURM template, CI

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

tools/import_studio_export.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""Bring a studio export (``python -m src.pipeline.run export``) into this repository.
+
+Images once, labels per dataset: every frame lands in ``data/pool/images/`` exactly
+once (LFS), and the dataset directory holds symlinks into the pool in the layout
+Ultralytics expects, so labels resolve next to them:
+
+    data/datasets/<name>/
+    ├── images/{train,val}/tpi_<id>.jpg -> ../../../../pool/images/tpi_<id>.jpg
+    ├── labels/{train,val}/tpi_<id>.txt
+    ├── dataset.yaml          (path: . ; resolved to an absolute path at train time)
+    ├── manifest.csv          (the studio's provenance: tpid, split, reason, beamline, ...)
+    └── source.json           (where it came from, counts, classes, gold manifest hash)
+
+A second dataset re-lists the frames it shares with earlier ones and adds only the
+new images to the pool, so ``git diff dataset/A dataset/B`` shows what changed.
+
+    python tools/import_studio_export.py /path/to/dataset_20261002_1200 --tag alc \\
+        [--gold /path/to/studio/data/gold] [--commit] [--git-tag]
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import os
+import re
+import shutil
+import subprocess
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional
+
+ROOT = Path(__file__).resolve().parents[1]
+IMG_EXT = (".jpg", ".jpeg", ".png")
+
+
+def _sha256(p: Path) -> str:
+    h = hashlib.sha256()
+    with p.open("rb") as fh:
+        for chunk in iter(lambda: fh.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def read_names(dataset_yaml: Path) -> Dict[int, str]:
+    """``names:`` block of a YOLO dataset.yaml, without PyYAML."""
+    names: Dict[int, str] = {}
+    in_names = False
+    for line in dataset_yaml.read_text(encoding="utf-8").splitlines():
+        if re.match(r"^names\s*:", line):
+            in_names = True
+            continue
+        if in_names:
+            m = re.match(r"^\s+(\d+)\s*:\s*(\S+)", line)
+            if m:
+                names[int(m.group(1))] = m.group(2)
+            elif line.strip() and not line.startswith(" "):
+                break
+    return names
+
+
+def write_dataset_yaml(dest: Path, names: Dict[int, str]) -> Path:
+    body = "\n".join(f"  {i}: {names[i]}" for i in sorted(names))
+    p = dest / "dataset.yaml"
+    p.write_text(
+        "# path is resolved to this directory's absolute path by tools/train_run.py\n"
+        f"path: .\ntrain: images/train\nval: images/val\n\nnames:\n{body}\n",
+        encoding="utf-8",
+    )
+    return p
+
+
+def add_to_pool(image: Path, pool: Path) -> tuple[Path, bool]:
+    """Copy ``image`` into the pool unless an identical file is already there."""
+    pool.mkdir(parents=True, exist_ok=True)
+    dest = pool / image.name
+    if dest.exists():
+        if dest.stat().st_size == image.stat().st_size:
+            return dest, False
+        raise SystemExit(f"{dest} exists with a different size than {image}; a pool frame never changes")
+    shutil.copy2(image, dest)
+    return dest, True
+
+
+def link_into(dataset_images: Path, pool_file: Path) -> Path:
+    dataset_images.mkdir(parents=True, exist_ok=True)
+    link = dataset_images / pool_file.name
+    rel = os.path.relpath(pool_file, dataset_images)
+    if link.is_symlink() or link.exists():
+        link.unlink()
+    link.symlink_to(rel)
+    return link
+
+
+def import_export(export_dir: Path, name: str, *, root: Path = ROOT, gold_dir: Optional[Path] = None) -> Dict:
+    export_dir = Path(export_dir)
+    if not (export_dir / "dataset.yaml").exists():
+        raise SystemExit(f"{export_dir} is not a studio export (no dataset.yaml)")
+    names = read_names(export_dir / "dataset.yaml")
+    pool = root / "data" / "pool" / "images"
+    dest = root / "data" / "datasets" / name
+    if dest.exists():
+        raise SystemExit(f"{dest} exists; datasets are immutable, pick another name")
+    counts: Dict[str, Dict[str, int]] = {}
+    added = 0
+    for split in ("train", "val"):
+        src_images = export_dir / split / "images"
+        src_labels = export_dir / split / "labels"
+        n_img = n_lbl = n_empty = 0
+        for img in sorted(p for p in src_images.glob("*") if p.suffix.lower() in IMG_EXT):
+            pool_file, new = add_to_pool(img, pool)
+            added += int(new)
+            link_into(dest / "images" / split, pool_file)
+            lbl = src_labels / f"{img.stem}.txt"
+            out_lbl = dest / "labels" / split / f"{img.stem}.txt"
+            out_lbl.parent.mkdir(parents=True, exist_ok=True)
+            if lbl.exists():
+                shutil.copy2(lbl, out_lbl)
+                n_lbl += 1
+                if not lbl.read_text(encoding="utf-8").strip():
+                    n_empty += 1
+            else:
+                out_lbl.write_text("", encoding="utf-8")  # background frame
+                n_empty += 1
+            n_img += 1
+        counts[split] = {"images": n_img, "labels": n_lbl, "empty": n_empty}
+    write_dataset_yaml(dest, names)
+    if (export_dir / "manifest.csv").exists():
+        shutil.copy2(export_dir / "manifest.csv", dest / "manifest.csv")
+    gold_info = None
+    if gold_dir:
+        gold_info = import_gold(Path(gold_dir), root=root)
+    source = {
+        "name": name,
+        "imported_at": datetime.now().isoformat(timespec="seconds"),
+        "export_dir": str(export_dir.resolve()),
+        "classes": names,
+        "counts": counts,
+        "new_pool_images": added,
+        "manifest_sha256": _sha256(dest / "manifest.csv") if (dest / "manifest.csv").exists() else None,
+        "gold": gold_info,
+    }
+    (dest / "source.json").write_text(json.dumps(source, indent=2) + "\n", encoding="utf-8")
+    return source
+
+
+def import_gold(gold_dir: Path, *, root: Path = ROOT) -> Dict:
+    """The studio's frozen gold set: images into the pool, symlinks + labels + manifest under data/gold."""
+    pool = root / "data" / "pool" / "images"
+    dest = root / "data" / "gold"
+    n = 0
+    for img in sorted(p for p in (gold_dir / "images").glob("*") if p.suffix.lower() in IMG_EXT):
+        pool_file, _ = add_to_pool(img, pool)
+        link_into(dest / "images", pool_file)
+        lbl = gold_dir / "labels" / f"{img.stem}.txt"
+        (dest / "labels").mkdir(parents=True, exist_ok=True)
+        if lbl.exists():
+            shutil.copy2(lbl, dest / "labels" / lbl.name)
+        n += 1
+    manifest = gold_dir / "gold_manifest.json"
+    info = {"images": n, "manifest_sha256": None}
+    if manifest.exists():
+        shutil.copy2(manifest, dest / "gold_manifest.json")
+        info["manifest_sha256"] = _sha256(manifest)
+    return info
+
+
+def git(root: Path, *args: str) -> str:
+    return subprocess.check_output(["git", "-C", str(root), *args], text=True).strip()
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("export_dir")
+    ap.add_argument("--tag", required=True, help="short dataset name; directory = <date>_<tag>")
+    ap.add_argument("--name", default=None, help="full dataset name instead of <date>_<tag>")
+    ap.add_argument("--gold", default=None, help="studio gold directory (images/, labels/, gold_manifest.json)")
+    ap.add_argument("--root", default=str(ROOT))
+    ap.add_argument("--commit", action="store_true", help="git add + commit the dataset (and pool additions)")
+    ap.add_argument("--git-tag", action="store_true", help="also tag the commit dataset/<name>")
+    args = ap.parse_args(argv)
+    root = Path(args.root)
+    name = args.name or f"{datetime.now().strftime('%Y-%m-%d')}_{args.tag}"
+    source = import_export(Path(args.export_dir), name, root=root, gold_dir=Path(args.gold) if args.gold else None)
+    c = source["counts"]
+    print(f"dataset {name}: train {c['train']['images']} / val {c['val']['images']} images, "
+          f"{source['new_pool_images']} new in the pool")
+    if args.commit:
+        git(root, "add", "data/pool/images", f"data/datasets/{name}")
+        if args.gold:
+            git(root, "add", "data/gold")
+        msg = (f"dataset: {name} ({c['train']['images']} train / {c['val']['images']} val, "
+               f"{source['new_pool_images']} new pool images)")
+        git(root, "commit", "-q", "-m", msg)
+        print(f"committed: {msg}")
+        if args.git_tag:
+            git(root, "tag", f"dataset/{name}")
+            print(f"tagged dataset/{name}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())