"""Pseudo-label the unlabeled pool with an nnU-Net round model (multi-GPU predict). 1. Symlinks remaining pool volumes (not already in the dataset) into a predict folder. 2. Runs one nnUNetv2_predict per GPU (-num_parts/-part_id), --save_probabilities. 3. Selection gates (same as scripts/06_pseudo_label.py): pos: p_tumor >= tau_pos, median cleanup, largest-CC fraction >= min_cc_frac, volume within [vol_p2, vol_p98] of labeled tumor volumes (data/vols.json) neg: >= neg_frac of interior (vol > 0.02) voxels have p_bg >= tau_neg 4. Per-subject longitudinal consistency filter over accepted positive timepoints (head-relative tumor centroid distance + volume ratio; absolute patient-space grids are not comparable across acquisitions). Writes /rows.jsonl, /_label.nii.gz, /accepted.jsonl, /summary.json. Usage: python scripts_nnu/04_nnu_pseudo_label.py --pool data/manifests/unlabeled_pool.jsonl \ --out data/pseudo_nnu/round1 --gpus 3 """ import argparse import os import sys import json import shutil import numpy as np 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, read_nii_arr, write_arr import SimpleITK as sitk from nnu_common import (ROOT, d, nnu_env, predict_cmd, tumor_probs_from_npz, write_label, pos_mask, neg_frac_bg, consistency_filter, load_voxel_stats) def build_input_folder(rows, in_dir): shutil.rmtree(in_dir, ignore_errors=True) os.makedirs(in_dir, exist_ok=True) for r in rows: os.symlink(os.path.abspath(r["pimg"]), os.path.join(in_dir, r["key"] + "_0000.nii.gz")) def find_outputs(out_dir, ext): found = {} for part in os.listdir(out_dir): pd = os.path.join(out_dir, part) if not os.path.isdir(pd): continue for fn in os.listdir(pd): if fn.endswith(ext): found[fn[:-len(ext)]] = os.path.join(pd, fn) return found def main(): ap = argparse.ArgumentParser() ap.add_argument("--pool", required=True) ap.add_argument("--out", required=True) ap.add_argument("--already", default=None, help="jsonl of case keys already added to the dataset") ap.add_argument("--gpus", type=int, default=3) ap.add_argument("--chk", default="checkpoint_best.pth") ap.add_argument("--no-tta", action="store_true") ap.add_argument("--npp", type=int, default=2, help="predict subprocesses per GPU") 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") ap.add_argument("--skip-predict", action="store_true", help="reuse existing prediction outputs") ap.add_argument("--keep-pred", action="store_true") args = ap.parse_args() outname = os.path.basename(os.path.normpath(args.out)) 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)} rows = [r for r in pool_rows if r["key"] not in already] print(f"[nnu:pseudo:{outname}] pool={len(pool_rows)} already_in_dataset={len(already)} to_predict={len(rows)}", flush=True) 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) tag = outname in_dir = os.path.join(d("nnu/pred"), f"in_pool_{tag}") out_dir = os.path.join(d("nnu/pred"), f"out_pool_{tag}") if rows: build_input_folder(rows, in_dir) if not args.skip_predict: import subprocess os.makedirs(out_dir, exist_ok=True) env = nnu_env() procs = [] for g in range(args.gpus): e = dict(env) e["CUDA_VISIBLE_DEVICES"] = str(g) logf = os.path.join(d("logs/nnu"), f"pred_{tag}_part{g}.log") with open(logf, "wb") as lf: procs.append(subprocess.Popen(predict_cmd(in_dir, os.path.join(out_dir, f"part{g}"), args.gpus, g, chk=args.chk, tta=not args.no_tta, npp=args.npp), stdout=lf, stderr=subprocess.STDOUT, env=e, cwd=ROOT)) for g, p in enumerate(procs): rc = p.wait() if rc != 0: print(f"[nnu:pseudo:{outname}] part{g} failed rc={rc}; retrying once (skips finished cases)", flush=True) e = dict(env) e["CUDA_VISIBLE_DEVICES"] = str(g) logf = os.path.join(d("logs/nnu"), f"pred_{tag}_part{g}_retry.log") with open(logf, "wb") as lf: p2 = subprocess.Popen(predict_cmd(in_dir, os.path.join(out_dir, f"part{g}"), args.gpus, g, chk=args.chk, tta=not args.no_tta, npp=args.npp), stdout=lf, stderr=subprocess.STDOUT, env=e, cwd=ROOT) rc2 = p2.wait() if rc2 != 0: raise RuntimeError(f"prediction part{g} failed twice; see {logf}") npz = find_outputs(out_dir, ".npz") if os.path.isdir(out_dir) else {} rows_out, n_err = [], 0 for r in rows: 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 = tumor_probs_from_npz(npz[key]) 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_label(mask, lp, 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"[nnu:pseudo:{outname}] {key} ERR {e!r}", flush=True) rows_out.append(entry) n_pos0 = sum(1 for x in rows_out if x["role"] == "pos") n_neg0 = sum(1 for x in rows_out if x["role"] == "neg") n_rej = consistency_filter(rows_out, out, args.max_rel_dist, args.vol_ratio) for x in rows_out: if x["role"] == "rej": x["label"] = None accepted = [x for x in rows_out if x["role"] in ("pos", "neg")] save_jsonl(rows_out, os.path.join(out, "rows.jsonl")) save_jsonl(accepted, os.path.join(out, "accepted.jsonl")) posv = [x["vol_mm3"] for x in rows_out if x["role"] == "pos"] summ = { "n_pool_predicted": len(rows_out), "n_predicted": len(npz), "n_pos": sum(1 for x in rows_out if x["role"] == "pos"), "n_neg": sum(1 for x in rows_out if x["role"] == "neg"), "n_pos_before_consistency": n_pos0, "n_rejected_consistency": n_rej, "n_other_rej": sum(1 for x in rows_out if x["role"] == "rej"), "n_error": n_err, "pos_vol_mm3": {"med": float(np.median(posv)) if posv else 0, "p5": float(float(np.percentile(posv, 5))) if posv else 0, "p95": float(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, "chk": args.chk, "tta": not args.no_tta, } with open(os.path.join(out, "summary.json"), "w") as f: json.dump(summ, f, indent=1) print(f"[nnu:pseudo] round {outname}: {json.dumps(summ)}", flush=True) if not args.keep_pred and n_err == 0: shutil.rmtree(in_dir, ignore_errors=True) shutil.rmtree(out_dir, ignore_errors=True) elif n_err > 0: print(f"[nnu:pseudo:{outname}] kept prediction outputs ({in_dir}, {out_dir}) due to {n_err} errors; " f"rerun with --skip-predict --keep-pred after fixing", flush=True) if __name__ == "__main__": main()