refactor(core): improve data pipeline robustness and manifest generation
Refactor the data loading and preprocessing pipeline to handle edge cases in medical imaging data, including NaN/Inf values, shape mismatches, and numerical instability during training. - Update `01_build_ntuh_manifest.py` with improved regex for T1c detection, spine exclusion, and deduplication logic based on acquisition timestamps. - Enhance `PatchDataset` in `src/dataset.py` to handle NaN/Inf values, clip intensity ranges, and ensure label/image shape alignment via padding/trimming. - Add a zero-gradient fallback in `src/training.py` to prevent DDP synchronization failures when encountering NaN/Inf losses. - Add `scripts/scan_procs.py` for process monitoring. - Increase DataLoader timeout to prevent hangs during heavy I/O.
This commit is contained in:
parent
b6fa62a763
commit
33bc3c603f
4 changed files with 85 additions and 23 deletions
|
|
@ -1,22 +1,26 @@
|
|||
"""Build NTUH2022G4 labeled manifest: native T1c + tumor seg pairs from register_inv.
|
||||
"""Build NTUH2022G4 labeled manifest: brain native T1c + tumor seg pairs from register_inv.
|
||||
|
||||
Fast listing-only pass (no volume decoding); content validated during preprocessing.
|
||||
Rules:
|
||||
- series name must contain T1 and +C
|
||||
- exclude MRA/FLAIR/TOF/T2/SWI/DWI, _ROI1/_ROI re-exports, and spinal levels (e.g. T4-T8, T11-L3)
|
||||
- dedupe per (subject, case, acquisition timestamp): prefer _MPR_Tra, else lowest series number
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
import re
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from src.common import ROOT, save_jsonl
|
||||
|
||||
T1C_RE = re.compile(r"T1.*\+C")
|
||||
EXCL_RE = re.compile(r"MRA|FLAIR|TOF|T2|SWI|DWI", re.I)
|
||||
T1C_RE = re.compile(r"T1.*\+C|\+C.*T1")
|
||||
EXCL_RE = re.compile(r"MRA|FLAIR|TOF|T2|SWI|DWI|_ROI|ROI1", re.I)
|
||||
SPINE_RE = re.compile(r"[TLC]\s?\d+\s?[-–]\s?[TLC]\s?\d+|[TLC]\d{2}", re.I)
|
||||
|
||||
|
||||
def main(limit=None):
|
||||
out = os.path.join(ROOT, "data", "manifests", "ntuh.jsonl")
|
||||
reg_inv = "/mnt/pve/SRS/NTUH2022G4/register_inv"
|
||||
rows = []
|
||||
cands = {} # (subj, case, ts) -> list of (pref, ser, fname)
|
||||
n_subj = 0
|
||||
for s in sorted(os.listdir(reg_inv)):
|
||||
fp = os.path.join(reg_inv, s)
|
||||
|
|
@ -31,18 +35,27 @@ def main(limit=None):
|
|||
for f in sorted(fl):
|
||||
if not f.endswith(".nii.gz") or f.endswith((".seg.nii.gz", ".label.nii.gz")):
|
||||
continue
|
||||
if not T1C_RE.search(f) or EXCL_RE.search(f):
|
||||
if not T1C_RE.search(f) or EXCL_RE.search(f) or SPINE_RE.search(f):
|
||||
continue
|
||||
seg = f[: -len(".nii.gz")] + ".seg.nii.gz"
|
||||
if seg not in fl:
|
||||
if f[: -len(".nii.gz")] + ".seg.nii.gz" not in fl:
|
||||
continue
|
||||
rows.append({"key": f"ntuh_{s}_{c}", "subject": s, "case": c,
|
||||
"date": c2date(c), "img": os.path.join(cp, f),
|
||||
"label": os.path.join(cp, seg), "source": "ntuh"})
|
||||
mt = re.search(r"_(\d{14})_(\d+)\.nii\.gz$", f)
|
||||
ts = mt.group(1) if mt else f
|
||||
ser = int(mt.group(2)) if mt else 9999
|
||||
pref = 0 if "MPR_Tra" in f else 1
|
||||
cands.setdefault((s, c, ts), []).append((pref, ser, f))
|
||||
if limit and n_subj >= limit:
|
||||
break
|
||||
rows = []
|
||||
for (s, c, ts), lst in sorted(cands.items()):
|
||||
lst.sort(key=lambda x: (x[0], x[1]))
|
||||
f = lst[0][2]
|
||||
rows.append({"key": f"ntuh_{s}_{c}_{ts}", "subject": s, "case": c,
|
||||
"date": c2date(c), "img": os.path.join(reg_inv, s, c, f),
|
||||
"label": os.path.join(reg_inv, s, c, f[: -len(".nii.gz")] + ".seg.nii.gz"),
|
||||
"source": "ntuh", "dedup": len(lst)})
|
||||
save_jsonl(rows, out)
|
||||
print(f"subjects={n_subj} rows={len(rows)}")
|
||||
print(f"subjects={n_subj} rows={len(rows)} (deduped {sum(r['dedup'] for r in rows) - len(rows)})")
|
||||
|
||||
|
||||
def c2date(c):
|
||||
|
|
|
|||
40
scripts/scan_procs.py
Normal file
40
scripts/scan_procs.py
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
"""Scan processed volumes for non-finite/corrupt data. Writes data/bad_procs.json."""
|
||||
import os
|
||||
import sys
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
import glob
|
||||
import json
|
||||
import numpy as np
|
||||
import SimpleITK as sitk
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
|
||||
|
||||
def check(f):
|
||||
try:
|
||||
x = sitk.GetArrayFromImage(sitk.ReadImage(f))
|
||||
return os.path.basename(f), bool(np.isfinite(x).all()), float(x.max()) if x.size else -1.0
|
||||
except Exception as e: # noqa
|
||||
return os.path.basename(f), False, repr(e)
|
||||
|
||||
|
||||
def main():
|
||||
files = sorted(glob.glob(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "proc", "*.nii.gz")))
|
||||
print("checking", len(files), flush=True)
|
||||
bad = []
|
||||
with ProcessPoolExecutor(max_workers=48) as ex:
|
||||
futs = [ex.submit(check, f) for f in files]
|
||||
for i, fu in enumerate(as_completed(futs), 1):
|
||||
k, ok, mx = fu.result()
|
||||
if not ok or not np.isfinite(mx):
|
||||
bad.append(k)
|
||||
if i % 1000 == 0:
|
||||
print(i, "checked", len(bad), "bad", flush=True)
|
||||
print("BAD volumes:", len(bad))
|
||||
for b in sorted(bad):
|
||||
print(" ", b)
|
||||
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "bad_procs.json"), "w") as f:
|
||||
json.dump(sorted(bad), f, indent=1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -29,19 +29,27 @@ class PatchDataset(Dataset):
|
|||
row = self.rows[idx]
|
||||
w = float(row.get("w", 1.0))
|
||||
img = self._read(row["pimg"]).astype(np.float32)
|
||||
img = np.nan_to_num(img, nan=0.0, posinf=1.0, neginf=0.0)
|
||||
img = np.clip(img, 0.0, 1.5)
|
||||
# pad to multiple of 8 and at least the patch size (training) / 256 max (eval irrelevant)
|
||||
p = tuple(max((img.shape[i] + 7) // 8 * 8, self.patch[i] if self.train else 8) for i in range(3))
|
||||
if any(a < b for a, b in zip(img.shape, p)):
|
||||
img = np.pad(img, [(0, b - a) for a, b in zip(img.shape, p)])
|
||||
label = None
|
||||
if row.get("is_neg") and not (row.get("plabel") and os.path.exists(row["plabel"])):
|
||||
label = np.zeros(img.shape, dtype=np.int64)
|
||||
if row.get("plabel") and os.path.exists(row["plabel"]):
|
||||
label = self._read(row["plabel"]).astype(np.int64)
|
||||
q = tuple((label.shape[i] + 7) // 8 * 8 for i in range(3))
|
||||
label = np.pad(label, [(0, b - a) for a, b in zip(label.shape, q)])
|
||||
# pad/trim label to the same shape as img
|
||||
if label.shape != img.shape:
|
||||
lab_p = []
|
||||
for i in range(3):
|
||||
pad_i = 0
|
||||
pad_j = max(img.shape[i] - label.shape[i], 0)
|
||||
lab_p.append((pad_i, pad_j))
|
||||
label = np.pad(label, lab_p)
|
||||
label = label[: img.shape[0], : img.shape[1], : img.shape[2]]
|
||||
label = np.clip(np.rint(label).astype(np.int64), 0, 1)
|
||||
elif row.get("is_neg"):
|
||||
label = np.zeros(img.shape, dtype=np.int64)
|
||||
|
||||
if self.train:
|
||||
s0 = [self.rng.randint(0, img.shape[i] - self.patch[i] + 1) if img.shape[i] >= self.patch[i] else 0
|
||||
|
|
@ -98,5 +106,5 @@ def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persi
|
|||
ds = PatchDataset(rows, patch=patch, train=train, seed=seed)
|
||||
dl = DataLoader(ds, batch_size=batch_size, shuffle=train, num_workers=num_workers,
|
||||
collate_fn=_collate, drop_last=train and len(ds) > batch_size, pin_memory=True,
|
||||
persistent_workers=(train and num_workers > 0 and persistent))
|
||||
persistent_workers=(train and num_workers > 0 and persistent), timeout=900)
|
||||
return ds, dl
|
||||
|
|
@ -141,9 +141,10 @@ def train_ddp(rows, val_rows, args):
|
|||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
logits = ddp(img)
|
||||
losses = per_sample_loss(logits.float(), one_hot(lab, 2))
|
||||
if torch.isnan(losses).any() or torch.isinf(losses).any():
|
||||
# zero-gradient fallback: keeps DDP collectives in sync
|
||||
losses = torch.zeros_like(losses)
|
||||
loss = (losses * wts).sum() / wts.sum().clamp(min=1e-6)
|
||||
if torch.isnan(loss):
|
||||
continue
|
||||
opt.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(model.parameters(), 5.0)
|
||||
|
|
|
|||
Loading…
Reference in a new issue