365 lines
No EOL
13 KiB
Python
365 lines
No EOL
13 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", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
sys.path.insert(0, ROOT)
|
|
from src.common import DATA, d, load_jsonl, save_jsonl, read_nii_arr, head_mask_from_image, largest_cc # 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 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) ----------------
|
|
#
|
|
# Pipeline B feeds nnU-Net the NATIVE volumes (no resampling / cropping /
|
|
# intensity normalization beforehand): nnUNetv2_plan_and_preprocess does all of
|
|
# that itself (nonzero-bbox crop, resample to the plan's target spacing,
|
|
# per-channel normalization, .b2nd storage).
|
|
|
|
def same_frame(a, b):
|
|
return (a.GetSize() == b.GetSize() and a.GetSpacing() == b.GetSpacing()
|
|
and a.GetOrigin() == b.GetOrigin() and a.GetDirection() == b.GetDirection())
|
|
|
|
|
|
def native_label_target(img_path, lab_path, link_path):
|
|
"""Label target for a raw-dataset row: 0/1 on the image's native grid.
|
|
|
|
Symlinks the label when it already sits on the image grid with values
|
|
{0, 1}; otherwise writes a converted copy to link_path (nearest-neighbor
|
|
warp to the image grid + binarization to 0/1). This is label fix-up only —
|
|
no common-grid resampling, cropping, or normalization.
|
|
"""
|
|
img_itk = sitk.ReadImage(img_path)
|
|
lab_itk = sitk.ReadImage(lab_path)
|
|
if same_frame(img_itk, lab_itk):
|
|
vals = set(np.unique(sitk.GetArrayFromImage(lab_itk)).tolist())
|
|
if vals <= {0, 1}:
|
|
os.symlink(os.path.abspath(lab_path), link_path)
|
|
return
|
|
warped = sitk.Resample(lab_itk, img_itk, sitk.Transform(), sitk.sitkNearestNeighbor, 0.0)
|
|
out = sitk.GetImageFromArray((sitk.GetArrayFromImage(warped) > 0).astype(np.uint8))
|
|
out.CopyInformation(img_itk)
|
|
os.makedirs(os.path.dirname(link_path), exist_ok=True)
|
|
sitk.WriteImage(out, link_path, True)
|
|
|
|
|
|
def make_raw_dataset(rows, channel="0000"):
|
|
"""rows: list of {key, img, label} — img/label are native (unprocessed) NIfTI
|
|
paths. Rebuilds imagesTr/labelsTr + dataset.json. Images are symlinked
|
|
as-is; labels via native_label_target. See section note."""
|
|
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["img"], 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"))
|
|
native_label_target(img, 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) ----------------
|
|
#
|
|
# Prediction outputs are ALREADY on each case's native grid: nnUNetv2_predict
|
|
# (--save_probabilities) resamples tumor probabilities and the argmax seg back
|
|
# to the original input shape. Gates below therefore run directly on the native
|
|
# volume; physical mm is recovered from the native voxel spacing.
|
|
|
|
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, voxel_vol_mm3=1.0):
|
|
"""Returns (mask, cc_frac, vol_mm3) if the positive gates pass, else None.
|
|
|
|
voxel_vol_mm3 is the native voxel volume in mm3 (Product of NIfTI spacing);
|
|
1.0 for 1mm processed volumes (Pipelines A/C)."""
|
|
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()) * voxel_vol_mm3
|
|
if cc_frac >= min_cc_frac and vol_lo <= vol_mm3 <= vol_hi:
|
|
return big, cc_frac, vol_mm3
|
|
return None
|
|
|
|
|
|
def _interior_mask(vol, min_interior=5000):
|
|
"""Head-interior mask. Processed [0,1] percentile-normalized volumes keep
|
|
the legacy `vol > 0.02` rule; native raw-intensity volumes use the robust
|
|
percentile / largest-CC head mask (see src/common.head_mask_from_image)."""
|
|
if float(vol.max()) <= 1.0 + 1e-6:
|
|
m = vol > 0.02
|
|
else:
|
|
m = head_mask_from_image(vol)
|
|
return m if m.sum() >= min_interior else None
|
|
|
|
|
|
def neg_frac_bg(p_tumor, vol, tau_neg, min_frac, min_interior=5000):
|
|
interior = _interior_mask(vol, min_interior)
|
|
if interior is None:
|
|
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, computed in
|
|
# each case's own grid with physical mm (native spacing for Pipeline B, 1mm for
|
|
# A/C):
|
|
# * tumor centroid relative to the head centroid (offset in mm)
|
|
# * tumor volume ratio (mm3)
|
|
# 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 = _interior_mask(vol)
|
|
if m is None:
|
|
return None
|
|
m = largest_cc(m)
|
|
if int(m.sum()) == 0:
|
|
return None
|
|
return np.array(ndimage.center_of_mass(m))
|
|
|
|
|
|
def rel_tumor_features(mask, vol, spacing=(1.0, 1.0, 1.0)):
|
|
"""Returns (centroid_rel_to_head [z,y,x] in mm, vol_mm3) or None.
|
|
|
|
spacing is the NIfTI voxel spacing in array order (z, y, x); (1, 1, 1) for
|
|
1mm processed volumes (Pipelines A/C)."""
|
|
if int(mask.sum()) == 0:
|
|
return None
|
|
spacing = np.asarray(spacing, dtype=np.float64)
|
|
tc = np.array(ndimage.center_of_mass(mask))
|
|
hc = head_centroid(vol)
|
|
rel = (tc - hc) * spacing if hc is not None else None
|
|
return rel, int(mask.sum()) * float(spacing.prod())
|
|
|
|
|
|
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, image_key="pimg"):
|
|
"""Greedy removal of positive timepoints inconsistent with accepted neighbors.
|
|
|
|
masks from <out_dir>/<key>_label.nii.gz, volumes + native spacing from
|
|
row[image_key] ("pimg": 1mm processed, Pipelines A/C; "img": native,
|
|
Pipeline B). 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[image_key]).astype("float32")
|
|
if mask.shape != vol.shape:
|
|
raise ValueError(f"mask shape {mask.shape} != volume {r[image_key]} shape {vol.shape} for {r['key']}")
|
|
sp = np.asarray(sitk.ReadImage(r[image_key]).GetSpacing()[::-1], dtype=np.float64)
|
|
feats[i] = rel_tumor_features(mask, vol, sp)
|
|
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(DATA, "vols.json")
|
|
return json.load(open(p)) if os.path.exists(p) else {} |