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.

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")
+T1C_RE = re.compile(r"T1.*\+C|\+C.*T1")
-EXCL_RE = re.compile(r"MRA|FLAIR|TOF|T2|SWI|DWI", re.I)
+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 f[: -len(".nii.gz")] + ".seg.nii.gz" not in fl:
-                if seg not in fl:
                     continue
-                rows.append({"key": f"ntuh_{s}_{c}", "subject": s, "case": c,
+                mt = re.search(r"_(\d{14})_(\d+)\.nii\.gz$", f)
-                             "date": c2date(c), "img": os.path.join(cp, f),
+                ts = mt.group(1) if mt else f
-                             "label": os.path.join(cp, seg), "source": "ntuh"})
+                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):

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

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))
+            # pad/trim label to the same shape as img
-            label = np.pad(label, [(0, b - a) for a, b in zip(label.shape, q)])
+            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

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)