experiments: loop-type classifier, first cross-validated test (broken loops AUROC 0.985)

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

experiments/looptype/README.md

@@ -0,0 +1,20 @@
+# Loop-type classifier, first test (2026-09-29)
+
+Question: can a small classifier tell the loop type, and above all a broken loop, from the
+frame? Data: the 364 labelled production frames whose reviewer picked a loop type in the
+annotator (nylon 123, broken 97, litholoop 88, mitegen microloop 30, micromesh 26).
+Method: `train.py` (a copy of `AareLC_studio/src/pipeline/looptype/train.py`): ResNet-18
+from ImageNet weights, 256 px crop around every labelled object widened by 40 %, five
+cross-validation folds that never split the frames of one sample, class-weighted loss.
+Result: `looptype_head_2026-09-29.json`.
+
+| | |
+| --- | --- |
+| accuracy, five classes | 0.849 |
+| broken: precision / recall | 0.88 / 0.95 |
+| broken vs rest AUROC | 0.985 overall, 0.993 within PXII, 0.81 within PXIII (4 broken frames only) |
+| broken recall at 90 % precision | 0.95 |
+| weakest classes | micromesh and mitegen microloop, confused with each other (26 and 30 frames) |
+
+Caveat: 93 of the 97 broken tags come from PXII, so the PXIII number rests on four
+frames; the next tags to collect are broken loops from PXIII.

experiments/looptype/looptype_head_2026-09-29.json

@@ -0,0 +1,88 @@
+{
+  "n": 364,
+  "accuracy": 0.8489010989010989,
+  "per_class": {
+    "broken": {
+      "n": 97,
+      "precision": 0.8846153846153846,
+      "recall": 0.9484536082474226
+    },
+    "litholoop": {
+      "n": 88,
+      "precision": 0.8690476190476191,
+      "recall": 0.8295454545454546
+    },
+    "micromesh": {
+      "n": 26,
+      "precision": 0.68,
+      "recall": 0.6538461538461539
+    },
+    "mitegen_microloop": {
+      "n": 30,
+      "precision": 0.6785714285714286,
+      "recall": 0.6333333333333333
+    },
+    "nylon": {
+      "n": 123,
+      "precision": 0.8780487804878049,
+      "recall": 0.8780487804878049
+    }
+  },
+  "confusion": {
+    "broken": {
+      "broken": 92,
+      "litholoop": 2,
+      "micromesh": 0,
+      "mitegen_microloop": 1,
+      "nylon": 2
+    },
+    "litholoop": {
+      "broken": 0,
+      "litholoop": 73,
+      "micromesh": 2,
+      "mitegen_microloop": 2,
+      "nylon": 11
+    },
+    "micromesh": {
+      "broken": 1,
+      "litholoop": 3,
+      "micromesh": 17,
+      "mitegen_microloop": 5,
+      "nylon": 0
+    },
+    "mitegen_microloop": {
+      "broken": 2,
+      "litholoop": 1,
+      "micromesh": 6,
+      "mitegen_microloop": 19,
+      "nylon": 2
+    },
+    "nylon": {
+      "broken": 9,
+      "litholoop": 5,
+      "micromesh": 0,
+      "mitegen_microloop": 1,
+      "nylon": 108
+    }
+  },
+  "broken_vs_rest": {
+    "auroc": 0.9854434534151898,
+    "n_broken": 97,
+    "recall_at_90_precision": 0.9484536082474226,
+    "auroc_PXII": 0.9927012056044314,
+    "auroc_PXIII": 0.8112745098039216
+  },
+  "classes": [
+    "broken",
+    "litholoop",
+    "micromesh",
+    "mitegen_microloop",
+    "nylon"
+  ],
+  "size": 256,
+  "epochs": 15,
+  "folds": 5,
+  "groups": 308,
+  "trained_at": "2026-09-29T18:17:44",
+  "config": "pipeline.prod.json"
+}

experiments/looptype/train.py

@@ -0,0 +1,228 @@
+"""Can a small classifier tell the loop type, and above all a broken loop, from the frame?
+
+Data: every labelled pool frame whose reviewer picked a loop type (the tag next to Y in the
+annotator). Input: the region around the mount, the box around every labelled object
+widened by 40 % (a broken loop often has no loop polygon, but always a pin), letterboxed
+to a square. Model: ResNet-18 from ImageNet weights, one head over the loop types.
+
+Honest numbers: grouped cross-validation, frames of the same sample never split across
+folds; per class precision and recall; broken-versus-rest AUROC and recall at 90 %
+precision; and the same within one beamline, because 96 % of the broken tags come from
+PXII and a classifier could learn the camera instead of the loop.
+
+    python -m src.pipeline.looptype.train --config pipeline.prod.json --out models/looptype_head.pt
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from collections import Counter, defaultdict
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+
+
+def group_key(item: Dict) -> str:
+    """Frames of one sample share a name up to the event suffix: keep them in one fold."""
+    name = str(item.get("filename") or item.get("tpid"))
+    name = re.sub(r"\.(jpe?g|png)$", "", name, flags=re.I)
+    name = re.sub(r"_(pre_alc|post_alc|mounted|centered|before_raster|post_raster.*|scan.*|-?\d+deg.*)$", "", name, flags=re.I)
+    return name.rsplit("_", 1)[0] if "_" in name else name
+
+
+def crop_region(image: np.ndarray, item: Dict, margin: float = 0.4, min_side: int = 320) -> np.ndarray:
+    h, w = image.shape[:2]
+    xs: List[float] = []; ys: List[float] = []
+    for s in item.get("segments") or []:
+        pts = s.get("points") or []
+        if len(pts) >= 6:
+            xs += [float(p) * w for p in pts[0::2]]; ys += [float(p) * h for p in pts[1::2]]
+        elif s.get("box") and len(s["box"]) == 4:
+            cx, cy, bw, bh = s["box"]
+            if max(cx, cy, bw, bh) <= 1.5:
+                cx, cy, bw, bh = cx * w, cy * h, bw * w, bh * h
+            xs += [cx - bw / 2, cx + bw / 2]; ys += [cy - bh / 2, cy + bh / 2]
+    if not xs:
+        return image
+    x0, x1, y0, y1 = min(xs), max(xs), min(ys), max(ys)
+    side = max(x1 - x0, y1 - y0) * (1 + 2 * margin)
+    side = max(side, min_side)
+    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
+    X0, Y0 = int(max(cx - side / 2, 0)), int(max(cy - side / 2, 0))
+    X1, Y1 = int(min(cx + side / 2, w)), int(min(cy + side / 2, h))
+    return image[Y0:Y1, X0:X1]
+
+
+def letterbox(image: np.ndarray, size: int) -> np.ndarray:
+    h, w = image.shape[:2]
+    s = size / max(h, w)
+    img = cv2.resize(image, (max(int(w * s), 1), max(int(h * s), 1)), interpolation=cv2.INTER_AREA)
+    out = np.full((size, size, 3), 114, dtype=np.uint8)
+    y, x = (size - img.shape[0]) // 2, (size - img.shape[1]) // 2
+    out[y:y + img.shape[0], x:x + img.shape[1]] = img
+    return out
+
+
+def load_records(cfg_path: str, cache_root: Path, log=print) -> List[Dict]:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
+    from src.pipeline import run as pipeline_run
+    from src.pipeline.config import PipelineConfig
+    cfg = PipelineConfig.load(cfg_path)
+    items = [it for it in pipeline_run._all_labels(cfg, include_empty=True) if it.get("loop_type")]
+    keys = [f"tpi_{int(it['tpid'])}" for it in items]
+    lookup = pipeline_run._image_lookup(cfg, keys)
+    out = []
+    for it in items:
+        p = lookup.get(f"tpi_{int(it['tpid'])}")
+        if p and Path(p).exists():
+            out.append({**it, "image": str(p)})
+    log(f"looptype: {len(out)} tagged frames with an image, {dict(Counter(it['loop_type'] for it in out))}")
+    return out
+
+
+def build_arrays(records: List[Dict], size: int) -> Tuple[np.ndarray, np.ndarray, List[str], List[str]]:
+    classes = sorted({r["loop_type"] for r in records})
+    X = np.zeros((len(records), size, size, 3), dtype=np.uint8); y = np.zeros(len(records), dtype=np.int64)
+    for i, r in enumerate(records):
+        img = cv2.imread(r["image"])
+        X[i] = letterbox(crop_region(img, r), size)
+        y[i] = classes.index(r["loop_type"])
+    groups = [group_key(r) for r in records]
+    return X, y, classes, groups
+
+
+def grouped_folds(groups: List[str], y: np.ndarray, k: int, seed: int) -> List[np.ndarray]:
+    """Greedy stratified group split: whole groups to folds, balancing class counts."""
+    rng = np.random.RandomState(seed)
+    by_group: Dict[str, List[int]] = defaultdict(list)
+    for i, g in enumerate(groups):
+        by_group[g].append(i)
+    order = list(by_group); rng.shuffle(order)
+    order.sort(key=lambda g: -len(by_group[g]))
+    fold_counts = [Counter() for _ in range(k)]; folds = [[] for _ in range(k)]
+    for g in order:
+        idx = by_group[g]; c = Counter(int(y[i]) for i in idx)
+        best = min(range(k), key=lambda f: sum(fold_counts[f][cl] * n for cl, n in c.items()) + len(folds[f]) * 0.01)
+        folds[best] += idx; fold_counts[best].update(c)
+    return [np.array(sorted(f)) for f in folds]
+
+
+def make_model(n_classes: int, pretrained: bool = True):
+    import torch.nn as nn
+    import torchvision
+    weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None
+    m = torchvision.models.resnet18(weights=weights)
+    m.fc = nn.Linear(m.fc.in_features, n_classes)
+    return m
+
+
+def train_eval(X, y, train_idx, val_idx, n_classes, device, epochs, batch, lr, seed):
+    import torch
+    import torch.nn as nn
+    torch.manual_seed(seed)
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32); std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+    def to_tensor(idx, augment=False):
+        imgs = X[idx][..., ::-1].astype(np.float32) / 255.0   # BGR -> RGB
+        if augment:
+            flip = np.random.rand(len(idx)) < 0.5
+            imgs[flip] = imgs[flip][:, ::-1]          # vertical flip only: the mount orientation is fixed
+            imgs = np.clip(imgs * np.random.uniform(0.8, 1.2, (len(idx), 1, 1, 1)), 0, 1)
+        imgs = ((imgs - mean) / std).astype(np.float32)
+        return torch.from_numpy(np.ascontiguousarray(imgs.transpose(0, 3, 1, 2))).to(device)
+
+    model = make_model(n_classes).to(device)
+    counts = np.bincount(y[train_idx], minlength=n_classes).astype(np.float32)
+    weight = torch.tensor((counts.sum() / np.maximum(counts, 1)) / n_classes, dtype=torch.float32, device=device)
+    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
+    loss_fn = nn.CrossEntropyLoss(weight=weight)
+    for ep in range(epochs):
+        model.train(); perm = np.random.permutation(train_idx)
+        for i in range(0, len(perm), batch):
+            idx = perm[i:i + batch]
+            opt.zero_grad(); out = model(to_tensor(idx, augment=True)); loss = loss_fn(out, torch.from_numpy(y[idx]).to(device))
+            loss.backward(); opt.step()
+        sched.step()
+    model.eval(); probs = []
+    with torch.no_grad():
+        for i in range(0, len(val_idx), batch):
+            idx = val_idx[i:i + batch]
+            probs.append(torch.softmax(model(to_tensor(idx)), dim=1).cpu().numpy())
+    return model, np.concatenate(probs) if probs else np.zeros((0, n_classes))
+
+
+def auroc(scores: np.ndarray, labels: np.ndarray) -> Optional[float]:
+    pos, neg = scores[labels == 1], scores[labels == 0]
+    if len(pos) == 0 or len(neg) == 0:
+        return None
+    ranks = np.argsort(np.argsort(np.concatenate([pos, neg]))) + 1
+    return float((ranks[:len(pos)].sum() - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg)))
+
+
+def report(y_true, probs, classes, beamlines) -> Dict:
+    pred = probs.argmax(1)
+    out: Dict = {"n": int(len(y_true)), "accuracy": float((pred == y_true).mean()), "per_class": {}, "confusion": {}}
+    for ci, c in enumerate(classes):
+        tp = int(((pred == ci) & (y_true == ci)).sum()); fp = int(((pred == ci) & (y_true != ci)).sum()); fn = int(((pred != ci) & (y_true == ci)).sum())
+        out["per_class"][c] = {"n": int((y_true == ci).sum()), "precision": tp / (tp + fp) if tp + fp else 0.0, "recall": tp / (tp + fn) if tp + fn else 0.0}
+        out["confusion"][c] = {classes[cj]: int(((y_true == ci) & (pred == cj)).sum()) for cj in range(len(classes))}
+    if "broken" in classes:
+        bi = classes.index("broken"); s = probs[:, bi]; lab = (y_true == bi).astype(int)
+        out["broken_vs_rest"] = {"auroc": auroc(s, lab), "n_broken": int(lab.sum())}
+        # recall at 90 % precision, if reachable
+        order = np.argsort(-s); tp = np.cumsum(lab[order]); prec = tp / (np.arange(len(s)) + 1); rec = tp / max(lab.sum(), 1)
+        ok = np.where(prec >= 0.9)[0]
+        out["broken_vs_rest"]["recall_at_90_precision"] = float(rec[ok].max()) if len(ok) else 0.0
+        for bl in sorted(set(beamlines)):
+            m = np.array([b == bl for b in beamlines])
+            if lab[m].sum() >= 3 and (1 - lab[m]).sum() >= 3:
+                out["broken_vs_rest"][f"auroc_{bl}"] = auroc(s[m], lab[m])
+    return out
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--config", default="pipeline.prod.json")
+    ap.add_argument("--out", default="models/looptype_head.pt")
+    ap.add_argument("--size", type=int, default=256)
+    ap.add_argument("--epochs", type=int, default=15)
+    ap.add_argument("--batch", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--folds", type=int, default=5)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--device", default=None)
+    args = ap.parse_args(argv)
+    import torch
+    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
+    records = load_records(args.config, Path("data"))
+    X, y, classes, groups = build_arrays(records, args.size)
+    beamlines = [str(r.get("beamline")) for r in records]
+    folds = grouped_folds(groups, y, args.folds, args.seed)
+    all_probs = np.zeros((len(y), len(classes)), dtype=np.float32)
+    for k, val_idx in enumerate(folds):
+        train_idx = np.array(sorted(set(range(len(y))) - set(val_idx.tolist())))
+        _, probs = train_eval(X, y, train_idx, val_idx, len(classes), device, args.epochs, args.batch, args.lr, args.seed + k)
+        all_probs[val_idx] = probs
+        print(f"fold {k + 1}/{args.folds}: {len(train_idx)} train, {len(val_idx)} val, acc {(probs.argmax(1) == y[val_idx]).mean():.3f}", flush=True)
+    rep = report(y, all_probs, classes, beamlines)
+    rep.update({"classes": classes, "size": args.size, "epochs": args.epochs, "folds": args.folds, "groups": len(set(groups)),
+                "trained_at": datetime.now().isoformat(timespec="seconds"), "config": args.config})
+    print(json.dumps(rep, indent=2))
+    # final model on everything, for use in the loop later
+    model, _ = train_eval(X, y, np.arange(len(y)), np.arange(0), len(classes), device, args.epochs, args.batch, args.lr, args.seed)
+    out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"state_dict": model.state_dict(), "classes": classes, "size": args.size, "report": rep}, out)
+    out.with_suffix(".json").write_text(json.dumps(rep, indent=2), encoding="utf-8")
+    print(f"saved {out} and {out.with_suffix('.json')}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())