From 4698dbeaf993d81eccdb491b7eb53fe7df44c3e5 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:06:06 +0200 Subject: [PATCH] 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 --- .gitea/workflows/ci.yml | 24 +-- .gitignore | 8 + README.md | 33 ++++ data/.gitattributes | 1 + models/.gitattributes | 1 + runs/README.md | 8 + slurm/README.md | 28 +++ slurm/train.sbatch | 41 +++++ src/eval/__init__.py | 4 + src/eval/evaluate.py | 88 +++++++++ src/eval/metrics.py | 136 ++++++++++++++ src/eval/yolo_labels.py | 123 +++++++++++++ tests/test_import_studio_export.py | 66 +++++++ tests/test_ledger.py | 48 +++++ tools/__init__.py | 0 tools/import_studio_export.py | 206 +++++++++++++++++++++ tools/ledger.py | 283 +++++++++++++++++++++++++++++ tools/train_run.py | 227 +++++++++++++++++++++++ 18 files changed, 1313 insertions(+), 12 deletions(-) create mode 100644 runs/README.md create mode 100644 slurm/README.md create mode 100644 slurm/train.sbatch create mode 100644 src/eval/__init__.py create mode 100644 src/eval/evaluate.py create mode 100644 src/eval/metrics.py create mode 100644 src/eval/yolo_labels.py create mode 100644 tests/test_import_studio_export.py create mode 100644 tests/test_ledger.py create mode 100644 tools/__init__.py create mode 100755 tools/import_studio_export.py create mode 100755 tools/ledger.py create mode 100644 tools/train_run.py diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml index 1b239b6..2e0850c 100644 --- a/.gitea/workflows/ci.yml +++ b/.gitea/workflows/ci.yml @@ -7,24 +7,24 @@ on: branches: ["main"] jobs: - test: + ledger-and-tools: runs-on: ubuntu-latest - steps: - - name: Checkout + - name: Checkout (no LFS objects needed) uses: actions/checkout@v4 + with: + lfs: false - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.11" -# - name: Install dependencies -# run: | -# #python -m pip install --upgrade pip -# #pip install -r requirements.txt -# #pip install pytest pytest-cov -# -# - name: Run unit tests -# run: | -# #pytest tests/ -v --tb=short --cov=src --cov-report=term-missing + - name: Install test dependencies + run: python -m pip install --upgrade pip pytest + + - name: Ledger and importer tests + run: pytest tests/test_ledger.py tests/test_import_studio_export.py -q + + - name: Every run.json valid + run: python tools/ledger.py check diff --git a/.gitignore b/.gitignore index ffbb530..956c041 100644 --- a/.gitignore +++ b/.gitignore @@ -429,3 +429,11 @@ models/registry/*/exports/*.engine models/base/ models/archived/ /training_runs/ + +# --- Training loop (2026-09-28): registry weights live in git via LFS, runs are tracked --- +!models/registry/*/weights.pt +!models/registry/*/gold_report.json +runs/*/ultralytics/ +runs/*/gold_predictions/ +runs/*/dataset.resolved.yaml +runs/slurm-*.out diff --git a/README.md b/README.md index 8a747a8..ea53b8e 100644 --- a/README.md +++ b/README.md @@ -273,3 +273,36 @@ A Gitea Actions workflow at `.gitea/workflows/ci.yml` runs the test suite automa - no pairs found: check that image names match label names (`.txt`) - missing classes in logs: verify label files contain valid class ids in first column - model download fails: use a local `.pt` model from `models/` + +## Training loop (from the annotation pipeline to a deployed model) + +Decided on 2026-09-28: images live once in `data/pool/images/` (LFS); a dataset is a +directory under `data/datasets/_/` with symlinks into the pool, its labels, +`manifest.csv`, `dataset.yaml` and `source.json`, tagged `dataset/`; every run is +a directory under `runs/` tracked in git with fixed states; every retraining starts +from the base weights on the whole labelled pool; the gate promotes on its own. + +```bash +# 1. the studio exports the cumulative labelled pool and hands it over (Spark): +python -m src.pipeline.run export --config pipeline.prod.json --train-repo ~/AareLC_train --train-tag alc +# = tools/import_studio_export.py --tag alc --gold --commit --git-tag +git push && git push --tags + +# 2. a run: exported -> queued -> training -> measured -> gate_passed|gate_failed -> promoted -> deployed +python tools/ledger.py new --dataset dataset/2026-10-02_alc --tag alc --base yolo26n-seg.pt +python tools/train_run.py 2026-10-02_alc --device 0 2>&1 | tee runs/2026-10-02_alc/train.log +python tools/ledger.py show # the table; `render` writes runs/README.md; `check` runs in CI + +# 3. on the cluster instead of the Spark, once slurm/README.md is filled in: +sbatch slurm/train.sbatch 2026-10-02_alc +``` + +`tools/train_run.py` trains with the recipe from `src/training/train_yolo_seg.py`, +copies `best.pt` to `models/registry/-seg--/weights.pt` (LFS) with +`metadata.json`, measures the weights on `data/gold/` with the studio's evaluation +code (`src/eval/`, vendored), and gates against the active model's +`gold_report.json`: no drop beyond 0.005 in overall F1 or mask IoU, an improvement in +at least one, no drop beyond 0.02 in F1 for loop_all or pin. A pass rewrites +`models/active/production.json`; deployment (AareLC_infer tag, cps-deployments bump, +TensorRT engine on the target GPU) stays a human step recorded with +`python tools/ledger.py set deployed`. diff --git a/data/.gitattributes b/data/.gitattributes index 4b84726..79a6659 100644 --- a/data/.gitattributes +++ b/data/.gitattributes @@ -1 +1,2 @@ datasets/** filter=lfs diff=lfs merge=lfs -text +pool/** filter=lfs diff=lfs merge=lfs -text diff --git a/models/.gitattributes b/models/.gitattributes index 05ca2fc..a8610cc 100644 --- a/models/.gitattributes +++ b/models/.gitattributes @@ -1 +1,2 @@ base/** filter=lfs diff=lfs merge=lfs -text +registry/**/*.pt filter=lfs diff=lfs merge=lfs -text diff --git a/runs/README.md b/runs/README.md new file mode 100644 index 0000000..3968103 --- /dev/null +++ b/runs/README.md @@ -0,0 +1,8 @@ +# Training runs + +Generated by `python tools/ledger.py render`; do not edit by hand. + +| Run | Status | Dataset | Base | Version | F1 | Mask IoU | Gate | Last change | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | + +States: exported → queued → training → measured → gate_passed → gate_failed → promoted → deployed (or aborted). diff --git a/slurm/README.md b/slurm/README.md new file mode 100644 index 0000000..985c54a --- /dev/null +++ b/slurm/README.md @@ -0,0 +1,28 @@ +# Training on the GPU cluster (template, not yet filled in) + +`train.sbatch` runs one ledger run on the `mx` partition. It needs these facts, +which replace the `__PLACEHOLDERS__` in the script: + +| Placeholder | What | Value | +| --- | --- | --- | +| `__PARTITION__` | SLURM partition | `mx` | +| `__ACCOUNT__` | SLURM account to charge | | +| `__WALLTIME__` | wall-time limit for 400 epochs of yolo26n at 640 px | | +| `__SHARED_FS__` | shared filesystem visible from login and compute nodes, holds the repo checkout | | +| `__MODULES__` / `__VENV__` / `__CONTAINER__` | how Python, CUDA and ultralytics are provided on compute nodes (one of the three) | | +| `__DEPLOY_KEY__` | SSH key on the cluster allowed to push to this repo, if the job reports back itself | | + +Still to confirm before the first job: the cluster and login node names; whether +compute nodes reach gitea.psi.ch (git + LFS) and PyPI, or whether a container has to +carry everything; the GPU type on the partition; who submits (a service account or +a person). Until then the reference runs happen on the DGX Spark with the same +`tools/train_run.py`. + +Workflow once filled in: + +```bash +python tools/ledger.py new --dataset dataset/ --tag --base yolo26n-seg.pt +git commit -am "run : exported" && git push +sbatch slurm/train.sbatch +python tools/ledger.py show +``` diff --git a/slurm/train.sbatch b/slurm/train.sbatch new file mode 100644 index 0000000..9968636 --- /dev/null +++ b/slurm/train.sbatch @@ -0,0 +1,41 @@ +#!/bin/bash +# Train one ledger run on the GPU cluster. Placeholders __LIKE_THIS__ are filled in +# once the cluster facts are known (see slurm/README.md). +# +# sbatch slurm/train.sbatch # run created with: python tools/ledger.py new ... +# +#SBATCH --job-name=aarelc-train +#SBATCH --partition=__PARTITION__ # mx +#SBATCH --account=__ACCOUNT__ +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=__WALLTIME__ # e.g. 08:00:00 +#SBATCH --output=__SHARED_FS__/AareLC_train/runs/slurm-%j.out +set -euo pipefail + +RUN_ID="${1:?usage: sbatch slurm/train.sbatch }" +REPO="${AARELC_TRAIN_REPO:-__SHARED_FS__/AareLC_train}" # a checkout on the shared filesystem + +# --- environment: keep exactly one of the three --------------------------------- +# module load __MODULES__ # e.g. cuda/12 python/3.12 +# source __VENV__/bin/activate # venv with ultralytics + opencv +# APPTAINER="apptainer exec --nv __CONTAINER__" # container with CUDA + ultralytics +APPTAINER="${APPTAINER:-}" + +cd "$REPO" +git fetch -q --tags +git pull -q --ff-only +DATASET_TAG="$(python -c "import json;print(json.load(open('runs/$RUN_ID/run.json'))['dataset'])")" +git checkout -q "$DATASET_TAG" -- data/datasets # the dataset exactly as tagged +git lfs pull -q -I "data/pool/images/**,data/gold/**" # only the labelled pool and the gold set + +python tools/ledger.py set "$RUN_ID" queued --kv "job=$SLURM_JOB_ID" "host=$(hostname)" "partition=$SLURM_JOB_PARTITION" +$APPTAINER python tools/train_run.py "$RUN_ID" --device 0 --workers "$SLURM_CPUS_PER_TASK" 2>&1 | tee "runs/$RUN_ID/train.log" +STATUS=${PIPESTATUS[0]} + +# --- report back: commit the run, the registry entry and, on a pass, the promotion -- +git add "runs/$RUN_ID" models +git commit -q -m "run $RUN_ID: $(python -c "import json;print(json.load(open('runs/$RUN_ID/run.json'))['status'])")" || true +# git push -q origin HEAD:main # needs a deploy key on the cluster (__DEPLOY_KEY__) +exit "$STATUS" diff --git a/src/eval/__init__.py b/src/eval/__init__.py new file mode 100644 index 0000000..10e79bf --- /dev/null +++ b/src/eval/__init__.py @@ -0,0 +1,4 @@ +"""Gold-set evaluation, vendored from AareLC_studio/src/pipeline (metrics, yolo_labels, evaluate) on 2026-09-28. + +Kept identical on purpose so a measure here and in the studio agree; update by re-copying. +""" diff --git a/src/eval/evaluate.py b/src/eval/evaluate.py new file mode 100644 index 0000000..1ba8cd0 --- /dev/null +++ b/src/eval/evaluate.py @@ -0,0 +1,88 @@ +"""Eval harness: compare predicted YOLO-seg labels against a gold set. + +Reports detection precision/recall/F1 per class (bbox IoU >= threshold to match) +and mean mask IoU over matched instances. Operates on label files only (no source +images needed), so it can score any pre-labeler / model against the frozen gold set. + + python -m src.pipeline.evaluate --pred OUT/labels --gold GOLD/labels [--names dataset.yaml] +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Dict, List + +from src.eval.metrics import aggregate, evaluate_image, full_report +from src.eval.yolo_labels import load_class_names, read_label_file + + +def _stems(label_dir: Path) -> Dict[str, Path]: + return {p.stem: p for p in label_dir.glob("*.txt")} + + +def evaluate_dirs(pred_dir, gold_dir, iou_thr: float = 0.5, + names: List[str] | None = None) -> Dict: + pred_dir, gold_dir = Path(pred_dir), Path(gold_dir) + gold_files = _stems(gold_dir) + pred_files = _stems(pred_dir) + + # Gold defines the evaluation set; a missing prediction counts as all-FN. + per_image = [] + missing_pred = 0 + for stem, gpath in sorted(gold_files.items()): + gold = read_label_file(gpath) + ppath = pred_files.get(stem) + pred = read_label_file(ppath) if ppath else [] + if ppath is None: + missing_pred += 1 + per_image.append(evaluate_image(gold, pred, iou_thr)) + + report = full_report(aggregate(per_image), names) + report["dataset"] = { + "gold_images": len(gold_files), + "pred_images": len(pred_files), + "missing_predictions": missing_pred, + "iou_threshold": iou_thr, + } + return report + + +def _print_report(report: Dict) -> None: + ds = report["dataset"] + print(f"\nGold images: {ds['gold_images']} | predictions: {ds['pred_images']} " + f"| missing: {ds['missing_predictions']} | match IoU >= {ds['iou_threshold']}") + header = f"{'class':<14}{'P':>7}{'R':>7}{'F1':>7}{'maskIoU':>9}{'TP':>6}{'FP':>6}{'FN':>6}" + print("\n" + header) + print("-" * len(header)) + + def row(name, r): + print(f"{name:<14}{r['precision']:>7.3f}{r['recall']:>7.3f}{r['f1']:>7.3f}" + f"{r['mean_mask_iou']:>9.3f}{r['tp']:>6}{r['fp']:>6}{r['fn']:>6}") + + for name, r in report["per_class"].items(): + row(name, r) + print("-" * len(header)) + row("OVERALL", report["overall"]) + + +def main() -> int: + ap = argparse.ArgumentParser(description="Evaluate predicted YOLO-seg labels vs. a gold set") + ap.add_argument("--pred", required=True, help="Directory of predicted .txt labels") + ap.add_argument("--gold", required=True, help="Directory of gold .txt labels") + ap.add_argument("--iou", type=float, default=0.5, help="bbox IoU threshold to match instances") + ap.add_argument("--names", default=None, help="Optional dataset.yaml for class names") + ap.add_argument("--report", default=None, help="Optional path to write the JSON report") + args = ap.parse_args() + + names = load_class_names(args.names) if args.names else [] + report = evaluate_dirs(args.pred, args.gold, iou_thr=args.iou, names=names) + _print_report(report) + if args.report: + Path(args.report).write_text(json.dumps(report, indent=2), encoding="utf-8") + print(f"\nWrote {args.report}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/eval/metrics.py b/src/eval/metrics.py new file mode 100644 index 0000000..16f6054 --- /dev/null +++ b/src/eval/metrics.py @@ -0,0 +1,136 @@ +"""Quality metrics for auto-annotation: detection P/R/F1 + per-instance mask IoU. + +Works entirely on normalized polygons (YOLO-seg form), so no source images are +needed at eval time: masks are rasterized into a fixed grid for IoU. Matching is +greedy within a class by bbox IoU; mask IoU is reported for matched pairs. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Sequence, Tuple + +import cv2 +import numpy as np + +from src.eval.yolo_labels import PolygonInstance + +RASTER = 512 # grid resolution for mask IoU + + +def bbox_iou(a: Sequence[float], b: Sequence[float]) -> float: + """IoU of two [x1, y1, x2, y2] boxes (any consistent units).""" + ax1, ay1, ax2, ay2 = a + bx1, by1, bx2, by2 = b + ix1, iy1 = max(ax1, bx1), max(ay1, by1) + ix2, iy2 = min(ax2, bx2), min(ay2, by2) + iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1) + inter = iw * ih + area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1) + area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1) + union = area_a + area_b - inter + return float(inter / union) if union > 0 else 0.0 + + +def _raster(points_norm, size: int = RASTER) -> np.ndarray: + m = np.zeros((size, size), dtype=np.uint8) + if len(points_norm) >= 3: + pts = np.array([(x * size, y * size) for x, y in points_norm], dtype=np.int32).reshape(-1, 1, 2) + cv2.fillPoly(m, [pts], 1) + return m + + +def mask_iou(points_a_norm, points_b_norm, size: int = RASTER) -> float: + """IoU of two normalized polygons, via rasterization into a fixed grid.""" + if len(points_a_norm) < 3 or len(points_b_norm) < 3: + return 0.0 + ma, mb = _raster(points_a_norm, size), _raster(points_b_norm, size) + inter = int(np.logical_and(ma, mb).sum()) + union = int(np.logical_or(ma, mb).sum()) + return float(inter / union) if union > 0 else 0.0 + + +@dataclass +class ClassCount: + tp: int = 0 + fp: int = 0 + fn: int = 0 + mask_ious: List[float] = field(default_factory=list) + + def merge(self, other: "ClassCount") -> None: + self.tp += other.tp + self.fp += other.fp + self.fn += other.fn + self.mask_ious.extend(other.mask_ious) + + +def _match_class(gold: List[PolygonInstance], pred: List[PolygonInstance], + iou_thr: float) -> ClassCount: + """Greedy bbox-IoU matching within a single class.""" + cc = ClassCount() + gold_boxes = [g.bbox() for g in gold] + pred_boxes = [p.bbox() for p in pred] + + pairs: List[Tuple[float, int, int]] = [] + for gi, gb in enumerate(gold_boxes): + for pi, pb in enumerate(pred_boxes): + iou = bbox_iou(gb, pb) + if iou >= iou_thr: + pairs.append((iou, gi, pi)) + pairs.sort(reverse=True) + + used_g, used_p = set(), set() + for iou, gi, pi in pairs: + if gi in used_g or pi in used_p: + continue + used_g.add(gi) + used_p.add(pi) + cc.tp += 1 + cc.mask_ious.append(mask_iou(gold[gi].points, pred[pi].points)) + + cc.fn = len(gold) - len(used_g) + cc.fp = len(pred) - len(used_p) + return cc + + +def evaluate_image(gold: List[PolygonInstance], pred: List[PolygonInstance], + iou_thr: float = 0.5) -> Dict[int, ClassCount]: + """Per-class counts for one image.""" + by_class: Dict[int, ClassCount] = {} + classes = {i.class_id for i in gold} | {i.class_id for i in pred} + for c in classes: + g = [i for i in gold if i.class_id == c] + p = [i for i in pred if i.class_id == c] + by_class[c] = _match_class(g, p, iou_thr) + return by_class + + +def aggregate(per_image: List[Dict[int, ClassCount]]) -> Dict[int, ClassCount]: + """Sum per-image per-class counts across the dataset.""" + total: Dict[int, ClassCount] = {} + for img in per_image: + for c, cc in img.items(): + total.setdefault(c, ClassCount()).merge(cc) + return total + + +def class_report(cc: ClassCount) -> Dict[str, float]: + precision = cc.tp / (cc.tp + cc.fp) if (cc.tp + cc.fp) else 0.0 + recall = cc.tp / (cc.tp + cc.fn) if (cc.tp + cc.fn) else 0.0 + f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0 + mean_iou = float(np.mean(cc.mask_ious)) if cc.mask_ious else 0.0 + return { + "tp": cc.tp, "fp": cc.fp, "fn": cc.fn, + "precision": precision, "recall": recall, "f1": f1, + "mean_mask_iou": mean_iou, "n_matched": len(cc.mask_ious), + } + + +def full_report(total: Dict[int, ClassCount], names: List[str] | None = None) -> Dict: + """Per-class + micro-averaged overall report.""" + per_class = {} + overall = ClassCount() + for c, cc in sorted(total.items()): + name = names[c] if names and 0 <= c < len(names) else str(c) + per_class[name] = class_report(cc) + overall.merge(cc) + return {"per_class": per_class, "overall": class_report(overall)} diff --git a/src/eval/yolo_labels.py b/src/eval/yolo_labels.py new file mode 100644 index 0000000..884b5e8 --- /dev/null +++ b/src/eval/yolo_labels.py @@ -0,0 +1,123 @@ +"""Read/write Ultralytics YOLO-segmentation label files. + +A YOLO-seg label line is:: + + x1 y1 x2 y2 ... xn yn + +with every coordinate normalized to [0, 1] (x / image_width, y / image_height). +This module keeps points in normalized form (matching the on-disk format) and +offers helpers to convert to/from absolute pixels. +""" +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Tuple + +Point = Tuple[float, float] + + +@dataclass +class PolygonInstance: + """One segmented object: a class id and a single normalized polygon ring.""" + + class_id: int + points: List[Point] # normalized (x, y) in [0, 1] + class_name: Optional[str] = None + + def bbox(self) -> Tuple[float, float, float, float]: + """Axis-aligned bounding box (normalized) of the polygon.""" + xs = [p[0] for p in self.points] + ys = [p[1] for p in self.points] + return min(xs), min(ys), max(xs), max(ys) + + +def to_yolo_seg_line(inst: PolygonInstance) -> str: + coords = " ".join(f"{c:.6f}" for p in inst.points for c in p) + return f"{inst.class_id} {coords}".strip() + + +def parse_yolo_seg_line(line: str) -> Optional[PolygonInstance]: + parts = line.split() + if len(parts) < 7: # class id + at least 3 (x, y) pairs + return None + try: + class_id = int(float(parts[0])) + nums = [float(x) for x in parts[1:]] + except ValueError: + return None + pts = [(nums[i], nums[i + 1]) for i in range(0, len(nums) - 1, 2)] + if len(pts) < 3: + return None + return PolygonInstance(class_id=class_id, points=pts) + + +def read_label_file(path) -> List[PolygonInstance]: + """Parse a YOLO-seg ``.txt`` file into polygon instances (empty if missing).""" + p = Path(path) + if not p.exists(): + return [] + out: List[PolygonInstance] = [] + for line in p.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + inst = parse_yolo_seg_line(line) + if inst is not None: + out.append(inst) + return out + + +def write_label_file(path, instances: List[PolygonInstance]) -> None: + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + body = "\n".join(to_yolo_seg_line(i) for i in instances) + p.write_text(body + ("\n" if instances else ""), encoding="utf-8") + + +def abs_to_norm(points_abs, img_w: int, img_h: int) -> List[Point]: + return [(x / img_w, y / img_h) for (x, y) in points_abs] + + +def norm_to_abs(points_norm, img_w: int, img_h: int) -> List[Point]: + return [(x * img_w, y * img_h) for (x, y) in points_norm] + + +def load_class_names(dataset_yaml) -> List[str]: + """Read class names from a dataset.yaml ``names:`` block (best effort). + + Supports both the ``names: {0: a, 1: b}`` mapping and ``names: [a, b]`` list + forms without requiring PyYAML. + """ + p = Path(dataset_yaml) + if not p.exists(): + return [] + try: + import yaml # optional + data = yaml.safe_load(p.read_text(encoding="utf-8")) + names = data.get("names") if isinstance(data, dict) else None + if isinstance(names, dict): + return [names[k] for k in sorted(names)] + if isinstance(names, list): + return list(names) + except Exception: + pass + + # Minimal fallback parser for " 0: loop_all" style lines. + names_map = {} + in_names = False + for raw in p.read_text(encoding="utf-8").splitlines(): + if raw.strip().startswith("names:"): + in_names = True + continue + if in_names: + stripped = raw.strip() + if not stripped or (raw[:1] not in (" ", "\t") and ":" in stripped): + break + if ":" in stripped: + key, _, val = stripped.partition(":") + try: + names_map[int(key.strip().lstrip("- "))] = val.strip().strip("'\"") + except ValueError: + continue + return [names_map[k] for k in sorted(names_map)] if names_map else [] diff --git a/tests/test_import_studio_export.py b/tests/test_import_studio_export.py new file mode 100644 index 0000000..2ea566c --- /dev/null +++ b/tests/test_import_studio_export.py @@ -0,0 +1,66 @@ +"""tools/import_studio_export.py: images once in the pool, symlinked per dataset.""" +import json +from pathlib import Path + +import pytest + +from tools import import_studio_export as imp + +YAML = "path: /x\ntrain: train/images\nval: val/images\n\nnames:\n 0: loop_all\n 1: pin\n" + + +def make_export(root: Path, name: str, train=("tpi_1", "tpi_2"), val=("tpi_3",), missing_label=()): + exp = root / name + for split, keys in (("train", train), ("val", val)): + (exp / split / "images").mkdir(parents=True) + (exp / split / "labels").mkdir(parents=True) + for k in keys: + (exp / split / "images" / f"{k}.jpg").write_bytes(b"\xff\xd8" + k.encode() * 20) + if k not in missing_label: + (exp / split / "labels" / f"{k}.txt").write_text("0 0.1 0.1 0.2 0.1 0.2 0.2\n" if k != "tpi_2" else "") + (exp / "dataset.yaml").write_text(YAML) + (exp / "manifest.csv").write_text("key,tpid,split\n" + "\n".join(f"{k},{k[4:]},train" for k in train) + "\n") + return exp + + +def test_import_builds_pool_and_symlinks(tmp_path): + repo = tmp_path / "repo"; repo.mkdir() + exp = make_export(tmp_path, "export_a", missing_label=("tpi_3",)) + src = imp.import_export(exp, "2026-10-02_alc", root=repo) + ds = repo / "data" / "datasets" / "2026-10-02_alc" + pool = repo / "data" / "pool" / "images" + assert sorted(p.name for p in pool.iterdir()) == ["tpi_1.jpg", "tpi_2.jpg", "tpi_3.jpg"] + link = ds / "images" / "train" / "tpi_1.jpg" + assert link.is_symlink() and link.resolve() == (pool / "tpi_1.jpg").resolve() + assert (ds / "labels" / "train" / "tpi_1.txt").read_text().startswith("0 0.1") + assert (ds / "labels" / "val" / "tpi_3.txt").read_text() == "" # missing label -> background + assert src["counts"]["train"] == {"images": 2, "labels": 2, "empty": 1} + assert src["counts"]["val"] == {"images": 1, "labels": 0, "empty": 1} + assert src["new_pool_images"] == 3 and src["classes"] == {0: "loop_all", 1: "pin"} + assert "path: ." in (ds / "dataset.yaml").read_text() and (ds / "manifest.csv").exists() + assert json.loads((ds / "source.json").read_text())["name"] == "2026-10-02_alc" + + +def test_second_dataset_reuses_the_pool_and_names_are_immutable(tmp_path): + repo = tmp_path / "repo"; repo.mkdir() + imp.import_export(make_export(tmp_path, "a"), "2026-10-02_alc", root=repo) + src = imp.import_export(make_export(tmp_path, "b", train=("tpi_1", "tpi_2", "tpi_9"), val=("tpi_3",)), + "2026-10-09_alc", root=repo) + assert src["new_pool_images"] == 1 + assert len(list((repo / "data" / "pool" / "images").iterdir())) == 4 + with pytest.raises(SystemExit): + imp.import_export(make_export(tmp_path, "c"), "2026-10-02_alc", root=repo) + + +def test_gold_import(tmp_path): + repo = tmp_path / "repo"; repo.mkdir() + gold = tmp_path / "gold" + (gold / "images").mkdir(parents=True); (gold / "labels").mkdir() + (gold / "images" / "tpi_7.jpg").write_bytes(b"\xff\xd8gold") + (gold / "labels" / "tpi_7.txt").write_text("1 0.5 0.5 0.6 0.5 0.6 0.6\n") + (gold / "gold_manifest.json").write_text(json.dumps({"keys": ["tpi_7"]})) + src = imp.import_export(make_export(tmp_path, "a"), "2026-10-02_alc", root=repo, gold_dir=gold) + assert src["gold"]["images"] == 1 and src["gold"]["manifest_sha256"] + assert (repo / "data" / "gold" / "images" / "tpi_7.jpg").is_symlink() + assert (repo / "data" / "gold" / "labels" / "tpi_7.txt").exists() + assert (repo / "data" / "pool" / "images" / "tpi_7.jpg").exists() diff --git a/tests/test_ledger.py b/tests/test_ledger.py new file mode 100644 index 0000000..5e72110 --- /dev/null +++ b/tests/test_ledger.py @@ -0,0 +1,48 @@ +"""tools/ledger.py: runs as git-tracked JSON with fixed states.""" +import json + +import pytest + +from tools import ledger + + +def test_new_set_metrics_render_check(tmp_path): + run = ledger.new_run("dataset/2026-10-02_alc", "yolo26n-seg.pt", "alc", note="first", root=tmp_path) + assert run["status"] == "exported" and run["id"].endswith("_alc") + assert (tmp_path / "runs" / run["id"] / "run.json").exists() + + ledger.set_status(run["id"], "training", kv={"host": "spark"}, root=tmp_path) + again = ledger.load(run["id"], root=tmp_path) + assert again["status"] == "training" and again["host"] == "spark" and len(again["history"]) == 2 + + m = tmp_path / "m.json" + m.write_text(json.dumps({"overall": {"f1": 0.91, "mean_mask_iou": 0.8}, "gate": {"pass": True, "baseline": "b"}})) + ledger.record_metrics(run["id"], m, root=tmp_path) + ledger.set_status(run["id"], "gate_passed", root=tmp_path) + again = ledger.load(run["id"], root=tmp_path) + assert again["metrics"]["f1"] == 0.91 and again["metrics"]["gate_pass"] is True + assert (tmp_path / "runs" / run["id"] / "metrics.json").exists() + + readme = ledger.render(root=tmp_path) + text = readme.read_text() + assert run["id"] in text and "gate_passed" in text and "0.910" in text + assert ledger.check(root=tmp_path) == [] + assert run["id"] in ledger.table(root=tmp_path) + + +def test_second_run_same_day_gets_a_suffix(tmp_path): + a = ledger.new_run("dataset/x", "yolo26n-seg.pt", "alc", root=tmp_path) + b = ledger.new_run("dataset/x", "yolo26n-seg.pt", "alc", root=tmp_path) + assert b["id"] == a["id"] + "-2" + + +def test_unknown_state_and_check_problems(tmp_path): + run = ledger.new_run("dataset/x", "yolo26n-seg.pt", "t", root=tmp_path) + with pytest.raises(SystemExit): + ledger.set_status(run["id"], "done", root=tmp_path) + # measured without a metrics.json is a problem the CI check reports + p = tmp_path / "runs" / run["id"] / "run.json" + d = json.loads(p.read_text()); d["status"] = "measured"; p.write_text(json.dumps(d)) + problems = ledger.check(root=tmp_path) + assert any("no metrics.json" in x for x in problems) + assert ledger.main(["--root", str(tmp_path), "check"]) == 1 diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/import_studio_export.py b/tools/import_studio_export.py new file mode 100755 index 0000000..e510203 --- /dev/null +++ b/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// + ├── images/{train,val}/tpi_.jpg -> ../../../../pool/images/tpi_.jpg + ├── labels/{train,val}/tpi_.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 = _") + ap.add_argument("--name", default=None, help="full dataset name instead of _") + 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/") + 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()) diff --git a/tools/ledger.py b/tools/ledger.py new file mode 100755 index 0000000..26b151d --- /dev/null +++ b/tools/ledger.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 +"""Run ledger: one directory per training run under ``runs/``, tracked in git. + +Every run is ``runs//run.json`` with a status that moves through fixed states; +each change appends to ``history`` and is meant to be committed, so +``git log runs//run.json`` is the timeline of that run. ``metrics.json`` next +to it holds the gold measure and the gate verdict, ``train.log`` the trainer's +output. + + python tools/ledger.py new --dataset dataset/2026-10-02_alc --tag alc --base yolo26n-seg.pt + python tools/ledger.py set 2026-10-02_alc training --kv host=mx-spark-01 job=12345 + python tools/ledger.py metrics 2026-10-02_alc path/to/metrics.json + python tools/ledger.py show + python tools/ledger.py render # runs/README.md + python tools/ledger.py check # CI: every run.json valid +""" + +from __future__ import annotations + +import argparse +import getpass +import json +import shutil +import socket +import sys +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional + +ROOT = Path(__file__).resolve().parents[1] +STATES = [ + "exported", # dataset tag exists, run directory created + "queued", # job submitted (Spark or cluster) + "training", # epochs running, train.log grows + "measured", # gold measure done, metrics.json written + "gate_passed", # measure beats the deployed model's baseline + "gate_failed", + "promoted", # models/active/production.json points at this run's weights + "deployed", # AareLC_infer tagged, cps-deployments bumped + "aborted", +] +REQUIRED = ("id", "dataset", "base_weights", "status", "created", "history") + + +def _now() -> str: + return datetime.now().isoformat(timespec="seconds") + + +def _who() -> str: + try: + return f"{getpass.getuser()}@{socket.gethostname()}" + except Exception: + return "unknown" + + +def runs_dir(root: Path = ROOT) -> Path: + return root / "runs" + + +def run_path(run_id: str, root: Path = ROOT) -> Path: + return runs_dir(root) / run_id / "run.json" + + +def load(run_id: str, root: Path = ROOT) -> Dict[str, Any]: + p = run_path(run_id, root) + if not p.exists(): + raise SystemExit(f"no such run: {run_id} ({p})") + return json.loads(p.read_text(encoding="utf-8")) + + +def save(run: Dict[str, Any], root: Path = ROOT) -> Path: + p = run_path(run["id"], root) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(run, indent=2, sort_keys=False) + "\n", encoding="utf-8") + return p + + +def _parse_kv(items: Optional[List[str]]) -> Dict[str, Any]: + out: Dict[str, Any] = {} + for it in items or []: + if "=" not in it: + raise SystemExit(f"expected key=value, got {it!r}") + k, v = it.split("=", 1) + try: + out[k] = json.loads(v) + except json.JSONDecodeError: + out[k] = v + return out + + +def new_run(dataset: str, base_weights: str, tag: str, *, note: str = "", root: Path = ROOT, + extra: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + stamp = datetime.now().strftime("%Y-%m-%d") + run_id = f"{stamp}_{tag}" + n = 2 + while run_path(run_id, root).exists(): + run_id = f"{stamp}_{tag}-{n}" + n += 1 + run = { + "id": run_id, + "dataset": dataset, + "base_weights": base_weights, + "status": "exported", + "created": _now(), + "created_by": _who(), + "hyperparams": {}, + "version": None, # models/registry/ once trained + "metrics": {}, # summary of metrics.json + "history": [{"status": "exported", "at": _now(), "by": _who(), "note": note}], + } + run.update(extra or {}) + save(run, root) + return run + + +def set_status(run_id: str, status: str, *, note: str = "", kv: Optional[Dict[str, Any]] = None, + root: Path = ROOT) -> Dict[str, Any]: + if status not in STATES: + raise SystemExit(f"unknown state {status!r}; one of {', '.join(STATES)}") + run = load(run_id, root) + run["status"] = status + for k, v in (kv or {}).items(): + run[k] = v + run["history"].append({"status": status, "at": _now(), "by": _who(), "note": note}) + save(run, root) + return run + + +def record_metrics(run_id: str, metrics_file: Path, root: Path = ROOT) -> Dict[str, Any]: + run = load(run_id, root) + data = json.loads(Path(metrics_file).read_text(encoding="utf-8")) + dest = run_path(run_id, root).parent / "metrics.json" + if Path(metrics_file).resolve() != dest.resolve(): + shutil.copy2(metrics_file, dest) + overall = data.get("overall") or {} + gate = data.get("gate") or {} + run["metrics"] = { + "f1": overall.get("f1"), + "mean_mask_iou": overall.get("mean_mask_iou"), + "gate_pass": gate.get("pass"), + "baseline": gate.get("baseline"), + "measured_at": _now(), + } + save(run, root) + return run + + +def all_runs(root: Path = ROOT) -> List[Dict[str, Any]]: + out = [] + for p in sorted(runs_dir(root).glob("*/run.json")): + try: + out.append(json.loads(p.read_text(encoding="utf-8"))) + except json.JSONDecodeError as e: + out.append({"id": p.parent.name, "status": f"INVALID JSON: {e}", "history": []}) + return out + + +def check(root: Path = ROOT) -> List[str]: + problems: List[str] = [] + for p in sorted(runs_dir(root).glob("*/run.json")): + try: + run = json.loads(p.read_text(encoding="utf-8")) + except json.JSONDecodeError as e: + problems.append(f"{p}: invalid JSON ({e})") + continue + for k in REQUIRED: + if k not in run: + problems.append(f"{p}: missing {k}") + if run.get("id") != p.parent.name: + problems.append(f"{p}: id {run.get('id')!r} does not match directory") + if run.get("status") not in STATES: + problems.append(f"{p}: unknown status {run.get('status')!r}") + for h in run.get("history", []): + if h.get("status") not in STATES: + problems.append(f"{p}: history entry with unknown status {h.get('status')!r}") + if run.get("status") in ("measured", "gate_passed", "gate_failed", "promoted", "deployed"): + if not (p.parent / "metrics.json").exists(): + problems.append(f"{p}: status {run['status']} but no metrics.json") + return problems + + +def _fmt(v: Any) -> str: + if v is None: + return "-" + if isinstance(v, float): + return f"{v:.3f}" + if isinstance(v, bool): + return "yes" if v else "no" + return str(v) + + +def table(root: Path = ROOT) -> str: + rows = [("run", "status", "dataset", "base", "version", "f1", "mask IoU", "gate", "last change")] + for r in all_runs(root): + m = r.get("metrics") or {} + last = (r.get("history") or [{}])[-1].get("at", "") + rows.append((r.get("id", "?"), r.get("status", "?"), r.get("dataset", "-"), r.get("base_weights", "-"), + _fmt(r.get("version")), _fmt(m.get("f1")), _fmt(m.get("mean_mask_iou")), + _fmt(m.get("gate_pass")), last[:16])) + widths = [max(len(str(row[i])) for row in rows) for i in range(len(rows[0]))] + lines = [] + for n, row in enumerate(rows): + lines.append(" ".join(str(c).ljust(widths[i]) for i, c in enumerate(row))) + if n == 0: + lines.append(" ".join("-" * w for w in widths)) + return "\n".join(lines) + + +def render(root: Path = ROOT) -> Path: + runs = all_runs(root) + md = ["# Training runs", "", "Generated by `python tools/ledger.py render`; do not edit by hand.", "", + "| Run | Status | Dataset | Base | Version | F1 | Mask IoU | Gate | Last change |", + "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"] + for r in runs: + m = r.get("metrics") or {} + last = (r.get("history") or [{}])[-1].get("at", "") + md.append(f"| {r.get('id')} | {r.get('status')} | {r.get('dataset', '-')} | {r.get('base_weights', '-')} | " + f"{_fmt(r.get('version'))} | {_fmt(m.get('f1'))} | {_fmt(m.get('mean_mask_iou'))} | " + f"{_fmt(m.get('gate_pass'))} | {last[:16]} |") + md.append("") + md.append("States: " + " → ".join(STATES[:-1]) + f" (or {STATES[-1]}).") + md.append("") + for r in runs: + md.append(f"## {r.get('id')}") + md.append("") + for h in r.get("history", []): + note = f" — {h['note']}" if h.get("note") else "" + md.append(f"- {h.get('at', '')[:19]} · {h.get('status')} · {h.get('by', '')}{note}") + md.append("") + out = runs_dir(root) / "README.md" + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text("\n".join(md), encoding="utf-8") + return out + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--root", default=str(ROOT), help="repository root (default: this checkout)") + sub = ap.add_subparsers(dest="cmd", required=True) + n = sub.add_parser("new", help="create a run in state 'exported'") + n.add_argument("--dataset", required=True, help="dataset tag, e.g. dataset/2026-10-02_alc") + n.add_argument("--tag", required=True, help="short name; run id = _") + n.add_argument("--base", default="yolo26n-seg.pt", help="base weights the run starts from") + n.add_argument("--note", default="") + n.add_argument("--kv", nargs="*", help="extra fields key=value (JSON values accepted)") + s = sub.add_parser("set", help="move a run to a state") + s.add_argument("run_id") + s.add_argument("status", choices=STATES) + s.add_argument("--note", default="") + s.add_argument("--kv", nargs="*", help="extra fields key=value") + m = sub.add_parser("metrics", help="record a metrics.json on a run") + m.add_argument("run_id") + m.add_argument("metrics_file") + sub.add_parser("show", help="print the table") + sub.add_parser("render", help="write runs/README.md") + sub.add_parser("check", help="validate every run.json (exit 1 on problems)") + args = ap.parse_args(argv) + root = Path(args.root) + + if args.cmd == "new": + run = new_run(args.dataset, args.base, args.tag, note=args.note, root=root, extra=_parse_kv(args.kv)) + print(run["id"]) + elif args.cmd == "set": + run = set_status(args.run_id, args.status, note=args.note, kv=_parse_kv(args.kv), root=root) + print(f"{run['id']}: {run['status']}") + elif args.cmd == "metrics": + run = record_metrics(args.run_id, Path(args.metrics_file), root=root) + print(json.dumps(run["metrics"])) + elif args.cmd == "show": + print(table(root)) + elif args.cmd == "render": + print(render(root)) + elif args.cmd == "check": + problems = check(root) + for p in problems: + print(p, file=sys.stderr) + print(f"{len(list(runs_dir(root).glob('*/run.json')))} runs, {len(problems)} problems") + return 1 if problems else 0 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/train_run.py b/tools/train_run.py new file mode 100644 index 0000000..ed27864 --- /dev/null +++ b/tools/train_run.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +"""Train one ledger run end to end: train → registry → gold measure → gate → promote. + + python tools/ledger.py new --dataset dataset/2026-10-02_alc --tag alc --base yolo26n-seg.pt + python tools/train_run.py 2026-10-02_alc --device 0 2>&1 | tee runs/2026-10-02_alc/train.log + +Every retraining starts from the base weights on the whole dataset (the cumulative +labelled pool), so runs stay comparable. The gate compares the gold measure with the +deployed model's own gold report (``models/registry//gold_report.json``): +a pass needs no drop beyond 0.005 in overall F1 and mask IoU, an improvement in at +least one of them, and no drop beyond 0.02 in F1 for loop_all or pin. On a pass the +run is promoted on its own: ``models/active/production.json`` is rewritten. The +first run, with no baseline, only records its report as the future baseline. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import shutil +import socket +import sys +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Optional + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from tools import ledger # noqa: E402 + +AUGMENT = dict(scale=0.7, translate=0.15, mosaic=0.5, close_mosaic=10, fliplr=0.5, flipud=0.3, + shear=5.0, perspective=0.0) +GATE_TOL = 0.005 +CLASS_TOL = 0.02 +GUARDED_CLASSES = ("loop_all", "pin") + + +def dataset_dir(run: Dict[str, Any]) -> Path: + name = str(run["dataset"]).split("/", 1)[-1] + d = ROOT / "data" / "datasets" / name + if not (d / "dataset.yaml").exists(): + raise SystemExit(f"dataset {run['dataset']} not found at {d} (git checkout the tag, git lfs pull)") + return d + + +def resolved_yaml(run_dir: Path, ds: Path) -> Path: + text = (ds / "dataset.yaml").read_text(encoding="utf-8").replace("path: .", f"path: {ds.resolve()}", 1) + out = run_dir / "dataset.resolved.yaml" + out.write_text(text, encoding="utf-8") + return out + + +def base_weights_path(spec: str) -> str: + p = Path(spec) + if p.exists(): + return str(p) + cand = ROOT / "models" / "base" / spec + return str(cand) if cand.exists() else spec # else an Ultralytics model name, downloaded + + +def results_summary(train_dir: Path) -> Dict[str, Any]: + csv_path = train_dir / "results.csv" + if not csv_path.exists(): + return {} + rows = list(csv.DictReader(csv_path.open(encoding="utf-8"))) + if not rows: + return {} + last = {k.strip(): v for k, v in rows[-1].items()} + out: Dict[str, Any] = {"epochs_run": len(rows)} + for k, v in last.items(): + if "metrics/" in k: + try: + out[k.replace("metrics/", "")] = float(v) + except ValueError: + pass + return out + + +def measure(weights: Path, run_dir: Path, names: list, conf: float = 0.25) -> Dict[str, Any]: + from src.eval.evaluate import evaluate_dirs + from src.eval.yolo_labels import PolygonInstance, write_label_file + from ultralytics import YOLO + + gold = ROOT / "data" / "gold" + img_dir, gold_labels = gold / "images", gold / "labels" + if not img_dir.exists() or not gold_labels.exists(): + raise SystemExit("no gold set under data/gold (import it with tools/import_studio_export.py --gold)") + pred_dir = run_dir / "gold_predictions" + pred_dir.mkdir(parents=True, exist_ok=True) + model = YOLO(str(weights)) + for img_path in sorted(p for p in img_dir.iterdir() if p.suffix.lower() in (".jpg", ".jpeg", ".png")): + res = model.predict(str(img_path), conf=conf, verbose=False)[0] + insts = [] + if res.masks is not None: + h, w = res.orig_shape + for poly, cls in zip(res.masks.xy, res.boxes.cls.tolist()): + pts = [(float(x) / w, float(y) / h) for x, y in poly] + if len(pts) >= 3: + insts.append(PolygonInstance(int(cls), pts)) + write_label_file(pred_dir / f"{img_path.stem}.txt", insts) + return evaluate_dirs(pred_dir, gold_labels, names=names) + + +def gate(report: Dict[str, Any], baseline: Optional[Path]) -> Optional[Dict[str, Any]]: + if not baseline or not Path(baseline).exists(): + return None + base = json.loads(Path(baseline).read_text(encoding="utf-8")) + f1_delta = float(report["overall"]["f1"]) - float(base["overall"]["f1"]) + iou_delta = float(report["overall"]["mean_mask_iou"]) - float(base["overall"]["mean_mask_iou"]) + per_class = report.get("per_class") or report.get("classes") or {} + base_classes = base.get("per_class") or base.get("classes") or {} + class_deltas = {} + for c in GUARDED_CLASSES: + a, b = per_class.get(c) or {}, base_classes.get(c) or {} + if "f1" in a and "f1" in b: + class_deltas[c] = float(a["f1"]) - float(b["f1"]) + ok = (f1_delta >= -GATE_TOL and iou_delta >= -GATE_TOL and (f1_delta > 0 or iou_delta > 0) + and all(d >= -CLASS_TOL for d in class_deltas.values())) + return {"baseline": str(baseline), "f1_delta": round(f1_delta, 4), "mask_iou_delta": round(iou_delta, 4), + "class_f1_delta": {k: round(v, 4) for k, v in class_deltas.items()}, "pass": bool(ok)} + + +def active_version() -> Optional[str]: + p = ROOT / "models" / "active" / "production.json" + if not p.exists(): + return None + return json.loads(p.read_text(encoding="utf-8")).get("version") + + +def promote(version: str, run_id: str, g: Dict[str, Any]) -> Path: + p = ROOT / "models" / "active" / "production.json" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps({ + "version": version, + "export": "weights.pt", # engine built on the target GPU at deploy time + "promoted_on": datetime.now().strftime("%Y-%m-%d"), + "promoted_by": "gate", + "run": run_id, + "gate": g, + }, indent=2) + "\n", encoding="utf-8") + return p + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("run_id") + ap.add_argument("--epochs", type=int, default=400) + ap.add_argument("--patience", type=int, default=50) + ap.add_argument("--imgsz", type=int, default=640) + ap.add_argument("--batch", type=int, default=16) + ap.add_argument("--device", default="0") + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--conf", type=float, default=0.25, help="confidence used for the gold measure") + ap.add_argument("--baseline", default=None, help="gold report to gate against (default: the active model's)") + ap.add_argument("--no-promote", action="store_true", help="gate, but leave production.json alone") + ap.add_argument("--skip-train", action="store_true", help="reuse runs//ultralytics/train/weights/best.pt") + args = ap.parse_args(argv) + + run = ledger.load(args.run_id) + run_dir = ledger.run_path(args.run_id).parent + ds = dataset_dir(run) + data_yaml = resolved_yaml(run_dir, ds) + names = [v for _, v in sorted(json.loads((ds / "source.json").read_text())["classes"].items(), key=lambda kv: int(kv[0]))] \ + if (ds / "source.json").exists() else ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"] + hyper = {"epochs": args.epochs, "patience": args.patience, "imgsz": args.imgsz, "batch": args.batch, + "device": args.device, "overlap_mask": False, **AUGMENT} + ledger.set_status(args.run_id, "training", note="train_run.py", kv={ + "hyperparams": hyper, "host": socket.gethostname(), "started": datetime.now().isoformat(timespec="seconds")}) + + best = run_dir / "ultralytics" / "train" / "weights" / "best.pt" + if not args.skip_train: + from ultralytics import YOLO + model = YOLO(base_weights_path(run["base_weights"])) + model.train(data=str(data_yaml), epochs=args.epochs, patience=args.patience, imgsz=args.imgsz, + batch=args.batch, device=args.device, workers=args.workers, overlap_mask=False, + project=str(run_dir / "ultralytics"), name="train", exist_ok=True, amp=False, **AUGMENT) + if not best.exists(): + ledger.set_status(args.run_id, "aborted", note=f"no best.pt at {best}") + return 1 + + family = Path(run["base_weights"]).stem.split("-")[0] + tag = args.run_id.split("_", 1)[-1] + version = f"{family}-seg-{tag}-{datetime.now().strftime('%Y-%m-%d')}" + reg = ROOT / "models" / "registry" / version + reg.mkdir(parents=True, exist_ok=True) + shutil.copy2(best, reg / "weights.pt") + + report = measure(reg / "weights.pt", run_dir, names, conf=args.conf) + baseline = Path(args.baseline) if args.baseline else None + if baseline is None: + av = active_version() + if av and (ROOT / "models" / "registry" / av / "gold_report.json").exists(): + baseline = ROOT / "models" / "registry" / av / "gold_report.json" + g = gate(report, baseline) + report["gate"] = g + (reg / "gold_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + metrics = {"overall": report["overall"], "gate": g, "training": results_summary(run_dir / "ultralytics" / "train"), + "version": version, "measured_at": datetime.now().isoformat(timespec="seconds")} + (run_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") + (reg / "metadata.json").write_text(json.dumps({ + "model_id": version, "family": family, "task": "segmentation", + "training_date": datetime.now().strftime("%Y-%m-%d"), "run": args.run_id, + "dataset": run["dataset"], "base_weights": run["base_weights"], "hyperparams": hyper, + "export_formats": ["pt"], "metrics": metrics["overall"], "gate": g, + }, indent=2) + "\n", encoding="utf-8") + ledger.record_metrics(args.run_id, run_dir / "metrics.json") + ledger.set_status(args.run_id, "measured", kv={"version": version}) + + if g is None: + print(f"measured; no baseline, {version} is the baseline for the next run") + return 0 + if g["pass"]: + ledger.set_status(args.run_id, "gate_passed", note=json.dumps(g)) + if not args.no_promote: + promote(version, args.run_id, g) + ledger.set_status(args.run_id, "promoted", note="by the gate") + print(f"PROMOTED {version}: {g}") + return 0 + ledger.set_status(args.run_id, "gate_failed", note=json.dumps(g)) + print(f"gate failed: {g}") + return 2 + + +if __name__ == "__main__": + raise SystemExit(main())