diff --git a/experiments/looptype/README.md b/experiments/looptype/README.md new file mode 100644 index 0000000..ed27df1 --- /dev/null +++ b/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. diff --git a/experiments/looptype/looptype_head_2026-09-29.json b/experiments/looptype/looptype_head_2026-09-29.json new file mode 100644 index 0000000..a51755f --- /dev/null +++ b/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" +} \ No newline at end of file diff --git a/experiments/looptype/train.py b/experiments/looptype/train.py new file mode 100644 index 0000000..c19b00b --- /dev/null +++ b/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())