diff --git a/scripts/01_build_ntuh_manifest.py b/scripts/01_build_ntuh_manifest.py index 2d1941b..ed2852b 100644 --- a/scripts/01_build_ntuh_manifest.py +++ b/scripts/01_build_ntuh_manifest.py @@ -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): diff --git a/scripts/scan_procs.py b/scripts/scan_procs.py new file mode 100644 index 0000000..c7db157 --- /dev/null +++ b/scripts/scan_procs.py @@ -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() \ No newline at end of file diff --git a/src/dataset.py b/src/dataset.py index 52c52e1..9883379 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -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 \ No newline at end of file diff --git a/src/training.py b/src/training.py index 807804f..f5ed2c8 100644 --- a/src/training.py +++ b/src/training.py @@ -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)