feat(registration): implement best-of-ANTs registration for CT to T1c alignment and enhance benchmarking scripts

BENCHMARKING.md

@@ -1,152 +1,96 @@
-# Benchmarking Pipelines A / B / C
+# Pipeline benchmark (A / B / C)
 
-`scripts/benchmark_pipelines.py` compares the three segmentation pipelines on
+Benchmark of the three segmentation pipelines — A in-house Unet3D, B nnU-Net v2
-**one GPU** with a unified protocol:
+(NTUHLPLTrainer), C MONAI UNet — driven by `scripts/benchmark_pipelines.py`
+(DICE, training time, inference time, params, VRAM). Protocol details live in
+the script docstring; this file records what has been run and what the machine
+currently supports.
 
-| pipeline | backbone | data protocol |
+## Data chain status (2026-10-04)
-|---|---|---|
-| A | in-house `Unet3D(base=16, depth=4)` | processed 1 mm head-cropped percentile-normalized volumes (`pimg`/`plabel`) |
-| B | nnU-Net v2 `3d_fullres` | native source volumes (`img`/`label`), per Pipeline B's design |
-| C | MONAI `UNet(16→128)` | processed 1 mm volumes (`pimg`/`plabel`) |
 
-Each pipeline is benchmarked with **its own standard settings** (batch, patch,
+The M6 labeled set was fully rebuilt on 2026-10-04 (see
-augmentation, TTA, optimizer), so the results reflect the pipelines as used in
+`scripts/02_build_m6_dataset.py`): T1c fixed, best-of-ANTs
-the study.
+(`ants_rigid` / `ants_affine` / `ants_syn`, winner by Mattes MI), TV warped
+onto the T1c grid with SMD+linear. After the repair pass:
 
-## Metrics
+- `data/manifests/m6_labeled.jsonl`: **1312** volumes (was 1094 pre-repair),
+  `m6_unlabeled.jsonl`: **229**
+- all 1541 m6 processed volumes (`data/proc/<key>{,_label}.nii.gz`, 1mm
+  cropped/normalized) regenerated, 0 errors
+- patient-level splits rebuilt (`scripts/05_build_splits.py`, seed 0):
+  **4200** labeled rows (ntuh + m6) → train 3570 / val 211 / test 419
+  volumes, 3893 subjects; unlabeled pool 13280
+- tumor volume stats: median 3116 mm3, p95 24.8k, max 3.2M, zero_frac 0.12%
 
-- **DICE** — tumor probability thresholded at 0.5 vs ground truth on the
+## Smoke run (2026-10-04)
-  held-out patient-level test split (`split_test.jsonl`), the primary-metric
-  convention of all three pipelines. B additionally reports the **hard-seg**
-  (argmax) Dice it exports natively.
-- **Training** — wall-clock total and s/epoch for `--train-epochs` on the
-  first `--max-rows` train rows (single GPU). For B, the one-time raw-dataset
-  build and `plan_and_preprocess` wall time are reported separately
-  (`raw_build_s`, `plan_preprocess_s`).
-- **Inference** — mean **s/volume** on `--max-eval-rows` held-out test volumes
-  (first `--eval-warmup` cases untimed), each pipeline's standard
-  sliding-window + TTA protocol (`--no-tta` disables TTA everywhere).
-- **Model size** — parameter count, checkpoint weight file size (MB), and peak
-  **VRAM** (MB) for the training run and for the inference run.
-  (A/C training VRAM via `torch.cuda.max_memory_allocated`; B via nvidia-smi
-  polling around the `nnUNetv2_train` subprocess.)
 
-## Prerequisites
+Small/fast end-to-end check that all three pipelines run on the rebuilt data
-
+(1 epoch, 24 train rows, 3 eval rows, TTA on, GPU 0):
-- The `longitudinal-313` conda env and the GPU(s) it sees:
-  ```bash
-  source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
-  ```
-- Data for Pipelines A and C (and split membership for all three):
-  `data/proc/*` + `data/manifests/split_{train,val,test}.jsonl`, i.e.
-  `scripts/preprocess.py` and `scripts/05_build_splits.py` have been run.
-- Data for Pipeline B (train/eval mode): the labeled source manifests
-  `data/manifests/{ntuh,m6_labeled}.jsonl` with reachable native NIfTIs
-  (the benchmark joins native `img`/`label` paths by case key).
-- `--skip-train` mode additionally needs the study checkpoints:
-  `runs/round0/best.pt` (A), `runs_monai/round0/best.pt` (C), and the study
-  nnU-Net results tree `nnu/results/Dataset210_NTUH_T1C_PL/...` (B).
-
-## Quick start
-
-```bash
-cd /mnt/b4/xfr/git26/longitudinal
-conda activate longitudinal-313
-
-# small/fast check (what the validation run used): ~minutes per pipeline
-python scripts/benchmark_pipelines.py \
-    --train-epochs 1 --max-rows 24 --nnu-cases 32 --max-eval-rows 4
-
-# default full benchmark: train 100 epochs on 80% of train cases (~1986 rows),
-# eval on 10% of test volumes (~30)
-python scripts/benchmark_pipelines.py
-```
-
-Outputs (repo root):
 
 ```
-results/benchmark_pipelines.json        # settings + all metrics, per pipeline
+python scripts/benchmark_pipelines.py --gpu 0 --max-rows 24 --nnu-cases 24 \
-results/benchmark_pipelines_per_row.jsonl
+    --max-eval-rows 3 --train-epochs 1 --out results/benchmark_pipelines_smoke.json
-results/benchmark_pipelines_bench.png   # 6-panel bar chart
 ```
 
-plus the artifacts it created (safe to delete): `runs_bench/{A,C}/best.pt`,
+| pipeline | wall | s/epoch | tr VRAM | s/vol (TTA) | in VRAM | params | weights | dice (1 epoch) |
-`nnu/{raw,preprocessed,results}/Dataset221_T1C_BENCH`,
+|---|---|---|---|---|---|---|---|---|
-`logs/nnu/bench_{plan,train}.log`.
+| A Unet3D(16, depth 4) | 17 s | 2.5 (8 steps) | 2.0 GB | 4.12 | 0.5 GB | 1.46 M | 5.9 MB | 0.0000 |
+| B nnU-Net 3d_fullres | 285 s (incl. 83 s plan&preprocess, 4 s raw build) | 73.3 | 9.8 GB | 34.73 | 1.8 GB | 30.8 M | 246 MB | 0.0000 |
+| C MONAI UNet(16→128) | 5 s | 2.5 (8 steps) | 1.7 GB | 0.33 | 0.5 GB | 1.19 M | 4.8 MB | 0.0005 |
 
-## All options
+DICE ≈ 0 is expected (models trained for a single epoch on 24 volumes); the
+point of the run was validating the full chain — new m6 native/processed
+manifests, split membership, B's `ntuh + m6_labeled` native join, nnU-Net
+isolated dataset 221 build, inference + dice scoring. All three passed.
+Eval rows for the smoke run were the first three `split_test` volumes (all
+ntuh). Artifacts: `results/benchmark_pipelines_smoke.json`,
+`_per_row.jsonl`, `_bench.png`.
 
-| flag | default | meaning |
+### Environment fixes made during the check
-|---|---|---|
-| `--pipelines` | `A B C` | space-separated subset, e.g. `"A C"` |
-| `--gpu` | `0` | physical GPU index; all pipelines run serially on this one GPU |
-| `--train-epochs` | `100` | training budget for A and C (and B via `NNU_PL_EPOCHS`) |
-| `--max-rows` | `80%` of `split_train` | train rows for A/C (first N of `split_train.jsonl` that have labels; default 80% of the split, ~1986 rows) |
-| `--nnu-cases` | `80%` of `split_train` | train cases for B (first N labeled native cases of the same split; default same as `--max-rows`) |
-| `--max-eval-rows` | `10%` of `split_test` | held-out test volumes for inference + DICE (default 10% of the split, ~30 rows) |
-| `--eval-warmup` | `1` | untimed warmup cases before inference timing |
-| `--workers` | `4` | data-loader workers (A/C) |
-| `--batch` | `3` | per-GPU batch (A/C) |
-| `--patch` | `96` | patch/sliding-window size (A/C) |
-| `--no-tta` | off | disable each pipeline's standard test-time augmentation |
-| `--skip-train` | off | evaluate existing study checkpoints instead of training |
-| `--nnu-dsid` | `221` train / `210` skip-train | nnU-Net dataset id: in training mode an isolated `Dataset{id}_T1C_BENCH` (the study's 210 is never touched); in `--skip-train` mode the study's 210 by default, or a previously benchmarked id to re-evaluate it |
-| `--nnu-skip-prep` | off | B (train mode): skip `plan_and_preprocess` and reuse the existing raw + preprocessed tree for `--nnu-dsid` (e.g. retime one more training epoch on an already prepared dataset) |
-| `--out` | `results/benchmark_pipelines.json` | output JSON path (the `_per_row.jsonl` / `_bench.png` are derived from it) |
 
-## Examples
+- `monai` was missing from `longitudinal-313` (pipeline C failed with
+  ModuleNotFoundError) — installed `pip install monai==1.6.1 --no-deps`
+  (torch 2.14+cu126 untouched); `import monai` verified.
 
-```bash
+## Resuming training (`--resume`)
-# longer training budget, bigger eval set (quality-focused run)
-python scripts/benchmark_pipelines.py --train-epochs 5 --max-rows 240 --max-eval-rows 20
 
-# timing only, no TTA
+All three pipelines can continue from their latest benchmark checkpoint:
-python scripts/benchmark_pipelines.py --no-tta --max-eval-rows 20
 
-# re-evaluate the study's trained checkpoints (no training)
+```
-python scripts/benchmark_pipelines.py --skip-train --max-eval-rows 24
+python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \
-
+    --train-epochs 200 --nnu-skip-prep
-# benchmark B only; keep the trained Dataset221 model for a second inference-only pass
-python scripts/benchmark_pipelines.py --pipelines B
-python scripts/benchmark_pipelines.py --pipelines B --skip-train --nnu-dsid 221
-
-# use GPU 1 and write to a custom location
-python scripts/benchmark_pipelines.py --gpu 1 --out results/bench_gpu1.json
 ```
 
-## Interpreting the results
+- `--train-epochs` is the **total budget** (not epochs added). Each pipeline
+  trains from where its checkpoint left off up to that total.
+- **A / C**: load `runs_bench/{A,C}/best.pt` (model + optimizer + scheduler
+  state — checkpoints written before 2026-10-04 carry model + epoch only and
+  fall back to fresh optimizer/scheduler) and run the remaining epochs; the
+  run is a no-op train (eval only) when the checkpoint already reached the
+  total.
+- **B**: appends nnU-Net's `--c` flag (continue `fold_0` from
+  `checkpoint_latest.pth`, native optimizer + LR-schedule restore) and
+  `NNUNet_pl_EPOCHS = --train-epochs`. Pair with `--nnu-skip-prep` so the
+  preprocessed tree isn't `--clean`ed; the isolated Dataset221 is rebuilt
+  otherwise and a finished run's `checkpoint_latest.pth` is gone, so a
+  completed run + a bigger budget retrains B from scratch (interrupted runs
+  keep `checkpoint_latest.pth` and resume correctly).
+- Verified 2026-10-04: A/C resumed epoch 1 → 3 (model restored, fresh
+  opt/sched for the old-format ckpt), re-run at 3/3 trained 0 epochs, B
+  continued from epoch 1 with LR restored mid-decay (0.00536 @ start of
+  epoch 2).
 
-- **"1 epoch" is pipeline-native.** A/C: one pass over the N train volumes
+## Readiness notes for a full run
-  (batch 3, 96³ random patches, AdamW 3e-4 + warmup/cosine, bf16).
-  B: one pass over the N train cases at nnU-Net's planned batch and patch size
-  (typically 2 and ~256³, AdamW 1e-2 PolyLR, fp16). Epoch times are therefore
-  comparable as *wall-clock cost per pass over the training set in each
-  pipeline's standard configuration*, not per identical number of patches.
-- **DICE reflects the training budget.** Short benchmarks measure the
-  machinery, not the final quality of the pipelines. Compare DICE across
-  pipelines only at the same `--train-epochs`/`--max-rows` budget, or use
-  `--skip-train` on fully trained study checkpoints.
-- **Inference times are single-GPU, standard-protocol numbers**
-  (A: 96³ step 48 + 4-flip TTA; B: nnU-Net gaussian sliding window +
-  mirroring, native volumes; C: MONAI gaussian 50% + 4-flip TTA).
-  B is the slowest per volume on this data set because it runs on uncropped
-  native volumes with full mirroring TTA.
-- **B's one-time cost.** `raw_build_s` + `plan_preprocess_s` are paid once per
-  (dataset, round) — in the iterative study every round repays the
-  plan/preprocess on the growing dataset; the benchmark reports it for
-  context only.
-- **VRAM notes.** Training VRAM includes optimizer states (AdamW) and the
-  largest batch; B's number is the nvidia-smi peak of the whole
-  `nnUNetv2_train` process tree on that GPU. Inference VRAM is per
-  pipeline as measured in-process (B: torch peak over the timed cases).
 
-## Troubleshooting
+- Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of
-
+  split_train ≈ 2856 for A/C and for B's `--nnu-cases`), `--max-eval-rows`
-| symptom | fix |
+  (default 10% of split_test ≈ 42), `--no-tta`.
-|---|---|
+- `--skip-train` is only usable for **A** (`runs/round0/best.pt`, trained
-| `split_train.jsonl has no labeled processed rows` | run `scripts/preprocess.py` over the manifests and `scripts/05_build_splits.py` first |
+  2026-09-25 on the *previous* split — a few patients moved between
-| `B: key ... not in ntuh/m6_labeled manifests` | build the source manifests (`scripts/01_build_ntuh_manifest.py`, `scripts/02_build_m6_dataset.py`) and confirm the native NIfTIs are reachable |
+  train/test in the rebuild, so its test dice carries a small leakage
-| `--skip-train: .../best.pt not found` | the study checkpoint for that pipeline doesn't exist yet — train a round first, or drop `--skip-train` |
+  caveat). No study checkpoints exist for B (Dataset210) or C
-| B `--skip-train` fails on the model folder | the study's `nnu/results/Dataset210_NTUH_T1C_PL/NTUHLPLTrainer__nnUNetPlans__3d_fullres/fold_0/` must exist (produced by `scripts_nnu/03_nnu_train.py`) |
+  (`runs_monai/round0/best.pt`); they must be trained.
-| OOM on B | B trains at the planned 256³-class patches and needs ~10–20 GB; free other GPU users or run with a smaller `--nnu-cases` set (patch size derives from the median case size) |
+- GPUs: 3× 80 GB available; the benchmark uses one (`--gpu`).
-| inflated/inconsistent timings | do not run two benchmarks on the same `--gpu` concurrently; note that data-loader (NFS) latency is included in training and inference wall times |
+- B's one-time cost scales with `--nnu-cases` (plan_and_preprocess ≈ 3.5 s/case
+  at 24 cases with 8 preprocess procs, plus training time at the planned
+  batch/patch).

scripts/02_build_m6_dataset.py

@@ -1,12 +1,20 @@
 """Build M6-2025 manifests: T1c volumes, labeled when GTV exists.
 
-GTV lives in CT space; T1c is registered to CT (rigid, Mattes MI on the full
+GTV lives in CT space; the T1c is the fixed image and the CT is registered
-volumes -- no mask, no crop / pad; multi-start initial tilts because
+to it with the best of three ANTsPy pipelines (src/registration_pipelines.py,
-inter-session head tilt reaches ~30 deg and tilted MI local minima can
+full volumes, no mask / crop / pad, each toolkit's native COM init):
-outscore the true alignment).
+    ants_rigid  : Rigid[0.25] (Mattes, 500x500x200, 4x2x1)
-The inverse transform then warps CT and GTV (nearest) into T1c native space.
+    ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1)
+    ants_syn    : SyN (Affine[0.25] + SyN[0.2,3,0])
+The winner is the candidate with the highest Mattes-style MI between the
+warped CT and the T1c (full volumes). All three return the
+transform in the unified contract direction (T1c -> CT), so resampling the
+CT / GTV onto the T1c grid with it places them in T1c native space -- the
+TV is warped with SMD + linear (signed distance to the TV boundary, linearly
+resampled, thresholded at the mean CT spacing halved; the validated recipe
+from scripts/qa_label_interpolation.py).
 Acceptance requires the warped CT skull ring to be concentric with the MR head
-silhouette (crop-independent check).
+silhouette.
 Outputs:
   data/manifests/m6_labeled.jsonl    {key, subject, date, img, label(warped GTV), source:'m6'}
   data/manifests/m6_unlabeled.jsonl  {key, subject, date, img, label:None, source:'m6'}
@@ -133,8 +141,8 @@ def _skull_ring_gap(warped_ct, ha):
         for a in np.linspace(0, 2 * np.pi, 72, endpoint=False):
             dy, dx = np.sin(a), np.cos(a)
             rr = np.arange(0.0, 160.0, 0.5)
-            yy = np.clip((cy0 + rr * dy).astype(int), 0, ha.shape[0] - 1)
+            yy = np.clip((cy0 + rr * dy).astype(int), 0, sl.shape[0] - 1)
-            xx = np.clip((cx0 + rr * dx).astype(int), 0, ha.shape[1] - 1)
+            xx = np.clip((cx0 + rr * dx).astype(int), 0, sl.shape[1] - 1)
             rh = rr[np.where(sl[yy, xx])[0]]
             rs = rr[np.where(sk[yy, xx])[0]]
             if len(rh) and len(rs):
@@ -321,6 +329,112 @@ def register_ct_to_t1c(ct, t1c, threads=12):
     return tr, info, tissue_w
 
 
+_QA = None
+
+
+def _load_qa():
+    """Import scripts/qa_label_interpolation.py (the validated SMD + linear
+    label-warp recipe: smd_threshold / make_tv_bin / warp_tv). Cached per
+    process (each worker process keeps its own copy)."""
+    global _QA
+    if _QA is None:
+        import importlib.util
+        p = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                         "qa_label_interpolation.py")
+        spec = importlib.util.spec_from_file_location("qa_label_interp", p)
+        m = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(m)
+        _QA = m
+    return _QA
+
+
+def register_ants_best(ct, t1c):
+    """Best-of-ANTs CT -> T1c registration with the T1c as the fixed image.
+
+    Runs the three ANTsPy pipelines from src/registration_pipelines.py
+    (fixed = T1c, moving = CT, full volumes, no mask / crop / pad, each
+    toolkit's native COM init):
+        ants_rigid  : Rigid[0.25] (Mattes, 500x500x200, 4x2x1)
+        ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1)
+        ants_syn    : SyN (Affine[0.25] + SyN[0.2,3,0])
+    Each returns a transform in the unified contract direction (T1c -> CT),
+    so sitk.Resample(ct, t1c, T, ...) / sitk.Resample(tv, t1c, T, ...) place
+    the CT / TV concentric with the MR head -- no inversion (the T1c was the
+    fixed image). The winner is the candidate with the highest Mattes-style
+    MI between the warped CT and the T1c (_mi_value, full volumes). The
+    per-candidate skull-ring gap is still recorded and gates acceptance.
+
+    Returns (transform, info, tissue_w) like register_ct_to_t1c.
+    """
+    import time
+    from scipy import ndimage
+    from src import registration_pipelines as RP
+    from src.common import head_mask_from_image
+    t_start = time.time()
+    ct = sitk.Cast(ct, sitk.sitkFloat32)
+    t1c = sitk.Cast(t1c, sitk.sitkFloat32)
+    calls = {
+        "ants_rigid": lambda: RP.register_antspyx(ct, t1c),
+        "ants_affine": lambda: RP.register_ants_affine_variant(ct, t1c, "antspyx"),
+        "ants_syn": lambda: RP.register_ants_variant_deformable(ct, t1c, "ants_syn"),
+    }
+    ha = sitk.GetArrayFromImage(t1c)
+    cands = {}
+    for name, fn in calls.items():
+        try:
+            T, meta = fn()
+        except Exception as e:  # noqa
+            cands[name] = {"error": repr(e)}
+            print(f"  {name}: FAILED {e!r}", flush=True)
+            continue
+        warped_ct = RP._warp_ct(ct, t1c, T)
+        mi = _mi_value(ha, warped_ct)
+        gap = RP._skull_gap(warped_ct, ha)
+        cands[name] = {"T": T, "meta": meta, "mi": mi, "gap": gap}
+        print(f"  {name}: mi={mi:.4f} gap={gap} ({meta.get('seconds_reg')}s)",
+              flush=True)
+    ok = {n: c for n, c in cands.items() if "T" in c}
+    if not ok:
+        raise RuntimeError("all ANTs candidates failed: "
+                           + repr({n: c.get("error") for n, c in cands.items()}))
+    best_name = max(ok, key=lambda n: ok[n]["mi"])
+    T, meta, gap = ok[best_name]["T"], ok[best_name]["meta"], ok[best_name]["gap"]
+
+    def _tp(x, y, z):
+        return np.array(T.TransformPoint((float(x), float(y), float(z))))
+
+    t0 = _tp(0, 0, 0)
+    mat = np.column_stack([_tp(1, 0, 0) - t0, _tp(0, 1, 0) - t0, _tp(0, 0, 1) - t0])
+    scale = [round(float(np.linalg.norm(mat[:, i])), 4) for i in range(3)]
+    hm = head_mask_from_image(ha)
+    warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
+    posw = warped_ct[warped_ct > 0]
+    tissue_w = ((warped_ct > max(2.0, float(np.percentile(posw, 2))))
+                if posw.size else np.zeros_like(warped_ct, bool))
+    tissue_w = ndimage.binary_opening(tissue_w, iterations=2)
+    coverage = float(tissue_w[hm].sum() / max(hm.sum(), 1))
+    info = {
+        "pipeline": "ants_best",
+        "fixed": "T1c", "moving": "CT",
+        "best": best_name,
+        "candidates": {
+            n: ({**c["meta"], "mi": round(c["mi"], 4),
+                "skull_ring_gap_mm": c["gap"], "selected": n == best_name}
+                if "T" in c else c)
+            for n, c in cands.items()},
+        "affine_stage": {"matrix": [[round(float(v), 4) for v in row] for row in mat.tolist()],
+                         "offset": [round(float(x), 3) for x in t0], "scale": scale,
+                         "transform_type": type(T).__name__},
+        "coverage_head_box": round(coverage, 3),
+        "skull_ring_gap_mm": gap,
+        "seconds_total": round(time.time() - t_start, 1),
+    }
+    print("reg done %.0fs best=%s mi=%.4f gap=%s cov %.2f" % (
+        info["seconds_total"], best_name, ok[best_name]["mi"], gap, coverage),
+          flush=True)
+    return T, info, tissue_w
+
+
 def reg_worker(row):
     reg_dir = d("data/m6reg")
     key = row["key"]
@@ -344,14 +458,17 @@ def reg_worker(row):
                       f, indent=1)
         return key, False, 0
     gtv = read_nii_img(row["label"])
-    tr, info, tissue_w = register_ct_to_t1c(ct, t1c)
+    tr, info, tissue_w = register_ants_best(ct, t1c)
     # QA files in m6reg, everything in (oriented) MR coordinates:
     #   _ct:  CT warped to the MR grid (linear), for overlaying _tv
     #   _mr:  the oriented MR volume itself
-    #   _tv:  original TV warped to the MR grid (nearest) — the study label
+    #   _tv:  original TV warped to the MR grid (SMD + linear) — the study label
     #   _reg.json: registration metrics + final image information
-    tv_mr = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0)
+    qa = _load_qa()
-    z = sitk.GetArrayFromImage(tv_mr)
+    thr = qa.smd_threshold([float(s) for s in gtv.GetSpacing()])[0]
+    z = qa.warp_tv(qa.make_tv_bin(gtv), t1c, tr, "smd_linear", thr)
+    tv_mr = sitk.GetImageFromArray(z)
+    tv_mr.CopyInformation(t1c)
     nvox = int((z > 0.5).sum())
     from scipy import ndimage
     tv_in_ct_frac = 1.0
@@ -369,6 +486,7 @@ def reg_worker(row):
     info.update({
         "key": key, "subject": row["subject"], "date": row["date"],
         "tv_nvox": nvox, "tv_in_ct_frac": round(tv_in_ct_frac, 3),
+        "tv_smd_thr_mm": round(thr, 3),
         "alignment": align,
         "tv_accepted": accepted,
         "accept_reason": ("ok" if accepted else

scripts/benchmark_pipelines.py

@@ -30,6 +30,9 @@ Usage (repo root, inside the `longitudinal-313` conda env):
   python scripts/benchmark_pipelines.py                          # train + eval A B C
 python scripts/benchmark_pipelines.py --pipelines A C
    python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240
+   python scripts/benchmark_pipelines.py --pipelines "A B C" --resume --train-epochs 200
+   # ^ continue from runs_bench/{A,C}/best.pt and the nnU-Net fold_0
+   #   checkpoint_latest.pth; --train-epochs is the TOTAL budget
   python scripts/benchmark_pipelines.py --no-tta                 # no test-time augmentation
   python scripts/benchmark_pipelines.py --skip-train             # eval existing study checkpoints
 Outputs: results/benchmark_pipelines.json (+ _per_row.jsonl, _bench.png)
@@ -81,6 +84,12 @@ def parse_args():
     ap.add_argument("--skip-train", action="store_true",
                     help="eval existing study checkpoints (runs/round0, runs_monai/round0, nnu/results 210) "
                          "instead of training")
+    ap.add_argument("--resume", action="store_true",
+                    help="continue training from the latest benchmark checkpoint. A/C: load "
+                         "runs_bench/{A,C}/best.pt (model + optimizer + scheduler when present) "
+                         "and train the remaining epochs. B: append nnUNet's --c flag (continue "
+                         "from fold_0/checkpoint_latest.pth; pair with --nnu-skip-prep). "
+                         "--train-epochs is the TOTAL budget, not epochs added")
     ap.add_argument("--nnu-dsid", type=int, default=None,
                 help="nnU-Net dataset id: training mode default 221 (isolated; the study's 210 is never "
                      "touched), --skip-train mode default 210 (study model); pass explicitly to evaluate "
@@ -218,10 +227,26 @@ def bench_A(args, train_rows, test_rows):
                   torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1),
                                                               eta_min=3e-4 * 0.05)],
             milestones=[warmup])
+        ckpt = os.path.join(ckpt_dir, "best.pt")
+        start_epoch = 0
+        if args.resume and os.path.exists(ckpt):
+            sd = torch.load(ckpt, map_location=dev)
+            model.load_state_dict(sd["model"])
+            if "optimizer" in sd:
+                opt.load_state_dict(sd["optimizer"])
+            if "scheduler" in sd:
+                sched.load_state_dict(sd["scheduler"])
+            start_epoch = int(sd.get("epoch", 0))
+            res["resumed_from_epoch"] = start_epoch
+            print(f"[A] resume: {ckpt} epoch {start_epoch}/{args.train_epochs} "
+                  f"(optimizer={'restored' if 'optimizer' in sd else 'fresh'}, "
+                  f"scheduler={'restored' if 'scheduler' in sd else 'fresh'})", flush=True)
+        elif args.resume:
+            print(f"[A] --resume: {ckpt} not found, training from scratch", flush=True)
         model.train()
         reset_peak(dev)
         t0 = time.time()
-        for ep in range(args.train_epochs):
+        for ep in range(start_epoch, args.train_epochs):
             if hasattr(dl.sampler, "set_epoch"):
                 dl.sampler.set_epoch(ep)
             loss_sum = n = 0.0
@@ -245,11 +270,11 @@ def bench_A(args, train_rows, test_rows):
             res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
         dt = time.time() - t0
         res["train_s"] = round(dt, 1)
-        res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
+        res["train_s_per_epoch"] = round(dt / max(args.train_epochs - start_epoch, 1), 1)
         res["train_vram_mb"] = peak_mb(dev)
         res["steps_per_epoch"] = steps
-        ckpt = os.path.join(ckpt_dir, "best.pt")
+        torch.save({"model": model.state_dict(), "epoch": args.train_epochs,
-        torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
+                    "optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, ckpt)
         del dl, opt, sched
     res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
     res["params"] = int(sum(p.numel() for p in model.parameters()))
@@ -326,10 +351,26 @@ def bench_C(args, train_rows, test_rows):
                                                               eta_min=3e-4 * 0.05)],
             milestones=[warmup])
         loss_f = WeightedDiceCELoss()
+        ckpt = os.path.join(ckpt_dir, "best.pt")
+        start_epoch = 0
+        if args.resume and os.path.exists(ckpt):
+            sd = torch.load(ckpt, map_location=dev)
+            model.load_state_dict(sd["model"])
+            if "optimizer" in sd:
+                opt.load_state_dict(sd["optimizer"])
+            if "scheduler" in sd:
+                sched.load_state_dict(sd["scheduler"])
+            start_epoch = int(sd.get("epoch", 0))
+            res["resumed_from_epoch"] = start_epoch
+            print(f"[C] resume: {ckpt} epoch {start_epoch}/{args.train_epochs} "
+                  f"(optimizer={'restored' if 'optimizer' in sd else 'fresh'}, "
+                  f"scheduler={'restored' if 'scheduler' in sd else 'fresh'})", flush=True)
+        elif args.resume:
+            print(f"[C] --resume: {ckpt} not found, training from scratch", flush=True)
         model.train()
         reset_peak(dev)
         t0 = time.time()
-        for ep in range(args.train_epochs):
+        for ep in range(start_epoch, args.train_epochs):
             if hasattr(dl.sampler, "set_epoch"):
                 dl.sampler.set_epoch(ep)
             loss_sum = n = 0.0
@@ -350,11 +391,11 @@ def bench_C(args, train_rows, test_rows):
             res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
         dt = time.time() - t0
         res["train_s"] = round(dt, 1)
-        res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
+        res["train_s_per_epoch"] = round(dt / max(args.train_epochs - start_epoch, 1), 1)
         res["train_vram_mb"] = peak_mb(dev)
         res["steps_per_epoch"] = steps
-        ckpt = os.path.join(ckpt_dir, "best.pt")
+        torch.save({"model": model.state_dict(), "epoch": args.train_epochs,
-        torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
+                    "optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, ckpt)
         del dl, opt, sched, loss_f
     res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
     res["params"] = int(sum(p.numel() for p in model.parameters()))
@@ -432,10 +473,21 @@ def bench_B(args, train_rows, test_rows):
 
         env = N.nnu_env(epoch=args.train_epochs, lr=1e-2)
         env["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+        bcmd = N.train_cmd(1)
+        if args.resume:
+            # nnUNet --c: continue fold_0 from checkpoint_latest.pth (optimizer +
+            # lr-schedule state are restored natively); trains up to PL_EPOCHS
+            # (= --train-epochs, the total budget)
+            bcmd.append("--c")
+            res["resumed_from_checkpoint_latest"] = True
+            if not args.nnu_skip_prep:
+                print("[B] --resume without --nnu-skip-prep: plan_and_preprocess "
+                      "still runs with --clean (the trained checkpoint is kept)",
+                      flush=True)
         poller = SmiPoller(args.gpu)
         poller.start()
         t0 = time.time()
-        N.run(N.train_cmd(1), os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
+        N.run(bcmd, os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
         res["train_s"] = round(time.time() - t0, 1)
         res["train_s_per_epoch"] = round(res["train_s"] / max(args.train_epochs, 1), 1)
         res["train_vram_mb"] = round(poller.max_mib * 1.048576, 1) if poller.max_mib else None