"""Orchestrates the nnU-Net iterative pseudo-labeling study (parallel to scripts/09). Round 0: train nnU-Net on the labeled patient-level train split (internal val = held-out patient-level val split, subject-disjoint). Round k: dataset grows with accepted pseudo-labels (pos + neg) from rounds 1..k; re-plan/preprocess; warm-start training (full weights) from round k-1 best checkpoint at a lower initial lr; then pseudo-label the remaining unlabeled pool (gates identical to scripts/06) and evaluate the model on the held-out patient-level test split. Usage: python scripts_nnu/06_nnu_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 load_jsonl, save_jsonl from nnu_common import ROOT, DATA, d, best_ckpt 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=250) ap.add_argument("--base-lr", type=float, default=1e-2) ap.add_argument("--round-epochs", type=int, default=75) ap.add_argument("--round-lr", type=float, default=1e-3) 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 mirroring TTA in 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(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") main_log = os.path.join(logdir, "nnu_iterative.log") pdir_root = d("data/pseudo_nnu") added_f = os.path.join(pdir_root, "added_keys.jsonl") base = [] for f in (train_f, val_f): for r in load_jsonl(f): if "img" not in r or not r.get("label"): raise SystemExit(f"{f} rows lack native img/label; rerun scripts/05_build_splits.py") base.append({"key": r["key"], "img": r["img"], "label": r["label"]}) save_jsonl(base, os.path.join(pdir_root, "base_rows.jsonl")) print(f"[nnu-orch] base dataset rows: {len(base)}", flush=True) def py(script, extra): print(f"[nnu-orch] $ python {script} {extra}", flush=True) with open(main_log, "a") as f: f.write(f"$ python {script} {extra}\n") subprocess.run([sys.executable, f"{ROOT}/scripts_nnu/{script}", *extra.split()], stdout=f, stderr=subprocess.STDOUT, check=True, cwd=ROOT) def prepare(rows): f = os.path.join(pdir_root, "current_rows.jsonl") save_jsonl(rows, f) py("01_nnu_prepare_dataset.py", f"--rows {f}") def plan(): py("02_nnu_plan_preprocess.py", f"--val {val_f}") def train(round, epochs, lr, warmstart): extra = f"--gpus {args.gpus} --epochs {epochs} --lr {lr} --log {os.path.join(logdir, f'nnu_train_r{round}.log')}" if warmstart: extra += f" --warmstart {warmstart}" py("03_nnu_train.py", extra) dst = os.path.join(ROOT, "runs_nnu", f"round{round}", "best_nnu.pth") os.makedirs(os.path.dirname(dst), exist_ok=True) shutil.copy2(best_ckpt(), dst) return dst def predict_pool(round, added): extra = (f"--pool {pool_f} --out {pdir_root}/round{round} --gpus {args.gpus} " 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 args.no_tta: extra += " --no-tta" if added: extra += f" --already {added_f}" py("04_nnu_pseudo_label.py", extra) return os.path.join(pdir_root, f"round{round}", "accepted.jsonl") def eval_test(round, out_name): extra = (f"--rows {test_f} --out {results_dir}/{out_name} --gpus {args.gpus}") if args.no_tta: extra += " --no-tta" py("05_nnu_eval_test.py", extra) return os.path.join(results_dir, out_name) # ---- round 0: baseline ---- prepare(base) plan() train(0, args.base_epochs, args.base_lr, None) accepted_f = [predict_pool(1, None)] save_jsonl(load_jsonl(accepted_f[0]), added_f) eval_test(0, "round0_test_nnu.json") for k in range(1, args.rounds + 1): rows = list(base) for f in accepted_f: for r in load_jsonl(f): if args.no_neg_pseudo and r["role"] == "neg": continue if all(r["key"] != x["key"] for x in rows): rows.append(r) prepare(rows) plan() train(k, args.round_epochs, args.round_lr, os.path.join(ROOT, "runs_nnu", f"round{k-1}", "best_nnu.pth")) accepted_f.append(predict_pool(k + 1, added_f)) save_jsonl([r for f in accepted_f for r in load_jsonl(f)], added_f) eval_test(k, f"round{k}_test_nnu.json") # ---- report (same layout as scripts/09) ---- table = [] for k in range(args.rounds + 1): resf = os.path.join(results_dir, f"round{k}_test_nnu.json") if not os.path.exists(resf): continue r = json.load(open(resf)) row = {"round": k, "test_dice": round(r["dice"], 4), "test_dice_hard": round(r["dice_hard"], 4), "n_test": r["n"], "ckpt": os.path.join(ROOT, "runs_nnu", f"round{k}", "best_nnu.pth")} 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_nnu.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 (nnU-Net)") 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_nnu.png"), dpi=150) except Exception as e: # noqa print("plot failed:", repr(e)) if __name__ == "__main__": main()