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.
307 lines
No EOL
10 KiB
Python
307 lines
No EOL
10 KiB
Python
"""Shared helpers for the nnU-Net iterative pseudo-labeling pipeline (scripts_nnu/).
|
|
|
|
Runs against the installed nnunetv2 2.8.1 build (checkpoint key `network_weights`,
|
|
probabilities saved as <case>.npz channel-first, folds in splits_final.json).
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
|
|
ROOT = os.environ.get("LONGITUDINAL_ROOT", "/mnt/b4/xfr/git26/longitudinal")
|
|
sys.path.insert(0, ROOT)
|
|
from src.common import load_jsonl, save_jsonl, read_nii_arr # noqa: E402
|
|
|
|
import numpy as np
|
|
import SimpleITK as sitk
|
|
from scipy import ndimage
|
|
|
|
DS_ID = 210
|
|
DS_NAME = f"Dataset{DS_ID:03d}_NTUH_T1C_PL"
|
|
TRAINER = "NTUHLPLTrainer"
|
|
CONFIG = "3d_fullres"
|
|
FOLD = 0
|
|
PLANS = "nnUNetPlans"
|
|
|
|
SCRIPTS_NNU = os.path.join(ROOT, "scripts_nnu")
|
|
TRAINERS_DIR = os.path.join(SCRIPTS_NNU, "trainers")
|
|
|
|
|
|
def nnu_root():
|
|
return d(os.path.join("nnu"))
|
|
|
|
|
|
def raw_ds():
|
|
return os.path.join(nnu_root(), "raw", DS_NAME)
|
|
|
|
|
|
def preproc_ds():
|
|
return os.path.join(nnu_root(), "preprocessed", DS_NAME)
|
|
|
|
|
|
def results_root():
|
|
return os.path.join(nnu_root(), "results")
|
|
|
|
|
|
def model_folder():
|
|
return os.path.join(results_root(), DS_NAME, f"{TRAINER}__{PLANS}__{CONFIG}")
|
|
|
|
|
|
def fold_dir():
|
|
return os.path.join(model_folder(), f"fold_{FOLD}")
|
|
|
|
|
|
def best_ckpt():
|
|
return os.path.join(fold_dir(), "checkpoint_best.pth")
|
|
|
|
|
|
def final_ckpt():
|
|
return os.path.join(fold_dir(), "checkpoint_final.pth")
|
|
|
|
|
|
def d(name):
|
|
p = os.path.join(ROOT, name)
|
|
os.makedirs(p, exist_ok=True)
|
|
return p
|
|
|
|
|
|
def nnu_env(epoch=None, lr=None, warmstart=None):
|
|
env = dict(os.environ)
|
|
env["nnUNet_raw"] = os.path.join(nnu_root(), "raw")
|
|
env["nnUNet_preprocessed"] = os.path.join(nnu_root(), "preprocessed")
|
|
env["nnUNet_results"] = results_root()
|
|
env["nnUNet_extTrainer"] = TRAINERS_DIR
|
|
env.pop("NNU_PL_EPOCHS", None)
|
|
env.pop("NNU_PL_LR", None)
|
|
env.pop("NNU_PL_WARMSTART", None)
|
|
env.pop("CUDA_VISIBLE_DEVICES", None)
|
|
if epoch is not None:
|
|
env["NNU_PL_EPOCHS"] = str(epoch)
|
|
if lr is not None:
|
|
env["NNU_PL_LR"] = repr(float(lr))
|
|
if warmstart:
|
|
env["NNU_PL_WARMSTART"] = str(warmstart)
|
|
return env
|
|
|
|
|
|
def run(cmd, log, env=None, retries=2):
|
|
os.makedirs(os.path.dirname(os.path.abspath(log)), exist_ok=True)
|
|
cmd_s = " ".join(str(c) for c in cmd)
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
with open(log, "a") as f:
|
|
f.write(f"$ (attempt {attempt + 1}) " + cmd_s + "\n")
|
|
f.flush()
|
|
print(f"$ (attempt {attempt + 1}) " + cmd_s, flush=True)
|
|
subprocess.run(cmd, cwd=ROOT, stdout=f, stderr=subprocess.STDOUT,
|
|
env=env, check=True)
|
|
return
|
|
except subprocess.CalledProcessError:
|
|
if attempt == retries:
|
|
raise
|
|
print(f"[nnu] command failed, retrying in 60s: {cmd_s}", flush=True)
|
|
time.sleep(60)
|
|
|
|
|
|
def train_cmd(gpus):
|
|
return ["nnUNetv2_train", str(DS_ID), CONFIG, str(FOLD), "-tr", TRAINER,
|
|
"-p", PLANS, "-num_gpus", str(gpus)]
|
|
|
|
|
|
def plan_preprocess_cmd(npp=8):
|
|
return ["nnUNetv2_plan_and_preprocess", "-d", str(DS_ID), "-c", CONFIG,
|
|
"-np", str(npp), "--clean", "--no_pbar"]
|
|
|
|
|
|
def predict_cmd(in_dir, out_dir, gpus, part, chk="checkpoint_best.pth", tta=True,
|
|
npp=2, nps=2):
|
|
cmd = ["nnUNetv2_predict", "-i", str(in_dir), "-o", str(out_dir), "-d", str(DS_ID),
|
|
"-c", CONFIG, "-tr", TRAINER, "-p", PLANS, "-f", str(FOLD), "-chk", chk,
|
|
"--save_probabilities", "--continue_prediction", "--disable_progress_bar",
|
|
"-num_parts", str(gpus), "-part_id", str(part),
|
|
"-npp", str(npp), "-nps", str(nps)]
|
|
if not tta:
|
|
cmd.append("--disable_tta")
|
|
return cmd
|
|
|
|
|
|
# ---------------- raw dataset (nnU-Net format) ----------------
|
|
|
|
def make_raw_dataset(rows, channel="0000"):
|
|
"""rows: list of {key, pimg, label}. Rebuilds imagesTr/labelsTr symlinks + dataset.json."""
|
|
img_dir = os.path.join(raw_ds(), "imagesTr")
|
|
lab_dir = os.path.join(raw_ds(), "labelsTr")
|
|
shutil.rmtree(img_dir, ignore_errors=True)
|
|
shutil.rmtree(lab_dir, ignore_errors=True)
|
|
os.makedirs(img_dir, exist_ok=True)
|
|
os.makedirs(lab_dir, exist_ok=True)
|
|
seen = set()
|
|
for r in rows:
|
|
key, img, lab = r["key"], r["pimg"], r["label"]
|
|
if key in seen:
|
|
raise ValueError(f"duplicate case id in dataset: {key}")
|
|
seen.add(key)
|
|
if not os.path.exists(img):
|
|
raise FileNotFoundError(f"image missing for {key}: {img}")
|
|
if not os.path.exists(lab):
|
|
raise FileNotFoundError(f"label missing for {key}: {lab}")
|
|
os.symlink(os.path.abspath(img), os.path.join(img_dir, f"{key}_{channel}.nii.gz"))
|
|
os.symlink(os.path.abspath(lab), os.path.join(lab_dir, f"{key}.nii.gz"))
|
|
ds_json = {"channel_names": {"0": "t1c"},
|
|
"labels": {"background": 0, "tumor": 1},
|
|
"numTraining": len(seen),
|
|
"file_ending": ".nii.gz"}
|
|
with open(os.path.join(raw_ds(), "dataset.json"), "w") as f:
|
|
json.dump(ds_json, f, indent=1)
|
|
return len(seen)
|
|
|
|
|
|
def dataset_case_keys():
|
|
img_dir = os.path.join(raw_ds(), "imagesTr")
|
|
out = []
|
|
if os.path.isdir(img_dir):
|
|
for fn in sorted(os.listdir(img_dir)):
|
|
if fn.endswith("_0000.nii.gz"):
|
|
out.append(fn[:-len("_0000.nii.gz")])
|
|
return out
|
|
|
|
|
|
def write_splits(train_keys, val_keys):
|
|
val_keys = sorted(set(val_keys))
|
|
train_keys = sorted(set(k for k in train_keys if k not in val_keys))
|
|
folds = [{"train": train_keys, "val": val_keys}] * 5
|
|
with open(os.path.join(preproc_ds(), "splits_final.json"), "w") as f:
|
|
json.dump(folds, f)
|
|
return len(train_keys), len(val_keys)
|
|
|
|
|
|
# ---------------- label writing / selection (same gates as scripts/06) ----------------
|
|
|
|
def write_label(arr_u01, path, ref_itk):
|
|
img = sitk.GetImageFromArray(arr_u01.astype(np.uint8))
|
|
img.CopyInformation(ref_itk)
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
sitk.WriteImage(img, path, True)
|
|
|
|
|
|
def pos_mask(p_tumor, tau_pos, min_cc_frac, vol_lo, vol_hi):
|
|
"""Returns (mask, cc_frac, vol_mm3) if the positive gates pass, else None."""
|
|
m = p_tumor >= tau_pos
|
|
if m.sum() == 0:
|
|
return None
|
|
m = ndimage.median_filter(m, size=(3, 3, 3))
|
|
if m.sum() == 0:
|
|
return None
|
|
lab, n = ndimage.label(m)
|
|
sizes = ndimage.sum(m, lab, range(1, n + 1))
|
|
big = (lab == (int(np.argmax(sizes)) + 1)).astype(np.uint8)
|
|
cc_frac = float(big.sum()) / float(m.sum())
|
|
vol_mm3 = int(big.sum())
|
|
if cc_frac >= min_cc_frac and vol_lo <= vol_mm3 <= vol_hi:
|
|
return big, cc_frac, vol_mm3
|
|
return None
|
|
|
|
|
|
def neg_frac_bg(p_tumor, vol, tau_neg, min_frac, min_interior=5000):
|
|
interior = vol > 0.02
|
|
if interior.sum() < min_interior:
|
|
return None
|
|
frac = float(((1.0 - p_tumor)[interior] >= tau_neg).mean())
|
|
return frac if frac >= min_frac else None
|
|
|
|
|
|
def tumor_probs_from_npz(npz_path):
|
|
"""<case>.npz probabilities: channel-first (C,H,W,D), C=2 (bg, tumor)."""
|
|
probs = np.load(npz_path, allow_pickle=False)["probabilities"]
|
|
if probs.shape[0] != 2:
|
|
raise ValueError(f"unexpected probability channels {probs.shape} in {npz_path}")
|
|
return probs[1]
|
|
|
|
|
|
# ---------------- longitudinal consistency (head-relative) ----------------
|
|
#
|
|
# The absolute patient-space resampling used by scripts/06 cannot work for this
|
|
# data: each acquisition has its own scanner/patient coordinate system (table
|
|
# offset, head orientation), so true labels from two visits of the same subject
|
|
# resample to ~0 dice. Instead we gate on frame-invariant features:
|
|
# * tumor centroid relative to the head centroid (crop space, 1mm voxels)
|
|
# * tumor volume ratio
|
|
# Two timepoints are compatible if they agree with at least one accepted
|
|
# neighbor on both; greedy removal until stable (same scheme as scripts/06).
|
|
|
|
def head_centroid(vol):
|
|
m = vol > 0.02
|
|
if m.sum() < 5000:
|
|
return None
|
|
lab, n = ndimage.label(m)
|
|
if n == 0:
|
|
return None
|
|
sizes = ndimage.sum(m, lab, range(1, n + 1))
|
|
return np.array(ndimage.center_of_mass(lab == (int(np.argmax(sizes)) + 1)))
|
|
|
|
|
|
def rel_tumor_features(mask, vol):
|
|
"""Returns (centroid_rel_to_head [z,y,x] in mm, vol_mm3) or None."""
|
|
if int(mask.sum()) == 0:
|
|
return None
|
|
tc = np.array(ndimage.center_of_mass(mask))
|
|
hc = head_centroid(vol)
|
|
rel = tc - hc if hc is not None else None
|
|
return rel, int(mask.sum())
|
|
|
|
|
|
def _compatible(fa, fb, max_rel_dist_mm, vol_ratio_max):
|
|
ra, va = fa
|
|
rb, vb = fb
|
|
if ra is not None and rb is not None and np.linalg.norm(ra - rb) > max_rel_dist_mm:
|
|
return False
|
|
if va > 0 and vb > 0:
|
|
ratio = max(va / vb, vb / va)
|
|
if ratio > vol_ratio_max:
|
|
return False
|
|
return True
|
|
|
|
|
|
def consistency_filter(rows, out_dir, max_rel_dist_mm=40.0, vol_ratio_max=10.0):
|
|
"""Greedy removal of positive timepoints inconsistent with accepted neighbors.
|
|
|
|
masks from <out_dir>/<key>_label.nii.gz, volumes from row["pimg"]. Returns n_rejected.
|
|
"""
|
|
by_subj = {}
|
|
for r in rows:
|
|
if r["role"] == "pos":
|
|
by_subj.setdefault(r["subject"], []).append(r)
|
|
rejected = 0
|
|
for subj, tps in by_subj.items():
|
|
if len(tps) < 2:
|
|
continue
|
|
tps.sort(key=lambda x: (x.get("date") or ""))
|
|
feats = {}
|
|
for i, r in enumerate(tps):
|
|
mask = read_nii_arr(os.path.join(out_dir, r["key"] + "_label.nii.gz")).astype(bool)
|
|
vol = read_nii_arr(r["pimg"]).astype("float32")
|
|
feats[i] = rel_tumor_features(mask, vol)
|
|
accepted = list(range(len(tps)))
|
|
while True:
|
|
changed = False
|
|
for ai in list(accepted):
|
|
neigh = [b for b in (ai - 1, ai + 1) if b in accepted]
|
|
if not neigh or feats[ai] is None:
|
|
continue
|
|
if not any(_compatible(feats[ai], feats[b], max_rel_dist_mm, vol_ratio_max)
|
|
for b in neigh if feats[b] is not None):
|
|
tps[ai]["role"] = "rejected"
|
|
accepted.remove(ai)
|
|
rejected += 1
|
|
changed = True
|
|
break
|
|
if not changed:
|
|
break
|
|
return rejected
|
|
|
|
|
|
def load_voxel_stats():
|
|
p = os.path.join(ROOT, "data/vols.json")
|
|
return json.load(open(p)) if os.path.exists(p) else {} |