Furen Xiao
695bbb6e11
Update the nnU-Net pipeline to operate on native (unprocessed) volumes instead of preprocessed ones. This allows nnU-Net to utilize its own `plan_and_preprocess` logic for resampling, cropping, and normalization, ensuring Pipeline B remains distinct from Pipelines A and C. - Update `scripts/05_build_splits.py` to include native `img` and `label` paths in the manifest rows. - Modify `scripts_nnu/01_nnu_prepare_dataset.py` to symlink native volumes and implement a label fix-up mechanism for non-conforming grids. - Update `scripts_nnu/04_nnu_pseudo_label.py` and `05_nnu_eval_test.py` to use native image paths and perform selection/evaluation in physical mm on the native grid. - Refactor `scripts_nnu/nnu_common.py` to handle native-grid label alignment and volume-based selection gates. - Update `README.md` to document the preprocessing differences between Pipelines A/C and Pipeline B.
201 lines
No EOL
10 KiB
Python
201 lines
No EOL
10 KiB
Python
"""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. Volumes are NATIVE (unresampled/uncropped/unnormalized); nnU-Net
|
|
applies the plan's preprocessing itself at predict time.
|
|
2. Runs one nnUNetv2_predict per GPU (-num_parts/-part_id), --save_probabilities.
|
|
nnU-Net applies the plan's preprocessing at predict time and resamples the
|
|
probabilities (and argmax seg) back to each case's native grid, so the
|
|
selection gates run directly on that grid, in physical mm:
|
|
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 head-interior voxels have p_bg >= tau_neg
|
|
4. Per-subject longitudinal consistency filter over accepted positive timepoints
|
|
(head-relative tumor centroid distance in mm + volume ratio; absolute
|
|
patient-space grids are not comparable across acquisitions).
|
|
Writes <out>/rows.jsonl, <out>/<key>_label.nii.gz (native grid), <out>/accepted.jsonl,
|
|
<out>/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["img"]), 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)
|
|
if any("img" not in r for r in pool_rows):
|
|
raise SystemExit(f"{args.pool} rows lack native 'img'; rerun scripts/05_build_splits.py")
|
|
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"), "img": r["img"], "pimg": r.get("pimg"), "label": None,
|
|
"role": "rej", "vol_mm3": 0, "maxp": None}
|
|
try:
|
|
img_itk = sitk.ReadImage(r["img"])
|
|
vol = read_nii_arr(r["img"]).astype("float32")
|
|
pt = tumor_probs_from_npz(npz[key])
|
|
if pt.shape != vol.shape:
|
|
raise ValueError(f"prob shape {pt.shape} != native image {vol.shape} for {key}")
|
|
entry["maxp"] = round(float(pt.max()), 4)
|
|
voxel_vol_mm3 = float(np.prod(np.asarray(img_itk.GetSpacing(), dtype=np.float64)))
|
|
got = pos_mask(pt, args.tau_pos, args.min_cc_frac, vol_lo, vol_hi, voxel_vol_mm3)
|
|
if got is not None:
|
|
mask, cc_frac, vol_mm3 = got
|
|
lp = os.path.join(out, key + "_label.nii.gz")
|
|
write_label(mask, lp, img_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=img_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, image_key="img")
|
|
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() |