xfr/longitudinal
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(core): improve data pipeline robustness and manifest generation

Browse source 
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:
Furen Xiao 2026-09-25 22:12:32 +08:00
parent b6fa62a763
commit 33bc3c603f
4 changed files with 85 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
scripts/01_build_ntuh_manifest.py
View file

@ -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"})
if limit and n_subj >= limit:
break
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
View 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()

20
src/dataset.py
View file

@ -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)])
label = label[: img.shape[0], : img.shape[1], : img.shape[2]]
# 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

5
src/training.py
View file

@ -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)