longitudinal/scripts_nnu/04_nnu_pseudo_label.py

"""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 <out>/rows.jsonl, <out>/<key>_label.nii.gz, <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["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()