xfr/longitudinal
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(monai): add MONAI pipeline C implementation

Browse source 
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.
This commit is contained in:
Furen Xiao 2026-09-26 11:04:20 +08:00
parent 8c813db209
commit 77adc2b3af
9 changed files with 961 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

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

6
AGENTS.md
View file

@ -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.
- No lint/test tooling is configured yet; run scripts directly with `python <script>`.
- 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>` from the repo root.

78
README.md
View file

@ -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
```

64
scripts_monai/01_monai_build_rows.py Normal file
View file

@ -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()

156
scripts_monai/02_monai_train.py Normal file
View file

@ -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()

169
scripts_monai/03_monai_pseudo_label.py Normal file
View file

@ -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()

103
scripts_monai/04_monai_eval_test.py Normal file
View file

@ -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()

180
scripts_monai/05_monai_run_iterative.py Normal file
View file

@ -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()

208
scripts_monai/monai_common.py Normal file
View file

@ -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()