Furen Xiao
8c813db209
Update README.md to include detailed project overview, environment requirements, data sources, and a comprehensive directory layout. Add nnU-Net pipeline documentation and directory descriptions. Update .gitignore to exclude nnU-Net specific directories and add new scripts directory for nnU-Net pipeline. Add initial nnU-Net pipeline scripts.
166 lines
No EOL
7.2 KiB
Python
166 lines
No EOL
7.2 KiB
Python
"""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, 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(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")
|
|
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):
|
|
base += [{"key": r["key"], "pimg": r["pimg"], "label": r["plabel"]} for r in load_jsonl(f)]
|
|
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() |