feat(monai): add MONAI pipeline C implementation

Introduce the third segmentation pipeline using MONAI (1.6.x) to allow
direct comparison with Pipelines A and B. This includes the implementation
of the iterative pseudo-labeling workflow, training scripts, and
inference protocols.

- Add `scripts_monai/` directory containing the MONAI pipeline scripts.
- Update documentation in `README.md` and `AGENTS.md` to include MONAI
  package requirements and pipeline details.
- Configure `.gitignore` to exclude MONAI-specific run directories.
- Update data directory descriptions to include MONAI pseudo-labels.

.gitignore

@@ -4,4 +4,5 @@ results/
 logs/
 nnu/
 runs_nnu/
+runs_monai/
 __pycache__/

AGENTS.md

@@ -10,7 +10,7 @@ Activate with:
 source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
 ```
 
-Key packages: torch 2.14 (+cu126), torchvision, numpy, scipy, pandas, scikit-learn, scikit-image, matplotlib, nibabel, SimpleITK.
+Key packages: torch 2.14 (+cu126), torchvision, monai 1.6 (pip, Pipeline C), numpy, scipy, pandas, scikit-learn, scikit-image, matplotlib, nibabel, SimpleITK.
 
 ## Project
 
@@ -19,5 +19,5 @@ Longitudinal (repeated-measures) analysis of medical imaging data. Repo is at an
 ## Conventions
 
 - All Python commands must run inside the `longitudinal` conda environment.
-- GPU is available (CUDA 12.6); use `torch.device('cuda')` when appropriate.
+- GPU is available (CUDA 12.6); use `torch.device('cuda')` when appropriate. Multi-GPU stages of Pipelines A and C run via `torchrun --standalone --nproc_per_node N` (rank = GPU).
-- No lint/test tooling is configured yet; run scripts directly with `python <script>`.
+- No lint/test tooling is configured yet; run scripts directly with `python <script>` from the repo root.

README.md

@@ -19,7 +19,8 @@ source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
 ```
 
 GPU (CUDA 12.6) is available; multi-GPU jobs use `torchrun` (in-house) or
-nnU-Net's own DDP (`-num_gpus`). Run scripts from the repo root.
+nnU-Net's own DDP (`-num_gpus`). MONAI is pip-installed in the env (Pipeline C).
+Run scripts from the repo root.
 
 ## Data sources
 
@@ -39,11 +40,15 @@ data/
   procmeta/<key>.json   geometry (origin/direction/crop_vox) + normalization
   vols.json             labeled tumor volume stats (mm3 percentiles)
   pseudo/roundK/        in-house pseudo-labels (rows.jsonl, masks, summary.json)
+  pseudo_nnu/roundK/    nnU-Net pseudo-labels (rows.jsonl, masks, summary.json)
+  pseudo_monai/roundK/  MONAI pseudo-labels (rows.jsonl, masks, summary.json)
 src/                    U-Net, dataset, losses, training/eval helpers
 scripts/                Pipeline A (in-house 3D U-Net)
 scripts_nnu/            Pipeline B (nnU-Net)
+scripts_monai/          Pipeline C (MONAI)
 runs/roundK/            in-house checkpoints (best.pt, final.pt, state.pt)
 runs_nnu/roundK/        nnU-Net checkpoint snapshots (best_nnu.pth)
+runs_monai/roundK/      MONAI checkpoints (best.pt, final.pt, state.pt)
 nnu/                    nnU-Net raw / preprocessed / results trees
 results/                evaluation JSON tables + plots
 logs/                   run logs
@@ -201,6 +206,75 @@ python scripts_nnu/05_nnu_eval_test.py --rows data/manifests/split_test.jsonl \
 ### Verified
 
 End-to-end smoke-tested on a scratch 4-case dataset: dataset build → planning →
-`splits_final.json` → training 1 epoch → warm-start → 2-way sharded prediction →
+`splits_final.json` → training 1 epoch → warm start → 2-way sharded prediction →
 selection gates → test eval. The consistency filter's keep / reject /
 single-timepoint paths are unit-tested with synthetic volumes.
+
+---
+
+## Pipeline C — MONAI iterative pseudo-labeling (`scripts_monai/`)
+
+The same study driven by **MONAI** (1.6.x, pip) as the segmentation backbone,
+keeping the identical patient-level splits, unlabeled pool, pseudo-label gates,
+frame-invariant consistency filter, and evaluation protocol so results are
+directly comparable to Pipelines A and B.
+
+- **Network:** `monai.networks.nets.UNet` (1 in / 2 out, channels 16→128,
+  ~1.2 M params) — the MONAI analogue of Pipeline A's `Unet3D(base=16, depth=4)`.
+- **Data:** MONAI `Dataset` + transform chain (load → min-pad → per-axis flips /
+  rot90 / brightness / noise → `RandCropByPosNegLabeld` 96³). Patch sampling is
+  foreground-aware (pos=neg=1) rather than Pipeline A's uniform random crop.
+- **Loss:** per-sample weighted Dice+CE on MONAI `DiceLoss` (background+foreground
+  averaged, dice weight 0.5) — the same convention as Pipeline A, so pseudo rows
+  are down-weighted sample-by-sample (`--pseudo-weight 0.3`).
+- **Schedule:** AdamW + linear warmup + cosine annealing (Pipeline A's schedule).
+- **Inference:** MONAI `sliding_window_inference` (gaussian blend, 50% overlap,
+  96³ windows) + the same 4-view flip TTA as Pipeline A; multi-GPU = torchrun
+  rank shards (like Pipeline A), DDP for training.
+- **Shared with the other pipelines:** round semantics, checkpoint layout
+  (`runs_monai/roundK/{best,final,state}.pt`), pseudo outputs
+  (`data/pseudo_monai/roundK/`), and the gate/consistency code itself is imported
+  from `scripts_nnu/nnu_common.py` so it cannot drift.
+
+| # | Script | Purpose |
+|---|---|---|
+| 01 | `01_monai_build_rows.py` | Assemble a round's training manifest (labeled + accepted pseudo rows, weighted, de-duped by key) |
+| 02 | `02_monai_train.py` | DDP training of the MONAI UNet for one round (torchrun), full-weight warm start, auto-resume |
+| 03 | `03_monai_pseudo_label.py` | Sliding-window + TTA prediction of the remaining pool, selection gates + consistency filter |
+| 04 | `04_monai_eval_test.py` | Held-out test evaluation (Dice from probability map @0.5) |
+| 05 | `05_monai_run_iterative.py` | Orchestrates rounds 0…K, table + plot |
+| — | `monai_common.py` | Network / weighted loss / transform chain / inference + distributed helpers |
+
+### Running
+
+Full study from the repo root:
+
+```bash
+python scripts_monai/05_monai_run_iterative.py --rounds 4 --gpus 3
+```
+
+Defaults: baseline 100 epochs @ 3e-4; warm-started rounds 30 epochs @ 1e-3.
+Pseudo-label gates match Pipelines A/B (`--tau-pos 0.95`, `--min-cc-frac 0.2`,
+`--neg-frac 0.9`, volume p2–p98). Optional flags: `--no-tta`,
+`--no-neg-pseudo`, and the gate overrides. Outputs to
+`results/round{k}_test_monai.json`, `results/iterative_table_monai.jsonl`,
+`results/iterative_dice_monai.png`; per-round checkpoints in
+`runs_monai/round{k}/best.pt`.
+
+Individual stages:
+
+```bash
+python scripts_monai/01_monai_build_rows.py --train data/manifests/split_train.jsonl \
+    --accepted data/pseudo_monai/round1/accepted.jsonl \
+    --out data/pseudo_monai/round1_rows.jsonl
+torchrun --standalone --nproc_per_node 3 scripts_monai/02_monai_train.py \
+    --rows data/pseudo_monai/round1_rows.jsonl --val data/manifests/split_val.jsonl \
+    --epochs 30 --lr 1e-3 --batch 3 --ckpt-dir runs_monai/round1 \
+    --pretrained runs_monai/round0/best.pt
+torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.py \
+    --ckpt runs_monai/round1/best.pt --pool data/manifests/unlabeled_pool.jsonl \
+    --out data/pseudo_monai/round2 --already data/pseudo_monai/added_keys.jsonl
+torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py \
+    --rows data/manifests/split_test.jsonl --ckpt runs_monai/round1/best.pt \
+    --out results/round1_test_monai.json
+```

scripts_monai/01_monai_build_rows.py

@@ -0,0 +1,64 @@
+"""Build the training manifest for one pseudo-labeling round.
+
+base = labeled patient-level train split (w=1.0) — the same train/val usage as
+Pipeline A: the val split is never trained on (it is the internal-validation
+set). Accepted pseudo-label rows (pos + neg) are appended with w=<pseudo-weight>
+(negative rows point at the zero mask written by stage 03), de-duplicated by
+key (first occurrence wins). Every row must point at existing pimg/plabel
+niftis.
+
+Usage:
+    python scripts_monai/01_monai_build_rows.py \
+        --train data/manifests/split_train.jsonl \
+        --accepted data/pseudo_monai/round1/accepted.jsonl \
+        --pseudo-weight 0.3 --out data/pseudo_monai/round1_rows.jsonl
+"""
+import argparse
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from src.common import load_jsonl, save_jsonl
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--train", required=True, help="labeled train split jsonl ({key,pimg,plabel})")
+    ap.add_argument("--accepted", action="append", default=None,
+                    help="accepted pseudo jsonl (repeatable); pos + neg rows")
+    ap.add_argument("--pseudo-weight", type=float, default=0.3)
+    ap.add_argument("--no-neg", action="store_true", help="drop negative pseudo rows")
+    ap.add_argument("--out", required=True)
+    args = ap.parse_args()
+
+    rows, seen = [], set()
+    for r in load_jsonl(args.train):
+        rows.append({"key": r["key"], "subject": r.get("subject"), "pimg": r["pimg"],
+                     "plabel": r["plabel"], "w": 1.0, "role": "labeled"})
+        seen.add(r["key"])
+    n_pos, n_neg = 0, 0
+    for f in (args.accepted or []):
+        for r in load_jsonl(f):
+            if not r.get("label") or r["key"] in seen:
+                continue
+            if r["role"] == "neg" and args.no_neg:
+                continue
+            seen.add(r["key"])
+            rows.append({"key": r["key"], "subject": r.get("subject"), "pimg": r["pimg"],
+                         "plabel": r["label"], "w": args.pseudo_weight, "role": r["role"]})
+            if r["role"] == "pos":
+                n_pos += 1
+            else:
+                n_neg += 1
+    missing = [r["key"] for r in rows
+               if not (os.path.exists(r["pimg"]) and os.path.exists(r["plabel"]))]
+    if missing:
+        raise FileNotFoundError(f"{len(missing)} rows missing pimg/plabel, e.g. {missing[:3]}")
+    save_jsonl(rows, args.out)
+    n_labeled = len(rows) - n_pos - n_neg
+    print(f"[monai:rows] {args.out}: total={len(rows)} labeled={n_labeled} pos={n_pos} neg={n_neg}",
+          flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts_monai/02_monai_train.py

@@ -0,0 +1,156 @@
+"""MONAI DDP training entrypoint for one pseudo-labeling round (torchrun).
+
+torchrun --standalone --nproc_per_node 3 scripts_monai/02_monai_train.py \
+    --rows data/pseudo_monai/round0_rows.jsonl --val data/manifests/split_val.jsonl \
+    --epochs 100 --lr 3e-4 --batch 3 --ckpt-dir runs_monai/round0 \
+    [--pretrained runs_monai/round0/best.pt]
+
+- backbone: monai.networks.nets.UNet (monai_common.build_model)
+- data: MONAI Dataset + transform chain (load → pad → flip/rotate90/intensity →
+  foreground-aware 96³ crop → typed tensors); row "w" carries the per-sample
+  loss weight (pseudo rows down-weighted)
+- loss: per-sample Dice+CE, weighted batch mean (Pipeline A's convention)
+- schedule: AdamW + linear warmup + cosine (Pipeline A's schedule)
+- checkpoints in --ckpt-dir: best.pt (max val dice), final.pt, state.pt
+  (auto-resume after a crash, as in Pipeline A)
+"""
+import argparse
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import numpy as np
+import torch
+from src.common import load_jsonl, read_nii_arr
+from monai_common import (build_model, WeightedDiceCELoss, make_dataloader,
+                          predict_probs, prob_dice, rank_info, init_dist, barrier, destroy_dist)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--rows", required=True, help="jsonl of {key,pimg,plabel,w}")
+    ap.add_argument("--val", required=True, help="held-out val split jsonl ({pimg,plabel})")
+    ap.add_argument("--epochs", type=int, default=100)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--batch", type=int, default=3, help="per-GPU batch")
+    ap.add_argument("--patch", type=int, default=96)
+    ap.add_argument("--workers", type=int, default=4)
+    ap.add_argument("--ckpt-dir", required=True)
+    ap.add_argument("--pretrained", default=None, help="full-weight warm-start checkpoint")
+    ap.add_argument("--val-every", type=int, default=10)
+    ap.add_argument("--val-limit", type=int, default=60)
+    ap.add_argument("--sw-batch", type=int, default=8, help="sliding-window batch at inference")
+    args = ap.parse_args()
+
+    rank, world, local_rank = rank_info()
+    init_dist()
+    torch.manual_seed(0)
+    np.random.seed(0)
+    torch.cuda.set_device(local_rank)
+    device = f"cuda:{local_rank}"
+
+    rows = load_jsonl(args.rows)
+    val_rows = load_jsonl(args.val)[:args.val_limit]
+
+    dl = make_dataloader(rows, win=args.patch, batch=args.batch, workers=args.workers)
+    steps_per_epoch = max(len(dl), 1)
+
+    model = build_model(device=device)
+    if args.pretrained:
+        sd0 = torch.load(args.pretrained, map_location=device, weights_only=True)
+        model.load_state_dict(sd0.get("model", sd0))
+        if rank == 0:
+            print(f"[monai:train:rank0] warm start (all weights) from {args.pretrained}", flush=True)
+    ddp = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank]) if world > 1 else model
+    total_steps = steps_per_epoch * args.epochs
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
+    warmup = min(300, max(10, total_steps // 10))
+    base_sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1),
+                                                            eta_min=args.lr * 0.05)
+    sched = torch.optim.lr_scheduler.SequentialLR(
+        opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup), base_sched],
+        milestones=[warmup])
+    loss_f = WeightedDiceCELoss()
+
+    best, best_epoch, start_epoch = -1.0, -1, 0
+    if rank == 0:
+        os.makedirs(args.ckpt_dir, exist_ok=True)
+    state_f = os.path.join(args.ckpt_dir, "state.pt")
+    if os.path.exists(state_f):
+        st = torch.load(state_f, map_location="cpu", weights_only=True)
+        model.load_state_dict(st["model"])
+        opt.load_state_dict(st["opt"])
+        sched.load_state_dict(st["sched"])
+        best, best_epoch, start_epoch = st["best"], st["best_epoch"], st["epoch"]
+        if rank == 0:
+            print(f"[monai:train] auto-resuming from state.pt @epoch {start_epoch} (best={best:.4f})", flush=True)
+    if rank == 0:
+        print(f"[monai:train] rows={len(rows)} val={len(val_rows)} epochs={args.epochs} "
+              f"steps/epoch={steps_per_epoch} world={world}", flush=True)
+    barrier()
+
+    def save_state():
+        tmp = state_f + ".tmp"
+        torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(),
+                    "best": best, "best_epoch": best_epoch, "epoch": epoch + 1}, tmp)
+        os.replace(tmp, state_f)
+
+    for epoch in range(start_epoch, args.epochs):
+        if hasattr(dl.sampler, "set_epoch"):
+            dl.sampler.set_epoch(epoch)
+        model.train()
+        run_loss, run_n = 0.0, 0
+        for batch in dl:
+            img = batch["pimg"].to(device, non_blocking=True)
+            lab = batch["plabel"].squeeze(1).to(device, non_blocking=True)
+            wts = batch["w"].to(device, non_blocking=True)
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                logits = ddp(img)
+            loss = loss_f(logits, lab, wts)
+            opt.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+            opt.step()
+            sched.step()
+            run_loss += float(loss.detach())
+            run_n += 1
+        if rank == 0:
+            print(f"  epoch {epoch+1}/{args.epochs} loss={run_loss/max(run_n,1):.4f} "
+                  f"lr={opt.param_groups[0]['lr']:.2e}", flush=True)
+        if rank == 0 and (epoch + 1) % max(1, args.val_every // 2) == 0:
+            save_state()
+        if val_rows and rank == 0 and (epoch + 1) % args.val_every == 0:
+            dv, n_ok = 0.0, 0
+            for r in val_rows:
+                try:
+                    vol = read_nii_arr(r["pimg"]).astype("float32")
+                    labv = read_nii_arr(r["plabel"]).astype("uint8")
+                    p = predict_probs(model, vol, device, args.patch, tta=True, sw_batch=args.sw_batch)
+                    if p.shape == labv.shape:
+                        dv += prob_dice(p, labv)
+                        n_ok += 1
+                except Exception as e:  # noqa
+                    print("  val err", r.get("key"), repr(e))
+            dv = dv / max(n_ok, 1)
+            print(f"  [val] epoch {epoch+1} dice={dv:.4f}", flush=True)
+            if dv > best:
+                best, best_epoch = dv, epoch + 1
+                torch.save({"model": model.state_dict(), "epoch": epoch + 1, "val_dice": best},
+                           os.path.join(args.ckpt_dir, "best.pt"))
+                save_state()
+        barrier()
+    if rank == 0:
+        torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": best},
+                   os.path.join(args.ckpt_dir, "final.pt"))
+        best_f = os.path.join(args.ckpt_dir, "best.pt")
+        if not os.path.exists(best_f):
+            torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": 0.0}, best_f)
+            print("[monai:train] no val run; best.pt = final weights", flush=True)
+        print(f"[monai:train] done best_val_dice={best:.4f}@{best_epoch}", flush=True)
+    destroy_dist()
+
+
+if __name__ == "__main__":
+    main()

scripts_monai/03_monai_pseudo_label.py

@@ -0,0 +1,169 @@
+"""Pseudo-label the unlabeled pool with the MONAI round model (torchrun; rank = GPU).
+
+torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.py \
+    --ckpt runs_monai/round0/best.pt --pool data/manifests/unlabeled_pool.jsonl \
+    --out data/pseudo_monai/round1 [--already data/pseudo_monai/added_keys.jsonl]
+
+Each rank (GPU) takes rows[rank::world] and computes MONAI sliding-window
+(gaussian blend, 50% overlap) + 4-flip-TTA tumor probabilities, then applies
+the study gates (identical to Pipelines A/B, shared from scripts_nnu/nnu_common):
+  pos: p_tumor >= tau_pos, median cleanup, largest-CC fraction >= min_cc_frac,
+       volume within the labeled-tumor [p2, p98] range (data/vols.json)
+  neg: >= neg_frac of interior (vol > 0.02) voxels have p_bg >= tau_neg
+then the per-subject frame-invariant longitudinal consistency filter over
+accepted positive timepoints (head-relative centroid shift + volume ratio).
+
+Writes <out>/rows.jsonl, <out>/<key>_label.nii.gz (tumor mask for pos, zero
+mask for neg), <out>/accepted.jsonl, <out>/summary.json. Per-rank shards are
+flushed to <out>/part{rank}.jsonl so interrupted runs resume.
+
+Usage:
+    python scripts_monai/03_monai_pseudo_label.py <same flags, single GPU>
+    torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.py ...
+"""
+import argparse
+import os
+import sys
+import json
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import numpy as np
+import SimpleITK as sitk
+from src.common import load_jsonl, save_jsonl, read_nii_arr, write_arr, d
+from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model,
+                          predict_probs, pos_mask, neg_frac_bg, consistency_filter,
+                          load_voxel_stats)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--pool", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--already", default=None, help="jsonl of case keys already consumed by the dataset")
+    ap.add_argument("--no-tta", action="store_true")
+    ap.add_argument("--win", type=int, default=96)
+    ap.add_argument("--overlap", type=float, default=0.5)
+    ap.add_argument("--sw-batch", type=int, default=8)
+    ap.add_argument("--tau-pos", type=float, default=0.95)
+    ap.add_argument("--tau-neg", type=float, default=0.98)
+    ap.add_argument("--neg-frac", type=float, default=0.90)
+    ap.add_argument("--vol-qp", type=float, nargs=2, default=[2, 98])
+    ap.add_argument("--min-cc-frac", type=float, default=0.2)
+    ap.add_argument("--max-rel-dist", type=float, default=40.0,
+                    help="max mm of head-relative tumor centroid shift between compatible timepoints")
+    ap.add_argument("--vol-ratio", type=float, default=10.0,
+                    help="max tumor volume ratio between compatible timepoints")
+    args = ap.parse_args()
+
+    import torch
+    rank, world, local_rank = rank_info()
+    init_dist()
+    torch.cuda.set_device(local_rank)
+    device = f"cuda:{local_rank}"
+
+    out = d(args.out)
+    pool_rows = load_jsonl(args.pool)
+    already = set()
+    if args.already and os.path.exists(args.already):
+        already = {r["key"] for r in load_jsonl(args.already)}
+    done = set()
+    for k in range(world):
+        pf = os.path.join(out, f"part{k}.jsonl")
+        if os.path.exists(pf):
+            done |= {r["key"] for r in load_jsonl(pf)}
+    rows = [r for r in pool_rows if r["key"] not in already and r["key"] not in done]
+    shard = rows[rank::world]
+
+    vstats = load_voxel_stats()
+    vol_lo = vstats.get(f"p{args.vol_qp[0]:.0f}", 1.0)
+    vol_hi = vstats.get(f"p{args.vol_qp[1]:.0f}", 50000.0)
+    if rank == 0:
+        print(f"[monai:pseudo:rank0] pool={len(pool_rows)} already={len(already)} "
+              f"done={len(done)} to_predict={len(rows)} shard={len(shard)} "
+              f"vol_range=[{vol_lo:.0f},{vol_hi:.0f}]mm3", flush=True)
+
+    model, _ = load_model(args.ckpt, device)
+    part = os.path.join(out, f"part{rank}.jsonl")
+    buf, n_err = [], 0
+    for i, r in enumerate(shard, 1):
+        key = r["key"]
+        entry = {"key": key, "subject": r["subject"], "date": r.get("date"),
+                 "source": r.get("source"), "pimg": r["pimg"], "label": None,
+                 "role": "rej", "vol_mm3": 0, "maxp": None}
+        try:
+            pimg_itk = sitk.ReadImage(r["pimg"])
+            vol = read_nii_arr(r["pimg"]).astype("float32")
+            pt = predict_probs(model, vol, device, args.win, args.overlap,
+                               tta=not args.no_tta, sw_batch=args.sw_batch)
+            if pt.shape != vol.shape:
+                raise ValueError(f"prob shape {pt.shape} != image shape {vol.shape} for {key}")
+            entry["maxp"] = round(float(pt.max()), 4)
+            got = pos_mask(pt, args.tau_pos, args.min_cc_frac, vol_lo, vol_hi)
+            if got is not None:
+                mask, cc_frac, vol_mm3 = got
+                lp = os.path.join(out, key + "_label.nii.gz")
+                write_arr(mask, lp, itk_img=pimg_itk)
+                entry.update({"role": "pos", "label": lp, "vol_mm3": vol_mm3, "cc_frac": round(cc_frac, 3)})
+            else:
+                frac = neg_frac_bg(pt, vol, args.tau_neg, args.neg_frac)
+                if frac is not None:
+                    lp = os.path.join(out, key + "_label.nii.gz")
+                    write_arr(np.zeros(vol.shape, dtype="uint8"), lp, itk_img=pimg_itk)
+                    entry.update({"role": "neg", "label": lp, "neg_conf": round(frac, 4)})
+        except Exception as e:  # noqa
+            entry["role"] = "error"
+            n_err += 1
+            print(f"[monai:pseudo:rank{rank}] {key} ERR {e!r}", flush=True)
+        buf.append(entry)
+        if i % 20 == 0:
+            with open(part, "a") as f:
+                for b in buf:
+                    f.write(json.dumps(b) + "\n")
+            buf = []
+            print(f"[monai:pseudo:rank{rank}] {i}/{len(shard)}", flush=True)
+    if buf:
+        with open(part, "a") as f:
+            for b in buf:
+                f.write(json.dumps(b) + "\n")
+    barrier()
+    if rank != 0:
+        destroy_dist()
+        return
+
+    merged = []
+    for k in range(world):
+        pf = os.path.join(out, f"part{k}.jsonl")
+        if os.path.exists(pf):
+            merged.extend(load_jsonl(pf))
+    n_pos0 = sum(1 for x in merged if x["role"] == "pos")
+    n_rej = consistency_filter(merged, out, args.max_rel_dist, args.vol_ratio)
+    accepted = [x for x in merged if x["role"] in ("pos", "neg")]
+    save_jsonl(merged, os.path.join(out, "rows.jsonl"))
+    save_jsonl(accepted, os.path.join(out, "accepted.jsonl"))
+    posv = [x["vol_mm3"] for x in merged if x["role"] == "pos"]
+    summ = {
+        "n_pool_predicted": len(merged),
+        "n_pos": sum(1 for x in merged if x["role"] == "pos"),
+        "n_neg": sum(1 for x in merged if x["role"] == "neg"),
+        "n_pos_before_consistency": n_pos0,
+        "n_rejected_consistency": n_rej,
+        "n_other_rej": sum(1 for x in merged if x["role"] == "rej"),
+        "n_error": n_err,
+        "pos_vol_mm3": {"med": float(np.median(posv)) if posv else 0,
+                        "p5": float(np.percentile(posv, 5)) if posv else 0,
+                        "p95": float(np.percentile(posv, 95)) if posv else 0},
+        "tau_pos": args.tau_pos, "tau_neg": args.tau_neg, "neg_frac": args.neg_frac,
+        "vol_range": [vol_lo, vol_hi], "max_rel_dist_mm": args.max_rel_dist,
+        "vol_ratio": args.vol_ratio, "win": args.win, "overlap": args.overlap,
+        "tta": not args.no_tta, "ckpt": args.ckpt,
+    }
+    with open(os.path.join(out, "summary.json"), "w") as f:
+        json.dump(summ, f, indent=1)
+    print(f"[monai:pseudo] round {os.path.basename(out)}: {json.dumps(summ)}", flush=True)
+    destroy_dist()
+
+
+if __name__ == "__main__":
+    main()

scripts_monai/04_monai_eval_test.py

@@ -0,0 +1,103 @@
+"""Held-out test evaluation of a MONAI round model (torchrun; rank = GPU).
+
+torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py \
+    --rows data/manifests/split_test.jsonl --ckpt runs_monai/round0/best.pt \
+    --out results/round0_test_monai.json [--no-tta]
+
+Per case: MONAI sliding-window (gaussian blend, 50% overlap) + 4-flip-TTA tumor
+probability map; Dice at threshold 0.5 vs the held-out label — the same
+convention as Pipeline A's 08_eval and Pipeline B's primary metric.
+
+Per-rank results are appended to <out>_part{rank}.jsonl for resume; rank 0
+merges them into <out> (json) + <out>_per_row.jsonl.
+
+Usage:
+    python scripts_monai/04_monai_eval_test.py <same flags, single GPU>
+    torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py ...
+"""
+import argparse
+import os
+import sys
+import json
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+from src.common import load_jsonl, save_jsonl, read_nii_arr
+from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model,
+                          predict_probs, prob_dice)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--rows", required=True)
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--no-tta", action="store_true")
+    ap.add_argument("--win", type=int, default=96)
+    ap.add_argument("--overlap", type=float, default=0.5)
+    ap.add_argument("--sw-batch", type=int, default=8)
+    args = ap.parse_args()
+
+    rank, world, local_rank = rank_info()
+    init_dist()
+    torch.cuda.set_device(local_rank)
+    device = f"cuda:{local_rank}"
+
+    rows = load_jsonl(args.rows)
+    outp = os.path.join(os.path.dirname(os.path.abspath(args.out)),
+                        os.path.basename(args.out) + f"_part{rank}.jsonl")
+    done = set()
+    if os.path.exists(outp):
+        done = {r["key"] for r in load_jsonl(outp)}
+    shard = [r for r in rows if r["key"] not in done][rank::world]
+    if rank == 0:
+        print(f"[monai:eval:rank0] n_test={len(rows)} done={len(done)} "
+              f"shard(world={world})={len(shard)}", flush=True)
+
+    model, _ = load_model(args.ckpt, device)
+    with open(outp, "a") as f:
+        for i, r in enumerate(shard, 1):
+            key = r["key"]
+            entry = {"key": key, "dice": None}
+            try:
+                vol = read_nii_arr(r["pimg"]).astype("float32")
+                lab = read_nii_arr(r.get("plabel") or r.get("label")).astype("uint8")
+                p = predict_probs(model, vol, device, args.win, args.overlap,
+                                  tta=not args.no_tta, sw_batch=args.sw_batch)
+                if p.shape == lab.shape:
+                    entry["dice"] = round(prob_dice(p, lab), 4)
+                else:
+                    raise ValueError(f"prob {p.shape} vs label {lab.shape}")
+            except Exception as e:  # noqa
+                print(f"[monai:eval:rank{rank}] {key} ERR {e!r}", flush=True)
+            f.write(json.dumps(entry) + "\n")
+            f.flush()
+            if i % 20 == 0:
+                print(f"[monai:eval:rank{rank}] {i}/{len(shard)}", flush=True)
+    barrier()
+    if rank != 0:
+        destroy_dist()
+        return
+
+    per = []
+    for k in range(world):
+        pf = os.path.join(os.path.dirname(os.path.abspath(args.out)),
+                          os.path.basename(args.out) + f"_part{k}.jsonl")
+        if os.path.exists(pf):
+            per.extend(load_jsonl(pf))
+    dice_ok = [e["dice"] for e in per if e["dice"] is not None]
+    res = {"ckpt": args.ckpt, "n": len(per), "n_ok": len(dice_ok),
+           "dice": float(sum(dice_ok) / len(dice_ok)) if dice_ok else 0.0,
+           "tta": not args.no_tta}
+    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
+    with open(args.out, "w") as f:
+        json.dump(res, f, indent=1)
+    save_jsonl(per, args.out.replace(".json", "_per_row.jsonl"))
+    print(json.dumps({k: res[k] for k in ("n", "n_ok", "dice")}, indent=1), flush=True)
+    print("saved", args.out, flush=True)
+    destroy_dist()
+
+
+if __name__ == "__main__":
+    main()

scripts_monai/05_monai_run_iterative.py

@@ -0,0 +1,180 @@
+"""Orchestrates the MONAI iterative pseudo-labeling study (parallel to
+scripts/09 and scripts_nnu/06).
+
+Round 0: train the MONAI UNet on the labeled patient-level train split (internal
+          validation = held-out patient-level val split, subject-disjoint).
+Round k: dataset grows with accepted pseudo-labels (pos + neg) from rounds 1..k
+          (de-duplicated by key, pseudo rows weighted); warm-start training
+          (full weights) from round k-1's best checkpoint at a lower LR; then
+          pseudo-label the remaining unlabeled pool (gates identical to
+          Pipelines A/B) and evaluate the model on the held-out test split.
+
+Usage: python scripts_monai/05_monai_run_iterative.py [--rounds 4] [--gpus 3]
+"""
+import os
+import sys
+import json
+import shutil
+import argparse
+import subprocess
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from src.common import ROOT, d, load_jsonl, save_jsonl
+
+S = os.path.dirname(os.path.abspath(__file__))   # code lives next to this file, not under ROOT
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--rounds", type=int, default=4)
+    ap.add_argument("--gpus", type=int, default=3)
+    ap.add_argument("--base-epochs", type=int, default=100)
+    ap.add_argument("--base-lr", type=float, default=3e-4)
+    ap.add_argument("--round-epochs", type=int, default=30)
+    ap.add_argument("--round-lr", type=float, default=1e-3)
+    ap.add_argument("--pseudo-weight", type=float, default=0.3)
+    ap.add_argument("--batch", type=int, default=3)
+    ap.add_argument("--val-every", type=int, default=10)
+    ap.add_argument("--val-limit", type=int, default=60)
+    ap.add_argument("--no-neg-pseudo", action="store_true", help="do not add negative pseudo cases to training")
+    ap.add_argument("--no-tta", action="store_true", help="disable flip TTA at inference")
+    ap.add_argument("--tau-pos", type=float, default=0.95)
+    ap.add_argument("--tau-neg", type=float, default=0.98)
+    ap.add_argument("--neg-frac", type=float, default=0.90)
+    ap.add_argument("--vol-qp", type=float, nargs=2, default=[2, 98])
+    ap.add_argument("--min-cc-frac", type=float, default=0.2)
+    ap.add_argument("--max-rel-dist", type=float, default=40.0)
+    ap.add_argument("--vol-ratio", type=float, default=10.0)
+    args = ap.parse_args()
+
+    man = os.path.join(ROOT, "data/manifests")
+    train_f = os.path.join(man, "split_train.jsonl")
+    val_f = os.path.join(man, "split_val.jsonl")
+    test_f = os.path.join(man, "split_test.jsonl")
+    pool_f = os.path.join(man, "unlabeled_pool.jsonl")
+    results_dir = d("results")
+    logdir = d("logs/monai")
+    main_log = os.path.join(logdir, "iterative.log")
+    pdir_root = d("data/pseudo_monai")
+    added_f = os.path.join(pdir_root, "added_keys.jsonl")
+
+    def run(cmd, log, retries=2):
+        for attempt in range(retries + 1):
+            try:
+                with open(log, "a") as f:
+                    f.write(f"$ (attempt {attempt + 1}) " + cmd + "\n")
+                    f.flush()
+                    print(f"$ (attempt {attempt + 1}) " + cmd, flush=True)
+                    subprocess.run(cmd, shell=True, cwd=ROOT, stdout=f, stderr=subprocess.STDOUT, check=True)
+                return
+            except subprocess.CalledProcessError:
+                if attempt == retries:
+                    raise
+                print(f"[monai-orch] command failed, retrying in 60s: {cmd}", flush=True)
+                import time
+                time.sleep(60)
+
+    def tr(script, extra):
+        t = shutil.which("torchrun") or f"{sys.executable} -m torch.distributed.run"
+        return f"{t} --standalone --nproc_per_node {args.gpus} {os.path.join(S, script)} {extra}"
+
+    def py(script, extra, log=main_log):
+        run(f"{sys.executable} {os.path.join(S, script)} {extra}", log)
+
+    def trrun(script, extra, log):
+        run(tr(script, extra), log)
+
+    def build_rows(out, accepted, no_neg=False):
+        extra = f"--train {train_f} --pseudo-weight {args.pseudo_weight} --out {out}"
+        if no_neg:
+            extra += " --no-neg"
+        for f in accepted:
+            extra += f" --accepted {f}"
+        py("01_monai_build_rows.py", extra)
+
+    def train(round, rows_f, epochs, lr, warmstart, val_every):
+        extra = (f"--rows {rows_f} --val {val_f} --epochs {epochs} --lr {lr} "
+                 f"--batch {args.batch} --ckpt-dir runs_monai/round{round} "
+                 f"--val-every {val_every} --val-limit {args.val_limit}")
+        if warmstart:
+            extra += f" --pretrained {warmstart}"
+        trrun("02_monai_train.py", extra, os.path.join(logdir, f"train_r{round}.log"))
+
+    def predict_pool(round, ckpt, added=None):
+        out = os.path.join(pdir_root, f"round{round}")
+        extra = (f"--ckpt {ckpt} --pool {pool_f} --out {out} "
+                 f"--tau-pos {args.tau_pos} --tau-neg {args.tau_neg} --neg-frac {args.neg_frac} "
+                 f"--vol-qp {args.vol_qp[0]} {args.vol_qp[1]} --min-cc-frac {args.min_cc_frac} "
+                 f"--max-rel-dist {args.max_rel_dist} --vol-ratio {args.vol_ratio}")
+        if added and os.path.exists(added):
+            extra += f" --already {added}"
+        if args.no_tta:
+            extra += " --no-tta"
+        trrun("03_monai_pseudo_label.py", extra, os.path.join(logdir, f"pseudo_r{round}.log"))
+        return os.path.join(out, "accepted.jsonl")
+
+    def eval_test(round, ckpt):
+        out = os.path.join(results_dir, f"round{round}_test_monai.json")
+        extra = f"--rows {test_f} --ckpt {ckpt} --out {out}"
+        if args.no_tta:
+            extra += " --no-tta"
+        trrun("04_monai_eval_test.py", extra, os.path.join(logdir, f"eval_r{round}.log"))
+        return out
+
+    def best(round):
+        return os.path.join(ROOT, "runs_monai", f"round{round}", "best.pt")
+
+    # ---- round 0: baseline ----
+    rows0 = os.path.join(pdir_root, "round0_rows.jsonl")
+    build_rows(rows0, [])
+    train(0, rows0, args.base_epochs, args.base_lr, None, args.val_every)
+    accepted_f = [predict_pool(1, best(0))]
+    save_jsonl(load_jsonl(accepted_f[0]), added_f)
+    eval_test(0, best(0))
+
+    for k in range(1, args.rounds + 1):
+        rows_f = os.path.join(pdir_root, f"round{k}_rows.jsonl")
+        build_rows(rows_f, accepted_f, no_neg=args.no_neg_pseudo)
+        train(k, rows_f, args.round_epochs, args.round_lr, best(k - 1), args.val_every)
+        accepted_f.append(predict_pool(k + 1, best(k), added_f))
+        save_jsonl([r for f in accepted_f for r in load_jsonl(f)], added_f)
+        eval_test(k, best(k))
+
+    # ---- report (same layout as scripts/09 and scripts_nnu/06) ----
+    table = []
+    for k in range(args.rounds + 1):
+        resf = os.path.join(results_dir, f"round{k}_test_monai.json")
+        if not os.path.exists(resf):
+            continue
+        r = json.load(open(resf))
+        row = {"round": k, "test_dice": round(r["dice"], 4), "n_test": r["n"], "ckpt": best(k)}
+        pf = os.path.join(pdir_root, f"round{k}", "summary.json")
+        if k > 0 and os.path.exists(pf):
+            s = json.load(open(pf))
+            row.update({"n_pos": s["n_pos"], "n_neg": s["n_neg"],
+                        "n_rej_cons": s["n_rejected_consistency"], "n_error": s["n_error"],
+                        "pos_vol_med_mm3": s["pos_vol_mm3"]["med"]})
+        table.append(row)
+    save_jsonl(table, os.path.join(results_dir, "iterative_table_monai.jsonl"))
+    print(json.dumps(table, indent=1))
+    try:
+        import matplotlib
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+        rs = [t["round"] for t in table]
+        ds = [t["test_dice"] for t in table]
+        plt.figure(figsize=(6, 4))
+        plt.plot(rs, ds, "o-")
+        plt.xlabel("pseudo-labeling round")
+        plt.ylabel("holdout tumor Dice (MONAI)")
+        for x, y in zip(rs, ds):
+            plt.annotate(f"{y:.3f}", (x, y), textcoords="offset points", xytext=(0, 8), fontsize=8)
+        plt.grid(alpha=0.3)
+        plt.tight_layout()
+        plt.savefig(os.path.join(results_dir, "iterative_dice_monai.png"), dpi=150)
+    except Exception as e:  # noqa
+        print("plot failed:", repr(e))
+
+
+if __name__ == "__main__":
+    main()

scripts_monai/monai_common.py

@@ -0,0 +1,208 @@
+"""Shared helpers for the MONAI iterative pseudo-labeling pipeline (scripts_monai/).
+
+Backbone: MONAI 1.6.x (pip) — monai.networks.nets.UNet, MONAI transforms /
+Dataset / DataLoader, and MONAI sliding-window inference. Everything
+study-specific (splits, gates, loss weighting, longitudinal consistency) is
+shared with the other pipelines so round results stay directly comparable:
+
+  * per-sample weighted Dice+CE loss — Pipeline A's convention (row "w";
+    pseudo-label rows down-weighted sample-by-sample)
+  * selection gates + head-relative consistency filter — Pipeline B's
+    implementation (scripts_nnu/nnu_common; absolute patient-space resampling
+    is unreliable across acquisitions, see README)
+  * DDP via torchrun for training, per-rank row sharding for inference —
+    Pipeline A's parallelism model
+
+Network: monai UNet 3D, 1 in / 2 out, channels 16→128 (~1.2M params) — the
+MONAI analogue of Pipeline A's Unet3D(base=16, depth=4).
+"""
+import os
+import sys
+
+# Data/run dirs: overridable for scratch runs. Code (src/, scripts_nnu/): next to this file.
+ROOT = os.environ.get("LONGITUDINAL_ROOT", "/mnt/b4/xfr/git26/longitudinal")
+_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, _REPO)
+sys.path.insert(0, os.path.join(_REPO, "scripts_nnu"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from src.common import dice  # noqa: E402
+
+# Selection gates / consistency filter shared with Pipeline B (pure numpy/scipy)
+from nnu_common import pos_mask, neg_frac_bg, consistency_filter, load_voxel_stats  # noqa: E402,F401
+
+WIN = 96
+
+
+# ---------------- network / loss ----------------
+
+def build_model(device):
+    from monai.networks.nets import UNet
+    net = UNet(
+        spatial_dims=3,
+        in_channels=1,
+        out_channels=2,
+        channels=(16, 32, 64, 128),
+        strides=(2, 2, 2),
+        num_res_units=2,
+        act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}),
+        norm="batch",
+    )
+    return net.to(device)
+
+
+class WeightedDiceCELoss(nn.Module):
+    """Per-sample (soft-dice + CE), weighted batch mean.
+
+    Same convention as Pipeline A's src/losses.py per_sample_loss (dice over
+    background+foreground averaged, dice_weight 0.5), built on the MONAI
+    DiceLoss. Per-sample reduction is required so accepted pseudo-label rows
+    can be down-weighted individually via row "w"; a batch-level loss
+    (MONAI DiceCELoss / nnU-Net CE+Dice) cannot express that.
+    """
+
+    def __init__(self, dice_weight=0.5):
+        super().__init__()
+        from monai.losses import DiceLoss
+        self.dice_weight = dice_weight
+        self.dice = DiceLoss(include_background=True, softmax=False, reduction="none")
+
+    def forward(self, logits, label, w):
+        logits = logits.float()
+        c = logits.size(1)
+        one = nn.functional.one_hot(label, c).permute(0, -1, *range(1, label.ndim)).float()
+        p = nn.functional.softmax(logits, dim=1)
+        dsc = self.dice(p, one)                      # (B, C, 1, 1, 1)
+        dsc = dsc.squeeze(-1).squeeze(-1).squeeze(-1).mean(dim=1)   # (B,)
+        ce = nn.functional.cross_entropy(logits, label, reduction="none")  # (B, z, y, x)
+        ce = ce.mean(dim=tuple(range(1, ce.ndim)))                      # (B,)
+        per = (1 - self.dice_weight) * ce + self.dice_weight * dsc
+        if torch.isnan(per).any() or torch.isinf(per).any():
+            per = torch.zeros_like(per)   # zero-gradient fallback: keeps DDP collectives in sync
+        return (per * w).sum() / w.sum().clamp(min=1e-6)
+
+
+# ---------------- data ----------------
+
+def train_transform(win=WIN):
+    """Channel-first MONAI chain: load → min-pad → augment → pos/neg crop → typed tensors.
+
+    Mirrors Pipeline A's augmentation set (per-axis flip p=0.5, rot90 in the
+    (y,x) plane p=0.1, brightness 1±0.1 p=0.3, Gaussian noise σ=0.01 p=0.4).
+    Patch sampling uses MONAI's foreground-aware RandCropByPosNegLabeld
+    (pos=neg=1) instead of Pipeline A's uniform random crop.
+    """
+    from monai.transforms import (
+        Compose, LoadImaged, EnsureChannelFirstd, EnsureTyped, SpatialPadd,
+        RandFlipd, RandRotate90d, RandScaleIntensityd, RandGaussianNoised,
+        RandCropByPosNegLabeld)
+    return Compose([
+        LoadImaged(keys=["pimg", "plabel"]),
+        EnsureChannelFirstd(keys=["pimg", "plabel"]),
+        SpatialPadd(keys=["pimg", "plabel"], spatial_size=(win, win, win)),
+        RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=0),
+        RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=1),
+        RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=2),
+        RandRotate90d(keys=["pimg", "plabel"], prob=0.1, max_k=3, spatial_axes=(1, 2)),
+        RandScaleIntensityd(keys=["pimg"], factors=0.1, prob=0.3),
+        RandGaussianNoised(keys=["pimg"], std=0.01, prob=0.4),
+        RandCropByPosNegLabeld(keys=["pimg", "plabel"], label_key="plabel",
+                               spatial_size=(win, win, win), pos=1, neg=1, num_samples=1),
+        EnsureTyped(keys=["pimg", "plabel"], dtype=[torch.float32, torch.long]),
+    ])
+
+
+def rows_for_dataset(rows):
+    out = []
+    for r in rows:
+        out.append({"key": r["key"], "pimg": r["pimg"], "plabel": r["plabel"],
+                    "w": float(r.get("w", 1.0))})
+    return out
+
+
+def _collate(batch):
+    # RandCropByPosNegLabeld(num_samples=1) yields a 1-item list per sample
+    items = [b[0] if isinstance(b, (list, tuple)) else b for b in batch]
+    return torch.utils.data.default_collate(items)
+
+
+def make_dataloader(rows, win=WIN, batch=3, workers=4):
+    from monai.data import Dataset
+    ds = Dataset(data=rows_for_dataset(rows), transform=train_transform(win))
+    return torch.utils.data.DataLoader(
+        ds, batch_size=batch, shuffle=True, num_workers=workers,
+        collate_fn=_collate, drop_last=len(ds) > batch, pin_memory=True,
+        persistent_workers=workers > 0)
+
+
+# ---------------- inference ----------------
+
+@torch.no_grad()
+def predict_probs(model, vol, device, win=WIN, overlap=0.5, tta=True, sw_batch=8):
+    """Tumor probability map (C, z, y, x → (z, y, x) numpy) for one volume.
+
+    MONAI sliding_window_inference (gaussian blend, `overlap`), plus the same
+    4-view flip TTA (identity + 3 axis flips) as Pipeline A.
+    """
+    from monai.inferers import sliding_window_inference
+    a = np.nan_to_num(np.asarray(vol, dtype=np.float32), nan=0.0, posinf=1.5, neginf=0.0)
+    a = np.clip(a, 0.0, 1.5)
+    t = torch.from_numpy(a).unsqueeze(0).unsqueeze(0).to(device)
+
+    def run(vt):
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            logits = sliding_window_inference(
+                vt, roi_size=(win, win, win), sw_batch_size=sw_batch,
+                predictor=model, mode="gaussian", overlap=overlap,
+                sw_device=device, device=device)
+        return torch.softmax(logits.float(), dim=1)[0, 1].cpu().numpy()
+
+    acc = run(t)
+    if not tta:
+        return acc
+    for ax in (0, 1, 2):
+        acc += np.flip(run(t.flip(ax + 2)), axis=ax)
+    return acc / 4.0
+
+
+def prob_dice(probs, lab, thr=0.5):
+    """Dice at probability threshold 0.5 — the convention of Pipeline A's
+    08_eval and Pipeline B's primary metric."""
+    return dice((probs >= thr).astype("uint8"), (lab > 0).astype("uint8"))
+
+
+def load_model(ckpt, device):
+    sd = torch.load(ckpt, map_location=device, weights_only=True)
+    net = build_model(device)
+    net.load_state_dict(sd.get("model", sd))
+    net.eval()
+    return net, sd
+
+
+# ---------------- distributed (torchrun, as in Pipeline A) ----------------
+
+def rank_info():
+    return (int(os.environ.get("RANK", 0)), int(os.environ.get("WORLD_SIZE", 1)),
+            int(os.environ.get("LOCAL_RANK", 0)))
+
+
+def init_dist():
+    if int(os.environ.get("WORLD_SIZE", 1)) > 1:
+        from datetime import timedelta
+        import torch.distributed as dist
+        dist.init_process_group("nccl", timeout=timedelta(minutes=30))
+
+
+def barrier():
+    if int(os.environ.get("WORLD_SIZE", 1)) > 1:
+        import torch.distributed as dist
+        dist.barrier()
+
+
+def destroy_dist():
+    if int(os.environ.get("WORLD_SIZE", 1)) > 1:
+        import torch.distributed as dist
+        dist.destroy_process_group()