feat: initial project structure

Add .gitignore, AGENTS.md, scripts directory, and src directory to initialize the repository.

.gitignore

@@ -0,0 +1,5 @@
+data/
+runs/
+results/
+logs/
+__pycache__/

AGENTS.md

@@ -0,0 +1,23 @@
+# AGENTS.md
+
+## Environment
+
+This project uses a conda environment named `longitudinal` (Python 3.14).
+
+Activate with:
+
+```bash
+source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
+```
+
+Key packages: torch 2.14 (+cu126), torchvision, numpy, scipy, pandas, scikit-learn, scikit-image, matplotlib, nibabel, SimpleITK.
+
+## Project
+
+Longitudinal (repeated-measures) analysis of medical imaging data. Repo is at an early stage — see README.md for any project notes.
+
+## Conventions
+
+- All Python commands must run inside the `longitudinal` conda environment.
+- GPU is available (CUDA 12.6); use `torch.device('cuda')` when appropriate.
+- No lint/test tooling is configured yet; run scripts directly with `python <script>`.

scripts/01_build_ntuh_manifest.py

@@ -0,0 +1,54 @@
+"""Build NTUH2022G4 labeled manifest: native T1c + tumor seg pairs from register_inv.
+
+Fast listing-only pass (no volume decoding); content validated during preprocessing.
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import re
+import json
+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)
+
+
+def main(limit=None):
+    out = os.path.join(ROOT, "data", "manifests", "ntuh.jsonl")
+    reg_inv = "/mnt/pve/SRS/NTUH2022G4/register_inv"
+    rows = []
+    n_subj = 0
+    for s in sorted(os.listdir(reg_inv)):
+        fp = os.path.join(reg_inv, s)
+        if not os.path.isdir(fp):
+            continue
+        n_subj += 1
+        for c in os.listdir(fp):
+            cp = os.path.join(fp, c)
+            if not os.path.isdir(cp):
+                continue
+            fl = set(os.listdir(cp))
+            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):
+                    continue
+                seg = f[: -len(".nii.gz")] + ".seg.nii.gz"
+                if seg 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
+    save_jsonl(rows, out)
+    print(f"subjects={n_subj} rows={len(rows)}")
+
+
+def c2date(c):
+    m = re.match(r"case?(\d{4})\.(\d{2})\.(\d{2})\.", c)
+    return f"{m.group(1)}-{m.group(2)}-{m.group(3)}" if m else c
+
+
+if __name__ == "__main__":
+    main()

scripts/02_build_m6_dataset.py

@@ -0,0 +1,147 @@
+"""Build M6-2025 manifests: T1c volumes, labeled when GTV exists.
+
+GTV lives in CT space; CT is registered to T1c (similarity->affine MI + BSpline)
+and GTV warped (nearest) into T1c native space.
+Outputs:
+  data/manifests/m6_labeled.jsonl    {key, subject, date, img, label(warped GTV), source:'m6'}
+  data/manifests/m6_unlabeled.jsonl  {key, subject, date, img, label:None, source:'m6'}
+Usage: python scripts/02_build_m6_dataset.py [--workers 4] [--skip-reg]
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import re
+import json
+import argparse
+import numpy as np
+import SimpleITK as sitk
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from src.common import ROOT, d, save_jsonl, read_nii_img
+
+BASE = "/mnt/pve/WORKSPACE/M6-2025/nii"
+T1C_RE = re.compile(r"T1.*\+C|fl3d.*\+.*c", re.I)
+EXCL_RE = re.compile(r"FLAIR|DTI|vibe|dixon|t2|SWI|T2|MPR_Cor", re.I)
+
+
+def pick_t1c(mrd):
+    cands = []
+    for f in sorted(os.listdir(mrd)):
+        if not f.endswith(".nii.gz"):
+            continue
+        b = f[: -len(".nii.gz")]
+        if T1C_RE.search(b) and not EXCL_RE.search(b):
+            pref = 0
+            if "MPR_Tra" in b:
+                pref = -2
+            if "_c" in b.lower() or "+C" in b:
+                pref -= 1
+            cands.append((pref, f))
+    cands.sort(key=lambda x: (x[0], x[1]))
+    return os.path.join(mrd, cands[0][1]) if cands else None
+
+
+def discover():
+    labeled, unlabeled = [], []
+    for pid in sorted(os.listdir(BASE)):
+        fp = os.path.join(BASE, pid)
+        if not os.path.isdir(fp):
+            continue
+        for date in sorted(os.listdir(fp)):
+            dp = os.path.join(fp, date)
+            if not os.path.isdir(dp) or not re.match(r"\d{8}$", date):
+                continue
+            mrd = os.path.join(dp, "MR")
+            if not os.path.isdir(mrd):
+                continue
+            t1c = pick_t1c(mrd)
+            if not t1c:
+                continue
+            gtv = os.path.join(dp, "RT", "TV", "Struct_GTV.nii.gz")
+            ct = os.path.join(dp, "RT", "ct_image.nii.gz")
+            key = f"m6_{pid}_{date}"
+            row = {"key": key, "subject": pid, "date": date, "img": t1c, "source": "m6"}
+            if os.path.exists(gtv) and os.path.exists(ct):
+                labeled.append({**row, "label": gtv, "ct": ct, "gtv_registered": False})
+            else:
+                unlabeled.append({**row, "label": None})
+    return labeled, unlabeled
+
+
+def register_ct_to_t1c(ct, t1c, threads=12):
+    import time
+    t_start = time.time()
+    try:  # noqa
+        sitk.CommonProperties.SetGlobalDefaultNumberOfThreads(threads)
+    except AttributeError:
+        pass
+    ct = sitk.Cast(ct, sitk.sitkFloat32)
+    t1c = sitk.Cast(t1c, sitk.sitkFloat32)
+    r = sitk.ImageRegistrationMethod()
+    r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32)
+    r.SetOptimizerAsRegularStepGradientDescent(
+        learningRate=10.0, minStep=0.5, numberOfIterations=600,
+        relaxationFactor=0.5, gradientMagnitudeTolerance=1e-5,
+        estimateLearningRate=sitk.ImageRegistrationMethod.EachIteration)
+    r.SetOptimizerScalesFromPhysicalShift()
+    r.SetInitialTransform(sitk.CenteredTransformInitializer(ct, t1c, sitk.Transform(3, sitk.sitkSimilarity)))
+    r.SetShrinkFactorsPerLevel([8, 4, 2, 1])
+    r.SetSmoothingSigmasPerLevel([4, 2, 1, 0])
+    r.SetInterpolator(sitk.sitkLinear)
+    aff = r.Execute(ct, t1c)
+    print("affine done %.0fs" % (time.time() - t_start), flush=True)
+    return aff
+
+
+def reg_worker(row):
+    reg_dir = d("data/m6reg")
+    key = row["key"]
+    t1c = read_nii_img(row["img"])
+    ct = read_nii_img(row["ct"])
+    gtv = read_nii_img(row["label"])
+    tr = register_ct_to_t1c(ct, t1c)
+    out = os.path.join(reg_dir, key + "_gtv.nii.gz")
+    warped = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0)
+    sitk.WriteImage(warped, out, True)
+    z = sitk.GetArrayFromImage(warped)
+    return key, os.path.exists(out) and (z > 0.5).sum() > 0, int((z > 0.5).sum())
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--workers", type=int, default=2)
+    ap.add_argument("--skip-reg", action="store_true")
+    args = ap.parse_args()
+    labeled, unlabeled = discover()
+    print(f"m6 labeled candidates: {len(labeled)}, unlabeled: {len(unlabeled)}")
+    reg_dir = d("data/m6reg")
+    if not args.skip_reg:
+        todo = [r for r in labeled if not os.path.exists(os.path.join(reg_dir, r["key"] + "_gtv.nii.gz"))]
+        print(f"registering {len(todo)} GTVs with {args.workers} workers")
+        with ProcessPoolExecutor(max_workers=args.workers) as ex:
+            futs = {ex.submit(reg_worker, r): r for r in todo}
+            for fu in as_completed(futs):
+                k = futs[fu]["key"]
+                try:
+                    key, ok, nvox = fu.result()
+                    print(f"  {key}: ok={ok} nvox={nvox}", flush=True)
+                except Exception as e:  # noqa
+                    print(f"  {k}: ERR {e}", flush=True)
+        # refresh manifest state
+        labeled, unlabeled = discover()
+    for r in labeled:
+        p = os.path.join(reg_dir, r["key"] + "_gtv.nii.gz")
+        if os.path.exists(p):
+            r["label"] = p
+        else:
+            r["label"] = None
+    keep = [r for r in labeled if r.get("label")]
+    unlabeled = [r for r in unlabeled if r["key"] not in {k["key"] for k in keep}]
+    for r in keep + unlabeled:
+        r.pop("ct", None)
+    save_jsonl(keep, os.path.join(ROOT, "data/manifests/m6_labeled.jsonl"))
+    save_jsonl(unlabeled, os.path.join(ROOT, "data/manifests/m6_unlabeled.jsonl"))
+    print(f"final labeled={len(keep)} unlabeled={len(unlabeled)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/03_scan_lee_t1c.py

@@ -0,0 +1,175 @@
+"""Scan lee longitudinal dataset for brain T1c (post-contrast T1) series.
+
+Each MR timepoint dir: <sid>/<YYYYMMDD>_MR_<n>/ contains top-level per-slice
+DICOM-dump txt files '<sid>_<date>_MR_<ser>_<slice>.txt' and one jpg folder
+named like '<description>_<sid>_<date>' holding '<sid>_<date>_MR_<ser>_<slice>_<inst>.jpg'.
+
+T1c series = contrasted (bolus tags) + T1-weighted 3D sequence.
+Writes: data/manifests/lee_t1c_raw.jsonl and (after subset selection)
+        data/manifests/lee_t1c_selected.jsonl
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import re
+import json
+import argparse
+import random
+from src.common import ROOT, save_jsonl, load_jsonl
+
+BASE = "/mnt/t24/Public/lee"
+T1_NAME_RE = re.compile(r"t1|tfl|spgr|mp2rage|tse3d|vfl|mpage", re.I)
+EXCL_RE = re.compile(r"flair|\bt2\b|dwi|dti|mra|angi|swi|bold|\bpp2d|\bpp3d|perf|t2\*|t2star", re.I)
+
+
+def parse_txt(p):
+    """Return dict tagtuple -> value (value column)."""
+    t = {}
+    try:
+        with open(p, errors="ignore") as f:
+            for line in f:
+                m = re.match(r'"\((\d{4}),(\d{4})\)","([A-Z]+)","(\d+)",(.*),"(.*)"\s*$', line.strip())
+                if m:
+                    t[(int(m.group(1), 16), int(m.group(2), 16))] = m.group(5).strip().strip('"')
+    except OSError:
+        pass
+    return t
+
+
+def is_t1c(t):
+    bolus = any(k in t for k in ((24, 4161), (24, 4168), (8, 307)))
+    if not bolus:
+        return False, "no bolus"
+    name = (t.get((24, 36), "") + " " + t.get((24, 33), "")).lower()
+    if EXCL_RE.search(name):
+        return False, "excluded name " + name[:30]
+    if not T1_NAME_RE.search(name):
+        seq = t.get((24, 32), "")
+        et = t.get((24, 129), "")
+        ok = seq in ("GR", "SE", "GR\\IR", "SE ", "GR ")
+        try:
+            et_ok = et.replace('"', "").split()[0].replace(" ", "")[:1] != "" and float(et.split("\\")[0].strip()) < 30
+        except (ValueError, IndexError):
+            et_ok = True
+        if not (ok and et_ok):
+            return False, "not T1W"
+    return True, name[:40]
+
+
+def scan_timepoint(sid, date, tpd):
+    """Returns list of candidate T1c series for one MR timepoint dir."""
+    txts = {}
+    jpg_dir = None
+    for f in os.listdir(tpd):
+        full = os.path.join(tpd, f)
+        if os.path.isdir(full):
+            jpg_dir = full
+            continue
+        m = re.match(rf"{sid}_{date}_MR_(\d+)_(\d+)\.txt$", f)
+        if m:
+            txts.setdefault(m.group(1), []).append((int(m.group(2)), full))
+    if jpg_dir is None:
+        return []
+    jpgs = set(os.listdir(jpg_dir))
+    cand = []
+    for ser, sl in txts.items():
+        if not sl:
+            continue
+        first = min(sl)[1]
+        t = parse_txt(first)
+        ok, why = is_t1c(t)
+        if not ok:
+            continue
+        rows = t.get((40, 16), ""); cols = t.get((40, 17), "")
+        # count jpgs of this series
+        pat = re.compile(rf"{sid}_{date}_MR_{ser}_(\d+)_(\d+)\.jpg$")
+        nslices = 0
+        for jf in jpgs:
+            if pat.match(jf):
+                nslices += 1
+                break
+        nslices_all = sum(1 for jf in jpgs if pat.match(jf))
+        cand.append({"ser": ser, "txt_first": first, "n_slices": nslices_all,
+                     "rows": rows, "cols": cols, "why": why,
+                     "jpg_dir": jpg_dir})
+    # prefer 3D (MRAcq) & most slices
+    def score(c):
+        t = parse_txt(c["txt_first"])
+        acq = "3" if "3D" in t.get((24, 35), "") else "2"
+        return (acq, c["n_slices"])
+    cand.sort(key=score, reverse=True)
+    return cand
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--limit-subj", type=int, default=0)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--max-multi-subj", type=int, default=60)
+    ap.add_argument("--max-single-subj", type=int, default=0)
+    ap.add_argument("--max-subj-timepoints", type=int, default=6)
+    ap.add_argument("--select-only", action="store_true")
+    args = ap.parse_args()
+    random.seed(args.seed)
+    raw_path = os.path.join(ROOT, "data/manifests/lee_t1c_raw.jsonl")
+    if args.select_only:
+        rows = load_jsonl(raw_path)
+        print(f"select-only: {len(rows)} raw rows")
+    else:
+        rows = []
+        sids = [s for s in os.listdir(BASE) if os.path.isdir(os.path.join(BASE, s)) and not s.endswith(".complete")]
+    if not args.select_only:
+        if args.limit_subj:
+            sids = sids[: args.limit_subj]
+        for i, sid in enumerate(sorted(sids)):
+            tpd_dir = os.path.join(BASE, sid)
+            for tp in os.listdir(tpd_dir):
+                if not re.match(r"\d{8}_MR_", tp) or tp.endswith(".complete"):
+                    continue
+                tpd = os.path.join(tpd_dir, tp)
+                if not os.path.isdir(tpd):
+                    continue
+                date = tp[:8]
+                for c in scan_timepoint(sid, date, tpd):
+                    rows.append({"sid": sid, "date": date, "tp": tp,
+                                 "jpg_dir": c["jpg_dir"], "ser": c["ser"], "n_slices": c["n_slices"],
+                                 "txt_first": c["txt_first"], "why": c["why"],
+                                 "key": f"lee_{sid}_{date}_s{c['ser']}"})
+            if i % 100 == 0 and i:
+                print(f"scanned {i}/{len(sids)} subjects, {len(rows)} T1c candidates", flush=True)
+        save_jsonl(rows, raw_path)
+        print(f"total T1c candidates: {len(rows)}")
+
+    # subset selection: prefer subjects with more timepoints (longitudinal consistency)
+    by_subj = {}
+    for r in rows:
+        by_subj.setdefault(r["sid"], []).append(r)
+    for v in by_subj.values():
+        v.sort(key=lambda x: x["date"])
+    multi = {k: v for k, v in by_subj.items() if len(v) >= 2}
+    single = {k: v for k, v in by_subj.items() if len(v) == 1}
+    rnd = random.Random(args.seed)
+    depth = {k: len(v) for k, v in multi.items()}
+    cand = [k for k in multi if depth[k] >= 3]
+    deep = sorted(cand, key=lambda k: -depth[k])
+    n_deep = max(1, args.max_multi_subj * 2 // 3)
+    pick = deep[:n_deep]
+    rest = [k for k in cand if k not in pick]
+    rnd.shuffle(rest)
+    pick += rest[: args.max_multi_subj - len(pick)]
+    sel = []
+    for k in pick:
+        v = multi[k]
+        if len(v) > args.max_subj_timepoints and args.max_subj_timepoints >= 2:
+            v = [v[int(j * (len(v) - 1) / (args.max_subj_timepoints - 1))] for j in range(args.max_subj_timepoints)]
+        sel.extend(v)
+    rnd.shuffle(list(single))
+    for k in list(single)[: args.max_single_subj]:
+        sel.extend(single[k])
+    sel.sort(key=lambda x: (x["sid"], x["date"], int(x["ser"])))
+    save_jsonl(sel, os.path.join(ROOT, "data/manifests/lee_t1c_selected.jsonl"))
+    print(f"selected: {len(sel)} timepoints from {len(set(r['sid'] for r in sel))} subjects")
+
+
+if __name__ == "__main__":
+    main()

scripts/04_reconstruct_lee.py

@@ -0,0 +1,190 @@
+"""Reconstruct lee T1c volumes from JPG slices + DICOM txt metadata.
+
+The timepoint dir holds only sampled per-slice txts (typically slices 1, 2, N);
+slice positions are reconstructed by fitting a linear IPP(s) model, validated
+against all available samples.
+Usage: python scripts/04_reconstruct_lee.py [--manifest ...selected.jsonl] [--workers 48]
+Output: data/lee_nii/<key>.nii.gz (uint8, native grid)
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import re
+import json
+import argparse
+import numpy as np
+import SimpleITK as sitk
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from PIL import Image
+from src.common import ROOT, d, load_jsonl
+
+
+LINE_RE = re.compile(
+    r'^\s*"\((\d{4}),(\d{4})\)","([A-Z]+)","(\d+)",(?:\"([^\"]*)\"|([^,]+)),"')
+
+
+def parse_txt(p):
+    ipp = spacing = iop = rows = cols = None
+    try:
+        with open(p, errors="ignore") as f:
+            for raw in f:
+                m = LINE_RE.match(raw)
+                if not m:
+                    continue
+                tag = (int(m.group(1), 16), int(m.group(2), 16))
+                if tag not in ((32, 50), (40, 48), (32, 55), (40, 16), (40, 17)):
+                    continue
+                val = (m.group(5) if m.group(5) is not None else m.group(6)).strip()
+                if not val:
+                    continue
+                try:
+                    if tag == (32, 50):
+                        ipp = np.array([float(x) for x in val.split("\\")[:3]])
+                    elif tag == (40, 48):
+                        spacing = np.array([float(x) for x in val.split("\\")[:2]])
+                    elif tag == (32, 55):
+                        iop = np.array([float(x) for x in val.split("\\")[:6]])
+                    elif tag == (40, 16):
+                        rows = int(float(val))
+                    elif tag == (40, 17):
+                        cols = int(float(val))
+                except (ValueError, IndexError):
+                    continue
+        if ipp is None or spacing is None or iop is None or rows is None or cols is None:
+            return None
+        return ipp, spacing, iop, rows, cols
+    except OSError:
+        return None
+
+
+def reconstruct(row):
+    sid, date, ser, jpg_dir = row["sid"], row["date"], row["ser"], row["jpg_dir"]
+    out = os.path.join(d("data/lee_nii"), row["key"] + ".nii.gz")
+    if os.path.exists(out):
+        return row["key"], True, "cached"
+    pat = re.compile(rf"^{re.escape(sid)}_{date}_MR_{ser}_(\d+)_(\d+)\.jpg$")
+    tpat = re.compile(rf"^{re.escape(sid)}_{date}_MR_{ser}_(\d+)\.txt$")
+    tp_dir = os.path.dirname(jpg_dir)
+    jpg_map, txt_samples = {}, []
+    try:
+        for f in os.listdir(jpg_dir):
+            m = pat.match(f)
+            if m:
+                sl, inst = int(m.group(1)), int(m.group(2))
+                if sl not in jpg_map or inst < jpg_map[sl][1]:
+                    jpg_map[sl] = (f, inst)
+        for f in os.listdir(tp_dir):
+            m = tpat.match(f)
+            if m:
+                p = parse_txt(os.path.join(tp_dir, f))
+                if p is not None:
+                    txt_samples.append((int(m.group(1)), p))
+    except OSError as e:
+        return row["key"], False, f"listdir fail {e!r}"
+    if not jpg_map or not txt_samples:
+        return row["key"], False, f"no jpg({len(jpg_map)}) or txt({len(txt_samples)})"
+    txt_samples.sort()
+    # majority (rows, cols, pixel-spacing) as reference geometry
+    from collections import Counter
+    geo = Counter((p[3], p[4], round(p[1][0], 4), round(p[1][1], 4)) for _, p in txt_samples)
+    (rows, cols, ps0, ps1) = geo.most_common(1)[0][0]
+    ref = [p for _, p in txt_samples if (p[3], p[4]) == (rows, cols) and np.allclose(p[1], (ps0, ps1), atol=1e-3)]
+    good_idx = [sl for sl, p in txt_samples if (p[3], p[4]) == (rows, cols)]
+    if len(good_idx) < 2 or not ref:
+        return row["key"], False, f"only {len(good_idx)} consistent geometry samples"
+    ipp0, ps, iop0, _, _ = ref[len(ref) // 2]
+    # linear IPP model (fit on geometry-consistent samples only)
+    good = [(sl, p) for sl, p in txt_samples if p[3] == rows and p[4] == cols and np.allclose(p[1], ps, atol=1e-3)]
+    if len(good) < 2:
+        return row["key"], False, "fewer than 2 consistent geometry samples"
+    (s_lo, p_lo), (s_hi, p_hi) = good[0], good[-1]
+    step = (p_hi[0] - p_lo[0]) / max(s_hi - s_lo, 1)
+    resid = max(np.linalg.norm(p[0] - p_lo[0] - (s - s_lo) * step) for s, p in good)
+    if resid > 0.35 * np.linalg.norm(step):
+        return row["key"], False, f"non-linear slice positions resid={resid:.3f}"
+    u, v = iop0[:3], iop0[3:]
+    n = np.cross(u, v)
+    # the stack direction must follow the nominal slice normal (rejects multiplanar 2D exports)
+    if abs(float(np.dot(step, n))) < 0.6 * float(np.linalg.norm(step)):
+        return row["key"], False, "stack not along slice normal"
+    sl_list = sorted(jpg_map)
+
+    def ipp_of(s):
+        return p_lo[0] + (s - s_lo) * step
+    ipp_all = {sl: ipp_of(sl) for sl in sl_list}
+    r0 = min(sl_list, key=lambda s: np.dot(ipp_all[s] - ipp_all[sl_list[0]], n))
+    p_ref = ipp_all[r0]
+    zs, pos = [], []
+    for sl in sl_list:
+        off = ipp_all[sl] - p_ref
+        zs.append(float(np.dot(off, n)))
+        pos.append((int(round(float(np.dot(off, u)) / ps[0])), int(round(float(np.dot(off, v)) / ps[1]))))
+    zs = np.array(zs)
+    dmed = np.median(np.abs(np.diff(np.sort(zs))))
+    if dmed <= 0 or not np.isfinite(dmed):
+        return row["key"], False, "bad slice spacing"
+    nz = int(round((zs.max() - zs.min()) / dmed)) + 1
+    ro_min = min(p[0] for p in pos)
+    ro_max = max(p[0] for p in pos)
+    co_min = min(p[1] for p in pos)
+    co_max = max(p[1] for p in pos)
+    vol = np.zeros((rows + (ro_max - ro_min), cols + (co_max - co_min), nz), dtype=np.uint8)
+    nread = 0
+    for idx, (sl, jf, (ro, co)) in enumerate(zip(sl_list, [jpg_map[s][0] for s in sl_list], pos)):
+        zi = int(round(zs[idx] / dmed))
+        try:
+            im = np.asarray(Image.open(os.path.join(jpg_dir, jf)).convert("L"), dtype=np.uint8)
+        except Exception:
+            continue
+        if im.shape[:2] != (rows, cols):
+            if im.shape[:2] == (cols, rows):
+                im = im.T.copy()
+            else:
+                continue
+        vol[ro - ro_min: ro - ro_min + rows, co - co_min: co - co_min + cols, zi] = im
+        nread += 1
+    if nread < 0.95 * len(sl_list) or vol.max() == 0:
+        return row["key"], False, f"only {nread}/{len(sl_list)} slices read"
+    ext = (vol.shape[0] * ps[0], vol.shape[1] * ps[1], vol.shape[2] * dmed)
+    if min(ext) < 60 or max(ext) > 350:
+        return row["key"], False, f"plausible extent failed {tuple(round(e, 1) for e in ext)}mm"
+    origin = p_ref + (ro_min * ps[0]) * u + (co_min * ps[1]) * v
+    direction = tuple(float(x) for x in np.concatenate([u, v, n]))
+    img = sitk.GetImageFromArray(vol)
+    img.SetOrigin(tuple(float(x) for x in origin))
+    img.SetSpacing((float(ps[0]), float(ps[1]), float(dmed)))
+    img.SetDirection(direction)
+    os.makedirs(os.path.dirname(out), exist_ok=True)
+    sitk.WriteImage(img, out, True)
+    return row["key"], True, f"{vol.shape} n={nread}"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--manifest", default=os.path.join(ROOT, "data/manifests/lee_t1c_selected.jsonl"))
+    ap.add_argument("--workers", type=int, default=48)
+    args = ap.parse_args()
+    rows = load_jsonl(args.manifest)
+    rows = [r for r in rows if not os.path.exists(os.path.join(d("data/lee_nii"), r["key"] + ".nii.gz"))]
+    print(f"todo={len(rows)} workers={args.workers}")
+    ok = err = 0
+    with ProcessPoolExecutor(max_workers=args.workers) as ex:
+        futs = {ex.submit(reconstruct, r): r for r in rows}
+        for i, fu in enumerate(as_completed(futs), 1):
+            try:
+                k, good, msg = fu.result()
+            except Exception as e:  # noqa
+                k, good, msg = futs[fu]["key"], False, repr(e)
+            if good:
+                ok += 1
+            else:
+                err += 1
+                if err <= 40:
+                    print(" ERR", k, msg, flush=True)
+            if i % 100 == 0:
+                print(f"  {i}/{len(futs)} ok={ok} err={err}", flush=True)
+    print(f"done ok={ok} err={err}")
+
+
+if __name__ == "__main__":
+    main()

scripts/05_build_splits.py

@@ -0,0 +1,121 @@
+"""Build patient-level train/val/test splits + unlabeled pool + tumor volume stats.
+
+Inputs: data/manifests/{ntuh,m6_labeled,m6_unlabeled,lee_t1c_selected}.jsonl + data/proc/*
+Outputs:
+  data/manifests/split_train.jsonl split_val.jsonl split_test.jsonl   (labeled rows, w=1.0)
+  data/manifests/unlabeled_pool.jsonl                                   (m6 + lee, labeled subjects removed)
+  data/vols.json                                                        (labeled tumor volume stats, mm3)
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import json
+import argparse
+import random
+import numpy as np
+from src.common import ROOT, d, load_jsonl, save_jsonl, read_nii_arr
+
+
+def proc_row(r, prefix):
+    key = r["key"]
+    p = os.path.join(d("data/proc"), key + ".nii.gz")
+    if not os.path.exists(p):
+        return None
+    row = {"key": key, "subject": f"{prefix}_{r['subject']}", "pimg": p}
+    lab = os.path.join(d("data/proc"), key + "_label.nii.gz")
+    if r.get("label") and os.path.exists(lab):
+        row["plabel"] = lab
+    return row
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--frac-val", type=float, default=0.05)
+    ap.add_argument("--frac-test", type=float, default=0.10)
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+    random.seed(args.seed)
+
+    labeled_rows = []
+    for name, prefix in (("ntuh", "ntuh"), ("m6_labeled", "m6")):
+        for r in load_jsonl(os.path.join(ROOT, "data", "manifests", name + ".jsonl")):
+            if not r.get("label"):
+                continue
+            p = proc_row(r, prefix)
+            if p:
+                labeled_rows.append(p)
+    print(f"labeled processed rows: {len(labeled_rows)}")
+
+    subjects = sorted({r["subject"] for r in labeled_rows})
+    random.shuffle(subjects)
+    n_test = max(1, int(len(subjects) * args.frac_test))
+    n_val = max(1, int(len(subjects) * args.frac_val))
+    test_subj = set(subjects[:n_test])
+    val_subj = set(subjects[n_test:n_test + n_val])
+    train_subj = set(subjects[n_test + n_val:])
+    print(f"subjects: {len(subjects)}  train={len(train_subj)} val={len(val_subj)} test={len(test_subj)}")
+
+    splits = {"train": [], "val": [], "test": []}
+    for r in labeled_rows:
+        r["w"] = 1.0
+        if r["subject"] in test_subj:
+            splits["test"].append(r)
+        elif r["subject"] in val_subj:
+            splits["val"].append(r)
+        else:
+            splits["train"].append(r)
+    for k, v in splits.items():
+        save_jsonl(v, os.path.join(ROOT, "data/manifests", f"split_{k}.jsonl"))
+        print(f"  split_{k}: {len(v)} volumes / {len(set(r['subject'] for r in v))} subjects")
+
+    labeled_subjects = train_subj | val_subj | test_subj
+
+    # unlabeled pool
+    pool = []
+    for r in load_jsonl(os.path.join(ROOT, "data/manifests/m6_unlabeled.jsonl")):
+        if f"m6_{r['subject']}" in labeled_subjects:
+            continue
+        p = proc_row(r, "m6")
+        if p:
+            p["date"] = r.get("date")
+            p["source"] = "m6"
+            pool.append(p)
+    lee_sel_path = os.path.join(ROOT, "data/manifests/lee_t1c_selected.jsonl")
+    lee_rows = {}
+    if os.path.exists(lee_sel_path):
+        for r in load_jsonl(lee_sel_path):
+            nii = os.path.join(d("data/lee_nii"), r["key"] + ".nii.gz")
+            proc = os.path.join(d("data/proc"), r["key"] + ".nii.gz")
+            if os.path.exists(proc):
+                p = {"key": r["key"], "subject": f"lee_{r['sid']}", "pimg": proc,
+                     "date": r["date"], "source": "lee"}
+                lee_rows[r["key"]] = p
+    pool.extend(lee_rows.values())
+    pool.sort(key=lambda r: (r["subject"], r.get("date", "")))
+    save_jsonl(pool, os.path.join(ROOT, "data/manifests/unlabeled_pool.jsonl"))
+    print(f"unlabeled pool: {len(pool)} volumes ({sum(1 for r in pool if r['source']=='m6')} m6, "
+          f"{sum(1 for r in pool if r['source']=='lee')} lee) from {len(set(r['subject'] for r in pool))} subjects")
+
+    # tumor volume stats (1mm voxels == mm3)
+    vols = []
+    for r in labeled_rows:
+        try:
+            l = read_nii_arr(r["plabel"])
+            n = int((l > 0.5).sum())
+            vols.append({"key": r["key"], "subject": r["subject"], "vol_mm3": n})
+        except Exception:
+            pass
+    v = np.array([x["vol_mm3"] for x in vols])
+    stats = {"n": int(v.size),
+             "p2": float(np.percentile(v, 2)), "p5": float(np.percentile(v, 5)),
+             "p50": float(np.percentile(v, 50)), "p95": float(np.percentile(v, 95)),
+             "p98": float(np.percentile(v, 98)), "max": float(v.max()),
+             "zero_frac": float((v == 0).mean())}
+    save_jsonl(vols, os.path.join(ROOT, "data/vols.jsonl"))
+    with open(os.path.join(ROOT, "data/vols.json"), "w") as f:
+        json.dump(stats, f, indent=1)
+    print("tumor volumes (mm3):", stats)
+
+
+if __name__ == "__main__":
+    main()

scripts/06_pseudo_label.py

@@ -0,0 +1,211 @@
+"""Iterative pseudo-labeling round.
+
+torchrun --standalone --nproc_per_node=3 scripts/06_pseudo_label.py \
+    --ckpt runs/round0/best.pt --unlabeled data/manifests/unlabeled_pool.jsonl --out data/pseudo/round1
+
+Per volume: sliding-window tumor probabilities (TTA). Selection:
+  pos: p_tumor >= tau_pos, largest-CC fraction >= min_cc_frac, volume within [vol_lo, vol_hi]
+  neg: >= neg_frac of interior voxels have p_bg >= tau_neg
+Then a per-subject longitudinal consistency filter over accepted positive timepoints.
+Writes: <out>/rows.jsonl, <out>/<key>.pseudo.nii.gz, <out>/summary.json
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import json
+import argparse
+import numpy as np
+import SimpleITK as sitk
+from scipy import ndimage
+from src.common import ROOT, d, load_jsonl, save_jsonl, read_nii_arr
+from src import training
+
+
+def grid_info(key):
+    p = os.path.join(d("data/procmeta"), key + ".json")
+    m = json.load(open(p))
+    origin = np.array(m["origin"])
+    R = np.array(m["direction"]).reshape(3, 3)  # row_dir, col_dir, slice_dir
+    cv = np.array(m["crop_vox"])
+    o = origin + cv[0] * R[0] + cv[1] * R[1] + cv[2] * R[2]
+    return o, R
+
+
+def grid_itk(shape, key):
+    o, R = grid_info(key)
+    img = sitk.GetImageFromArray(np.zeros(shape, np.uint8))
+    img.SetSpacing((1.0, 1.0, 1.0))
+    img.SetOrigin(tuple(float(x) for x in o))
+    img.SetDirection(tuple(float(x) for x in R.flatten()))
+    return img
+
+
+def load_pseudo(outdir, key):
+    return read_nii_arr(os.path.join(outdir, key + ".pseudo.nii.gz")).astype(bool)
+
+
+def dice_a_on_b(keyA, arrA, keyB, arrB):
+    a_itk = sitk.GetImageFromArray(arrA.astype(np.uint8) % 255)
+    a_itk.CopyInformation(grid_itk(arrA.shape, keyA))
+    ref = grid_itk(arrB.shape, keyB)
+    am = sitk.GetArrayFromImage(sitk.Resample(a_itk, ref, sitk.Transform(), sitk.sitkNearestNeighbor, 0)).astype(bool)
+    bm = arrB
+    inter = (am & bm).sum()
+    return float(2 * inter / max(int(am.sum()) + int(bm.sum()), 1))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--unlabeled", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--tau-pos", type=float, default=0.95)
+    ap.add_argument("--tau-neg", type=float, default=0.98)
+    ap.add_argument("--neg-frac", type=float, default=0.90)
+    ap.add_argument("--vol-qp", type=float, nargs=2, default=[2, 98])
+    ap.add_argument("--min-cc-frac", type=float, default=0.2)
+    ap.add_argument("--cons-dice", type=float, default=0.30)
+    ap.add_argument("--win", type=int, default=96)
+    ap.add_argument("--step", type=int, default=64)
+    args = ap.parse_args()
+    import torch
+    rank = int(os.environ.get("RANK", 0))
+    world = int(os.environ.get("WORLD_SIZE", 1))
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    if world > 1:
+        import torch.distributed as dist
+        dist.init_process_group("nccl")
+    torch.cuda.set_device(local_rank)
+    device = f"cuda:{local_rank}"
+
+    outname = os.path.basename(os.path.normpath(args.out))
+    out = d(os.path.join("data/pseudo", outname))
+    rows = load_jsonl(args.unlabeled)
+    done = set()
+    part = os.path.join(out, f"part{rank}.jsonl")
+    if os.path.exists(part):
+        for r in load_jsonl(part):
+            done.add(r["key"])
+    rows = [r for r in rows if r["key"] not in done]
+    shard = rows[rank::world]
+
+    vstats = {}
+    if os.path.exists(os.path.join(ROOT, "data/vols.json")):
+        vstats = json.load(open(os.path.join(ROOT, "data/vols.json")))
+    vol_lo = vstats.get(f"p{args.vol_qp[0]:.0f}", 1.0)
+    vol_hi = vstats.get(f"p{args.vol_qp[1]:.0f}", 50000.0)
+    if rank == 0:
+        print(f"[pseudo:rank0] pool={len(rows)} shard={len(shard)} vol_range=[{vol_lo:.0f},{vol_hi:.0f}]mm3", flush=True)
+
+    ckpt = torch.load(args.ckpt, map_location=device, weights_only=True)
+    model = training.build_model(device=device)
+    model.load_state_dict(ckpt["model"])
+    model.eval()
+
+    buf = []
+    for i, r in enumerate(shard, 1):
+        key = r["key"]
+        try:
+            vol = read_nii_arr(r["pimg"]).astype(np.float32)
+            p = training.sliding_window_probs(model, vol, device, args.win, args.step, tta=True)
+            out_row = {"key": key, "subject": r["subject"], "date": r.get("date"),
+                       "source": r.get("source"), "pimg": r["pimg"], "plabel": None,
+                       "role": "rej", "vol_mm3": 0, "maxp": round(float(p.max()), 4)}
+            m = p >= args.tau_pos
+            if m.sum() > 0:
+                m = ndimage.median_filter(m, size=(3, 3, 3))
+                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())
+                if cc_frac >= args.min_cc_frac and vol_lo <= vol_mm3 <= vol_hi:
+                    out_row.update({"role": "pos", "vol_mm3": vol_mm3, "cc_frac": round(cc_frac, 3)})
+                    sitk.WriteImage(sitk.GetImageFromArray(big), os.path.join(out, key + ".pseudo.nii.gz"), True)
+            interior = vol > 0.02
+            if out_row["role"] != "pos" and interior.sum() > 5000:
+                frac = float(((1.0 - p)[interior] >= args.tau_neg).mean())
+                if frac >= args.neg_frac:
+                    out_row.update({"role": "neg", "neg_conf": round(frac, 4)})
+            buf.append(out_row)
+        except Exception as e:  # noqa
+            print(f"[pseudo:rank{rank}] {key} ERR {e!r}", flush=True)
+        if i % 20 == 0:
+            with open(part, "a") as f:
+                for b in buf:
+                    f.write(json.dumps(b) + "\n")
+            buf = []
+            print(f"[pseudo:rank{rank}] {i}/{len(shard)}", flush=True)
+    if buf:
+        with open(part, "a") as f:
+            for b in buf:
+                f.write(json.dumps(b) + "\n")
+        buf = []
+    if world > 1:
+        dist.barrier()
+    if rank != 0:
+        dist.destroy_process_group() if world > 1 else None
+        return
+
+    merged = []
+    for k in range(world):
+        p = os.path.join(out, f"part{k}.jsonl")
+        if os.path.exists(p):
+            merged.extend(load_jsonl(p))
+
+    # longitudinal consistency filter across accepted positive timepoints
+    by_subj = {}
+    for r in merged:
+        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 ""))
+        accepted = list(range(len(tps)))
+        while True:
+            changed = False
+            for ai in list(accepted):
+                neigh = [ai - 1, ai + 1]
+                neigh = [b for b in neigh if b in accepted]
+                if not neigh:
+                    continue
+                ds = []
+                arrA = load_pseudo(out, tps[ai]["key"]).astype(np.uint8) * 255
+                for bi in neigh:
+                    arrB = load_pseudo(out, tps[bi]["key"]).astype(np.uint8) * 255
+                    ds.append(dice_a_on_b(tps[ai]["key"], arrA, tps[bi]["key"], arrB))
+                if max(ds) < args.cons_dice:
+                    tps[ai]["role"] = "rejected"
+                    accepted.remove(ai)
+                    rejected += 1
+                    changed = True
+                    break
+            if not changed:
+                break
+    for r in merged:
+        if r["role"] == "pos":
+            r["plabel"] = os.path.join(out, r["key"] + ".pseudo.nii.gz")
+    save_jsonl(merged, os.path.join(out, "rows.jsonl"))
+    posv = [r["vol_mm3"] for r in merged if r["role"] == "pos"]
+    summ = {
+        "n_pool": len(merged),
+        "n_pos": sum(1 for r in merged if r["role"] == "pos"),
+        "n_neg": sum(1 for r in merged if r["role"] == "neg"),
+        "n_rejected_consistency": rejected,
+        "n_other_rej": sum(1 for r in merged if r["role"] == "rej"),
+        "pos_vol_mm3": {"med": float(np.median(posv)) if posv else 0,
+                        "p5": float(np.percentile(posv, 5)) if posv else 0,
+                        "p95": float(np.percentile(posv, 95)) if posv else 0},
+        "tau_pos": args.tau_pos, "vol_range": [vol_lo, vol_hi],
+    }
+    with open(os.path.join(out, "summary.json"), "w") as f:
+        json.dump(summ, f, indent=1)
+    print(f"[pseudo] round {outname}: {json.dumps(summ)}", flush=True)
+    if world > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

scripts/07_train.py

@@ -0,0 +1,37 @@
+"""DDP training entrypoint (run via torchrun).
+
+torchrun --standalone --nproc_per_node=3 scripts/07_train.py \
+    --rows data/manifests/split_train.jsonl --val data/manifests/split_val.jsonl \
+    --epochs 40 --lr 3e-4 --ckpt-dir runs/round0
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import argparse
+from src.common import load_jsonl
+from src import training
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--rows", required=True, help="comma-separated jsonl (labeled + pseudo rows)")
+    ap.add_argument("--val", required=True)
+    ap.add_argument("--epochs", type=int, default=40)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--batch", type=int, default=3)
+    ap.add_argument("--patch", type=int, default=96)
+    ap.add_argument("--workers", type=int, default=4)
+    ap.add_argument("--ckpt-dir", required=True)
+    ap.add_argument("--resume", default=None, help="checkpoint to warm-restart weights from")
+    ap.add_argument("--val-every", type=int, default=2)
+    ap.add_argument("--val-limit", type=int, default=60)
+    args = ap.parse_args()
+    rows = []
+    for f in args.rows.split(","):
+        rows.extend(load_jsonl(f))
+    val_rows = load_jsonl(args.val)
+    training.train_ddp(rows, val_rows, args)
+
+
+if __name__ == "__main__":
+    main()

scripts/08_eval.py

@@ -0,0 +1,38 @@
+"""Full holdout evaluation (single GPU).
+
+python scripts/08_eval.py --rows data/manifests/split_test.jsonl --ckpt runs/round0/best.pt
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import argparse
+import json
+import torch
+from src.common import ROOT, load_jsonl, save_jsonl
+from src import training
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--rows", required=True)
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--patch", type=int, default=96)
+    ap.add_argument("--out", default=None)
+    args = ap.parse_args()
+    device = "cuda:0"
+    model, sd = training.load_model(args.ckpt, device)
+    rows = load_jsonl(args.rows)
+    d, per = training.evaluate(model, rows, device, args.patch, tta=True)
+    res = {"ckpt": args.ckpt, "n": len(per), "dice": d, "epoch": sd.get("epoch"),
+           "per_row": per}
+    print(json.dumps({k: res[k] for k in ("ckpt", "n", "dice", "epoch")}, indent=1))
+    out = args.out or os.path.join(ROOT, "results", os.path.basename(os.path.dirname(args.ckpt)) + "_test.json")
+    os.makedirs(os.path.dirname(out), exist_ok=True)
+    with open(out, "w") as f:
+        json.dump(res, f, indent=1)
+    save_jsonl(per, out.replace(".json", "_per_row.jsonl"))
+    print("saved", out)
+
+
+if __name__ == "__main__":
+    main()

scripts/09_run_iterative.py

@@ -0,0 +1,133 @@
+"""Orchestrates the iterative pseudo-labeling study.
+
+Round 0: baseline supervised training on labeled T1c (patient-level split).
+Round k: pseudo-label the unlabeled pool with round k-1 model (confidence +
+         volume plausibility + longitudinal consistency), then fine-tune with
+         labeled + pseudo data (warm restart, lower LR). After each round the
+         model is evaluated on the held-out patient-level test split.
+
+Usage: python scripts/09_run_iterative.py [--rounds 4] [--gpus 3]
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import argparse
+import json
+import shutil
+import subprocess
+from src.common import ROOT, load_jsonl, save_jsonl
+
+
+def run(cmd, log, retries=2):
+    for attempt in range(retries + 1):
+        try:
+            with open(log, "a") as f:
+                f.write(f"$ (attempt {attempt + 1}) " + cmd + "\n")
+                print(f"$ (attempt {attempt + 1}) " + cmd, flush=True)
+                subprocess.run(cmd, shell=True, cwd=ROOT, stdout=f, stderr=subprocess.STDOUT, check=True)
+            return
+        except subprocess.CalledProcessError:
+            if attempt == retries:
+                raise
+            print(f"[orch] command failed, retrying in 60s: {cmd}", flush=True)
+            import time
+            time.sleep(60)
+
+
+def torchrun(nproc, script, extra):
+    tr = shutil.which("torchrun") or (sys.executable + " -m torch.distributed.run")
+    return f"{tr} --standalone --nproc_per_node {nproc} {script} {extra}"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--rounds", type=int, default=4)
+    ap.add_argument("--gpus", type=int, default=3)
+    ap.add_argument("--base-epochs", type=int, default=40)
+    ap.add_argument("--round-epochs", type=int, default=12)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--round-lr", type=float, default=8e-5)
+    ap.add_argument("--pseudo-weight", type=float, default=0.3)
+    ap.add_argument("--batch", type=int, default=3)
+    args = ap.parse_args()
+
+    man = os.path.join(ROOT, "data/manifests")
+    train_f = os.path.join(man, "split_train.jsonl")
+    val_f = os.path.join(man, "split_val.jsonl")
+    test_f = os.path.join(man, "split_test.jsonl")
+    pool_f = os.path.join(man, "unlabeled_pool.jsonl")
+    results_dir = f"{ROOT}/results"
+    os.makedirs(results_dir, exist_ok=True)
+    logdir = f"{ROOT}/logs"
+    os.makedirs(logdir, exist_ok=True)
+    main_log = os.path.join(logdir, "iterative.log")
+
+    # ---- round 0: baseline ----
+    run(torchrun(args.gpus, "scripts/07_train.py",
+                 f"--rows {train_f} --val {val_f} --epochs {args.base_epochs} --lr {args.lr} "
+                 f"--batch {args.batch} --ckpt-dir runs/round0"), f"{main_log}")
+    run(f"python scripts/08_eval.py --rows {test_f} --ckpt runs/round0/best.pt", f"{main_log}")
+
+    for k in range(1, args.rounds + 1):
+        pdir = os.path.join(ROOT, "data/pseudo/round" + str(k))
+        run(torchrun(args.gpus, "scripts/06_pseudo_label.py",
+                     f"--ckpt runs/round{k-1}/best.pt --unlabeled {pool_f} --out {pdir}"), f"{main_log}")
+        # build round-k training manifest: labeled + accepted pseudo rows (weighted)
+        pr = load_jsonl(os.path.join(pdir, "rows.jsonl"))
+        used = []
+        for r in pr:
+            if r["role"] == "pos":
+                r["w"] = args.pseudo_weight
+                used.append(r)
+            elif r["role"] == "neg":
+                r["w"] = args.pseudo_weight
+                r["is_neg"] = True
+                r["plabel"] = None
+                used.append(r)
+        trf = os.path.join(pdir, "train_rows.jsonl")
+        save_jsonl(used, trf)
+        run(torchrun(args.gpus, "scripts/07_train.py",
+                     f"--rows {train_f},{trf} --val {val_f} --epochs {args.round_epochs} "
+                     f"--lr {args.round_lr} --batch {args.batch} --resume runs/round{k-1}/best.pt "
+                     f"--ckpt-dir runs/round{k} --val-every 1"), f"{main_log}")
+        run(f"python scripts/08_eval.py --rows {test_f} --ckpt runs/round{k}/best.pt", f"{main_log}")
+
+    # ---- report ----
+    table = []
+    for k in list(range(args.rounds + 1)):
+        f = os.path.join(results_dir, f"round{k}_test.json")
+        if os.path.exists(f):
+            r = json.load(open(f))
+            table.append({"round": k, "test_dice": round(r["dice"], 4), "n_test": r["n"],
+                          "ckpt_epoch": r.get("epoch")})
+            pf = os.path.join(ROOT, f"data/pseudo/round{k}/summary.json") if k > 0 else None
+            if pf and os.path.exists(pf):
+                s = json.load(open(pf))
+                table[-1].update({"n_pos": s["n_pos"], "n_neg": s["n_neg"],
+                                  "n_rej_cons": s["n_rejected_consistency"],
+                                  "pos_vol_med_mm3": s["pos_vol_mm3"]["med"]})
+        sf = os.path.join(ROOT, f"runs/round{k}/final.pt")
+        table[-1]["ckpt"] = sf if os.path.exists(sf) else ""
+    save_jsonl(table, os.path.join(results_dir, "iterative_table.jsonl"))
+    print(json.dumps(table, indent=1))
+    try:
+        import matplotlib
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+        rs = [t["round"] for t in table]
+        ds = [t["test_dice"] for t in table]
+        plt.figure(figsize=(6, 4))
+        plt.plot(rs, ds, "o-")
+        plt.xlabel("pseudo-labeling round")
+        plt.ylabel("holdout tumor Dice")
+        for x, y in zip(rs, ds):
+            plt.annotate(f"{y:.3f}", (x, y), textcoords="offset points", xytext=(0, 8), fontsize=8)
+        plt.grid(alpha=0.3)
+        plt.tight_layout()
+        plt.savefig(os.path.join(results_dir, "iterative_dice.png"), dpi=150)
+    except Exception as e:  # noqa
+        print("plot failed:", repr(e))
+
+
+if __name__ == "__main__":
+    main()

scripts/preprocess.py

@@ -0,0 +1,114 @@
+"""Preprocess source volumes (T1c + optional label) into a uniform 1mm cropped/normalized form.
+
+Usage:
+  python scripts/preprocess.py --manifest data/manifests/train_labeled.jsonl --workers 16
+Each row: {key, subject, date, img, label?, source}
+Outputs:  data/proc/<key>.nii.gz, data/proc/<key>_label.nii.gz, data/procmeta/<key>.json
+"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import json
+import argparse
+import numpy as np
+import SimpleITK as sitk
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from src.common import ROOT, d, read_nii_img, head_mask_from_image, crop_box, normalize_volume, write_arr, save_jsonl
+
+SPACING = (1.0, 1.0, 1.0)
+
+
+def resample_volume(itk, spacing, interp):
+    size = [max(1, int(round(s * sp / nsp))) for s, sp, nsp in zip(itk.GetSize(), itk.GetSpacing(), spacing)]
+    ptype = sitk.sitkFloat32 if interp != sitk.sitkNearestNeighbor else itk.GetPixelID()
+    return sitk.Resample(itk, size, sitk.Transform(), interp,
+                         itk.GetOrigin(),
+                         tuple(float(s) for s in spacing),
+                         itk.GetDirection(),
+                         0.0, ptype)
+
+
+def worker(args):
+    key, img_path, label_path, out_dir, meta_dir = args
+    try:
+        img_itk = read_nii_img(img_path)
+        img1 = resample_volume(img_itk, SPACING, sitk.sitkLinear)
+        a = sitk.GetArrayFromImage(img1)
+        label1 = None
+        if label_path and os.path.exists(label_path):
+            lbl = read_nii_img(label_path)
+            label1 = sitk.GetArrayFromImage(resample_volume(lbl, SPACING, sitk.sitkNearestNeighbor)).astype(np.uint8)
+            label1 = (label1 > 0).astype(np.uint8)
+        # ensure label1 grid == a grid
+        if label1 is not None and label1.shape != a.shape:
+            ref = sitk.GetImageFromArray(np.zeros(a.shape, dtype=np.uint8))
+            ref.CopyInformation(img1)
+            lbl2 = resample_volume(read_nii_img(label_path), SPACING, sitk.sitkNearestNeighbor)
+            label1 = sitk.GetArrayFromImage(sitk.Resample(lbl2, ref, sitk.Transform(), sitk.sitkNearestNeighbor)).astype(np.uint8)
+            label1 = (label1 > 0).astype(np.uint8)
+            if label1.shape != a.shape:
+                raise RuntimeError(f"label/image grid mismatch {label1.shape} vs {a.shape}")
+        if a.max() <= a.min() + 1e-6:
+            raise RuntimeError("empty volume")
+        box = crop_box(a, label=label1, margin=12, cap=216, spacing=1.0)
+        (s0, s1), (s2, s3), (s4, s5) = box
+        ac_arr = a[s0:s1, s2:s3, s4:s5]
+        lab_arr = label1[s0:s1, s2:s3, s4:s5] if label1 is not None else None
+        m = head_mask_from_image(ac_arr)
+        if m.sum() < 5000:
+            raise RuntimeError("head mask too small")
+        norm, (lo, hi) = normalize_volume(ac_arr, m)
+        norm = norm.astype(np.float32)
+        out = os.path.join(out_dir, key + ".nii.gz")
+        write_arr(norm, out)
+        meta = {"key": key, "img_spacing_in": list(img_itk.GetSpacing()), "img_size_in": list(img_itk.GetSize()),
+                "origin": list(img_itk.GetOrigin()), "direction": list(img_itk.GetDirection()),
+                "crop_vox": [int(s0), int(s2), int(s4)],
+                "norm_lo": float(lo), "norm_hi": float(hi),
+                "shape": list(norm.shape)}
+        with open(os.path.join(meta_dir, key + ".json"), "w") as f:
+            json.dump(meta, f)
+        if lab_arr is not None:
+            outl = os.path.join(out_dir, key + "_label.nii.gz")
+            write_arr(lab_arr.astype(np.uint8), outl)
+            meta["n_tumor_vox"] = int(lab_arr.sum())
+            with open(os.path.join(meta_dir, key + ".json"), "w") as f:
+                json.dump(meta, f)
+        return key, None
+    except Exception as e:  # noqa
+        return key, repr(e)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--manifest", required=True)
+    ap.add_argument("--workers", type=int, default=16)
+    ap.add_argument("--subset", type=int, default=0, help="only first N rows (0=all)")
+    args = ap.parse_args()
+    rows = [json.loads(l) for l in open(args.manifest) if l.strip()]
+    if args.subset:
+        rows = rows[: args.subset]
+    out_dir = d("data/proc")
+    meta_dir = d("data/procmeta")
+    todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows
+            if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz"))]
+    print(f"rows={len(rows)} todo={len(todo)}")
+    errs = []
+    with ProcessPoolExecutor(max_workers=args.workers) as ex:
+        futs = {ex.submit(worker, t): t[0] for t in todo}
+        done = 0
+        for fu in as_completed(futs):
+            k, err = fu.result()
+            done += 1
+            if err:
+                errs.append((k, err))
+                print("ERR", k, err, flush=True)
+            if done % 50 == 0:
+                print(f"  {done}/{len(todo)} done, {len(errs)} errs", flush=True)
+    print(f"finished. errors: {len(errs)}")
+    for k, e in errs[:20]:
+        print("  ", k, e)
+
+
+if __name__ == "__main__":
+    main()

src/common.py

@@ -0,0 +1,191 @@
+import os
+import json
+import glob
+import numpy as np
+import SimpleITK as sitk
+from scipy import ndimage
+
+ROOT = os.environ.get("LONGITUDINAL_ROOT", "/mnt/b4/xfr/git26/longitudinal")
+
+
+def d(name):
+    p = os.path.join(ROOT, name)
+    os.makedirs(p, exist_ok=True)
+    return p
+
+
+def load_npz_or_nii(path):
+    if path.endswith(".npz"):
+        return np.load(path, allow_pickle=False)
+    return sitk.GetArrayFromImage(sitk.ReadImage(path))
+
+
+def read_nii(path):
+    return sitk.ReadImage(path)
+
+
+def write_nii(itk_img, path):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    sitk.WriteImage(itk_img, path, True)
+
+
+def npy_to_itk(arr, origin=(0, 0, 0), spacing=(1, 1, 1), direction=(1, 0, 0, 0, 1, 0, 0, 0, 1)):
+    img = sitk.GetImageFromArray(arr.astype(np.float32 if arr.dtype != np.uint8 else np.uint8))
+    img.SetOrigin(tuple(origin))
+    img.SetSpacing(tuple(spacing))
+    img.SetDirection(tuple(direction))
+    return img
+
+
+def resample_to_spacing(itk_img, spacing, interpol=sitk.sitkLinear):
+    arr = sitk.GetArrayFromImage(itk_img)
+    out = np.zeros(arr.shape, dtype=np.float32)
+    if max(spacing) > 0:
+        out = sitk.GetArrayFromImage(
+            sitk.Resample(
+                itk_img,
+                itk_img,
+                sitk.Transform(),
+                interpol,
+                0.0,
+                tuple(float(s) for s in spacing),
+            )
+        )
+    return out
+
+
+def resample_array(arr, spacing_from, spacing_to, interp="linear"):
+    h, w, d = arr.shape
+    ref = sitk.GetImageFromArray(np.zeros((1, 1, 1), dtype=np.float32))
+    src = sitk.GetImageFromArray(arr.astype(np.float32))
+    src.SetSpacing(tuple(spacing_from))
+    ref2 = sitk.GetImageFromArray(arr.astype(np.float32))
+    ref2.SetSpacing(tuple(spacing_from))
+    new = sitk.Resample(
+        ref2,
+        ref2,
+        sitk.Transform(),
+        sitk.sitkLinear if interp == "linear" else sitk.sitkNearestNeighbor,
+        0.0,
+        tuple(spacing_to),
+    )
+    return sitk.GetArrayFromImage(new)
+
+
+def largest_cc(mask):
+    lab, n = ndimage.label(mask)
+    if n == 0:
+        return mask
+    sizes = ndimage.sum(mask, lab, range(1, n + 1))
+    return lab == (int(np.argmax(sizes)) + 1)
+
+
+def head_mask_from_image(arr):
+    """Foreground (head + tumor) mask from a T1c volume. Robust to 8-bit or raw MR scaling."""
+    a = arr.ravel()
+    a = a[a > 0]
+    if a.size == 0:
+        a = arr.ravel()
+    thr = np.percentile(a if a.size else arr.ravel(), 2)
+    m = arr > max(thr, np.finfo(arr.dtype).eps)
+    m = largest_cc(m)
+    return m
+
+
+def _axis_range(mask, i):
+    proj = np.any(mask, axis=tuple(j for j in range(mask.ndim) if j != i))
+    idx = np.where(proj)[0]
+    return int(idx.min()), int(idx.max()) + 1
+
+
+def crop_box(arr, label=None, margin=10, cap=216, spacing=1.0):
+    """Compute (start, end) per axis for a head crop. Includes label with margin when given."""
+    m = head_mask_from_image(arr)
+    m = ndimage.binary_dilation(m, iterations=max(1, int(margin / spacing)))
+    sl = [_axis_range(m, i) for i in range(m.ndim)]
+    sl = [slice(lo, hi) for lo, hi in sl]
+    if label is not None and (label > 0).any():
+        l = ndimage.binary_dilation(label > 0, iterations=max(1, int(margin / spacing)))
+        for i in range(arr.ndim):
+            lo, hi = _axis_range(l, i)
+            sl[i] = slice(min(sl[i].start, lo), max(sl[i].stop, hi))
+    start = [s.start for s in sl]
+    end = [s.stop for s in sl]
+    # cap each dim to `cap`, centered on head centroid (or label centroid if given)
+    cent = ndimage.center_of_mass(m)
+    if label is not None and (label > 0).any():
+        cent = ndimage.center_of_mass((label > 0) & (m | label))
+    for i in range(arr.ndim):
+        size = end[i] - start[i]
+        if size > cap:
+            c = int(round(cent[i]))
+            s0 = max(0, min(c - cap // 2, arr.shape[i] - cap))
+            start[i], end[i] = s0, s0 + cap
+    return [(s, e) for s, e in zip(start, end)]
+
+
+def normalize_volume(arr, mask):
+    a = arr[mask]
+    lo, hi = np.percentile(a, 2), np.percentile(a, 98)
+    if hi - lo < 1e-6:
+        hi = lo + 1.0
+    out = (arr - lo) / (hi - lo)
+    return np.clip(out, 0.0, 1.0).astype(np.float32), (lo, hi)
+
+
+def slices(a, b):
+    return slice(a, b)
+
+
+def load_jsonl(path):
+    with open(path) as f:
+        return [json.loads(l) for l in f if l.strip()]
+
+
+def save_jsonl(rows, path):
+    p = str(path)
+    os.makedirs(os.path.dirname(p), exist_ok=True)
+    with open(p, "w") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+
+
+def read_nii_arr(path):
+    return sitk.GetArrayFromImage(sitk.ReadImage(path))
+
+
+def read_nii_img(path):
+    return sitk.ReadImage(path)
+
+
+def write_arr(arr, path, itk_img=None):
+    """Write array to nii.gz, copying geometry from itk_img when provided (else identity 1mm)."""
+    if itk_img is None:
+        itk_img = sitk.GetImageFromArray(arr.astype(np.float32))
+        itk_img.SetSpacing((1, 1, 1))
+    else:
+        if arr.ndim == 3 and arr.shape == tuple(itk_img.GetSize()[::-1]):
+            img = sitk.GetImageFromArray(arr.astype(np.float32))
+            img.CopyInformation(itk_img)
+            itk_img = img
+        else:
+            raise ValueError(f"shape mismatch {arr.shape} vs {itk_img.GetSize()}")
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    sitk.WriteImage(itk_img, path, True)
+
+
+def dice(a, b, eps=1e-7):
+    a = a > 0
+    b = b > 0
+    if a.sum() == 0 and b.sum() == 0:
+        return 1.0
+    return float(2 * (a & b).sum() / (a.sum() + b.sum() + eps))
+
+
+def bbox3(mask):
+    sl = ndimage.find_objects(ndimage.label(mask)[0])
+    idx = np.argwhere(mask)
+    if idx.size == 0:
+        return None
+    lo, hi = idx.min(0), idx.max(0)
+    return lo, hi

src/dataset.py

@@ -0,0 +1,102 @@
+import os
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+
+class PatchDataset(Dataset):
+    """Random 3D patches from preprocessed nifti volumes.
+
+    row: {key, pimg, plabel?, w (loss weight, default 1.0)}
+    yields (img(1,H,W,D), label(H,W,D) or None, weight)
+    """
+
+    def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0):
+        self.rows = list(rows)
+        self.patch = tuple(patch)
+        self.train = train
+        self.rng = np.random.RandomState(seed)
+
+    def __len__(self):
+        return len(self.rows)
+
+    @staticmethod
+    def _read(path):
+        import SimpleITK as sitk
+        return sitk.GetArrayFromImage(sitk.ReadImage(path))
+
+    def __getitem__(self, idx):
+        row = self.rows[idx]
+        w = float(row.get("w", 1.0))
+        img = self._read(row["pimg"]).astype(np.float32)
+        # 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]]
+            label = np.clip(np.rint(label).astype(np.int64), 0, 1)
+
+        if self.train:
+            s0 = [self.rng.randint(0, img.shape[i] - self.patch[i] + 1) if img.shape[i] >= self.patch[i] else 0
+                  for i in range(3)]
+            img = img[s0[0]:s0[0] + self.patch[0], s0[1]:s0[1] + self.patch[1], s0[2]:s0[2] + self.patch[2]]
+            if label is not None:
+                label = label[s0[0]:s0[0] + self.patch[0], s0[1]:s0[1] + self.patch[1], s0[2]:s0[2] + self.patch[2]]
+            img, label = self._augment(img, label)
+        else:
+            if any(img.shape[i] < self.patch[i] for i in range(3)):
+                img = np.pad(img, [(0, max(self.patch[i] - img.shape[i], 0)) for i in range(3)])
+                if label is not None:
+                    label = np.pad(label, [(0, max(self.patch[i] - label.shape[i], 0)) for i in range(3)])
+
+        t = torch.from_numpy(img).unsqueeze(0)
+        lab = torch.from_numpy(label) if label is not None else None
+        return t, lab, torch.tensor(w, dtype=torch.float32)
+
+    @staticmethod
+    def _augment(img, label):
+        if np.random.rand() > 0.5:
+            img = img[:, :, ::-1].copy()
+            if label is not None: label = label[:, :, ::-1].copy()
+        if np.random.rand() > 0.5:
+            img = img[:, ::-1, :].copy()
+            if label is not None: label = label[:, ::-1, :].copy()
+        if np.random.rand() > 0.5:
+            img = img[::-1, :, :].copy()
+            if label is not None: label = label[::-1, :, :].copy()
+        r = np.random.rand()
+        if r < 0.1:
+            k = int(np.random.rand() * 4)
+            img = np.rot90(img, k, (1, 2)).copy()
+            if label is not None: label = np.rot90(label, k, (1, 2)).copy()
+        elif r < 0.4:
+            img = img * float(np.random.uniform(0.9, 1.1))
+        if np.random.rand() > 0.6:
+            img = img + np.random.normal(0, 0.01, img.shape).astype(np.float32)
+        return img, label
+
+
+def _collate(items):
+    imgs = torch.stack([it[0] for it in items])
+    labs = [it[1] for it in items]
+    if all(l is not None for l in labs):
+        labs = torch.stack(labs)
+    else:
+        labs = None
+    wts = torch.stack([it[2] for it in items])
+    return imgs, labs, wts
+
+
+def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persistent=True):
+    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))
+    return ds, dl

src/losses.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SoftDiceLoss(nn.Module):
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, logits, target):
+        p = F.softmax(logits, dim=1)
+        n = target.size(1)
+        losses = []
+        for c in range(n):
+            tc = target[:, c]
+            pc = p[:, c]
+            dim = tuple(range(1, pc.dim()))
+            inter = (pc * tc).sum(dim)
+            denom = pc.sum(dim) + tc.sum(dim)
+            losses.append(1 - (2 * inter + self.eps) / (denom + self.eps))
+        return torch.stack(losses, dim=0).mean(dim=0)
+
+
+def per_sample_loss(logits, target, dice_weight=0.5):
+    """target: (B,C,...) one-hot -> per-sample loss (B,)"""
+    bce = F.cross_entropy(logits, target.argmax(dim=1), reduction="none")
+    bce = bce.mean(dim=tuple(range(1, bce.dim())))
+    dcl = SoftDiceLoss()(logits, target)
+    return (1 - dice_weight) * bce + dice_weight * dcl
+
+
+def one_hot(label, num_classes=2):
+    one = torch.zeros(label.size(0), num_classes, *label.shape[1:], device=label.device, dtype=torch.long)
+    return one.scatter_(1, label.unsqueeze(1), 1)

src/training.py

@@ -0,0 +1,181 @@
+"""Training / evaluation / inference for the T1c brain-tumor 3D U-Net."""
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from src.unet3d import Unet3D
+from src.dataset import make_dataloader
+from src.losses import per_sample_loss, one_hot
+
+
+def build_model(base=16, out_ch=2, device="cuda"):
+    return Unet3D(in_ch=1, out_ch=out_ch, base=base, depth=4).to(device)
+
+
+@torch.no_grad()
+def sliding_window_probs(model, vol, device, win=96, step=48, tta=True):
+    """vol: (H,W,D) float32 -> tumor probability (H,W,D) numpy float32."""
+    model.eval()
+    h, w, d = vol.shape
+    vol = np.pad(vol, [(0, max(h, win) - h), (0, max(w, win) - w), (0, max(d, win) - d)])
+    h, w, d = vol.shape
+
+    def coords(n):
+        if n <= win:
+            return [0]
+        pos = list(range(0, n - win + 1, step))
+        if n - win > pos[-1]:
+            pos.append(n - win)
+        return pos
+    cs = (coords(h), coords(w), coords(d))
+    acc = np.zeros((h, w, d), dtype=np.float32)
+    cnt = np.zeros((h, w, d), dtype=np.float32)
+    flips = [lambda x: x, lambda x: x[:, ::-1].copy(), lambda x: x[:, :, ::-1].copy(),
+             lambda x: x[::-1, :, :].copy()]
+    if tta:
+        flips.append(lambda x: x[::-1, ::-1].copy())
+        flips.append(lambda x: x[::-1, :, ::-1].copy())
+        flips.append(lambda x: x[:, ::-1, ::-1].copy())
+        flips.append(lambda x: x[::-1, ::-1, ::-1].copy())
+    for cz in cs[2]:
+        for cy in cs[1]:
+            for cx in cs[0]:
+                sl = (slice(cx, cx + win), slice(cy, cy + win), slice(cz, cz + win))
+                for tf in flips:
+                    patch = tf(vol)[sl]
+                    t = torch.from_numpy(patch).float().unsqueeze(0).unsqueeze(0).to(device)
+                    with torch.autocast("cuda", dtype=torch.bfloat16):
+                        logits = model(t)
+                    prob = F.softmax(logits.float(), dim=1)[0, 1].cpu().numpy()
+                    acc[sl] += tf(prob)
+                    cnt[sl] += 1
+    return acc / np.maximum(cnt, 1e-6)
+
+
+def dice_np(probs, lab, thr=0.5):
+    pred = probs >= thr
+    l = lab > 0
+    if l.sum() == 0:
+        return 1.0 if pred.sum() == 0 else 0.0
+    return float(2 * (pred & l).sum() / (pred.sum() + l.sum()))
+
+
+@torch.no_grad()
+def evaluate(model, rows, device, patch=96, tta=True):
+    model.eval()
+    scores = []
+    for r in rows:
+        try:
+            vol = _read(r["pimg"]).astype(np.float32)
+            lab = _read(r["plabel"]).astype(np.uint8)
+            if vol.shape != lab.shape:
+                lab = np.pad(lab, [(0, max(vol.shape[i] - lab.shape[i], 0)) for i in range(3)])
+                lab = lab[: vol.shape[0], : vol.shape[1], : vol.shape[2]]
+            probs = sliding_window_probs(model, vol, device, patch, max(patch // 2, 16), tta)
+            scores.append((r["key"], dice_np(probs, lab)))
+        except Exception as e:  # noqa
+            print("  eval err", r.get("key"), repr(e))
+    if not scores:
+        return 0.0, []
+    return float(np.mean([s for _, s in scores])), scores
+
+
+def _read(path):
+    import SimpleITK as sitk
+    return sitk.GetArrayFromImage(sitk.ReadImage(path))
+
+
+def train_ddp(rows, val_rows, args):
+    import torch.distributed as dist
+    from torch.nn.parallel import DistributedDataParallel as DDP
+    from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
+
+    rank = int(os.environ.get("RANK", 0))
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world = int(os.environ.get("WORLD_SIZE", 1))
+    torch.manual_seed(0)
+    np.random.seed(0)
+
+    patch = (args.patch, args.patch, args.patch)
+    _, dl = make_dataloader(rows, patch, args.batch, True, num_workers=args.workers, seed=0 + rank,
+                            persistent=False)
+    steps_per_epoch = max(len(dl), 1)
+
+    if world > 1:
+        from datetime import timedelta
+        dist.init_process_group("nccl", timeout=timedelta(minutes=30))
+    torch.cuda.set_device(local_rank)
+    device = f"cuda:{local_rank}"
+
+    model = build_model(device=device)
+    if args.resume:
+        sd = torch.load(args.resume, map_location=device, weights_only=True)
+        model.load_state_dict(sd["model"])
+        if rank == 0:
+            print(f"[train:rank0] resumed weights from {args.resume}")
+    ddp = DDP(model, device_ids=[local_rank]) if world > 1 else model
+    total_steps = steps_per_epoch * args.epochs
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
+    warmup = min(300, max(10, total_steps // 10))
+    base_sched = CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1), eta_min=args.lr * 0.05)
+    sched = SequentialLR(opt, [LinearLR(opt, start_factor=0.1, total_iters=warmup), base_sched],
+                         milestones=[warmup])
+
+    best, best_epoch = -1.0, -1
+    ckpt_dir = args.ckpt_dir
+    if rank == 0:
+        os.makedirs(ckpt_dir, exist_ok=True)
+        print(f"[train] rows={len(rows)} val={len(val_rows)} epochs={args.epochs} steps/epoch={steps_per_epoch}", flush=True)
+
+    for epoch in range(args.epochs):
+        if hasattr(dl.sampler, "set_epoch"):
+            dl.sampler.set_epoch(epoch)
+        model.train()
+        run_loss, run_n = 0.0, 0
+        for img, lab, wts in dl:
+            img = img.to(device, non_blocking=True)
+            lab = lab.to(device, non_blocking=True)
+            wts = wts.to(device, non_blocking=True)
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                logits = ddp(img)
+            losses = per_sample_loss(logits.float(), one_hot(lab, 2))
+            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)
+            opt.step()
+            sched.step()
+            run_loss += float(loss.detach())
+            run_n += 1
+        if rank == 0:
+            print(f"  epoch {epoch+1}/{args.epochs} loss={run_loss/max(run_n,1):.4f} "
+                  f"lr={opt.param_groups[0]['lr']:.2e}", flush=True)
+        if val_rows and rank == 0 and (epoch + 1) % args.val_every == 0:
+            d, _ = evaluate(model, val_rows[: args.val_limit], device, patch[0], tta=True)
+            print(f"  [val] epoch {epoch+1} dice={d:.4f}", flush=True)
+            if d > best:
+                best = d
+                best_epoch = epoch + 1
+                torch.save({"model": model.state_dict(), "epoch": epoch + 1, "val_dice": best},
+                           os.path.join(ckpt_dir, "best.pt"))
+        if world > 1:
+            dist.barrier()
+    if rank == 0:
+        torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": best},
+                   os.path.join(ckpt_dir, "final.pt"))
+        print(f"[train] done best_val_dice={best:.4f}@{best_epoch}", flush=True)
+    if world > 1:
+        dist.destroy_process_group()
+    return best
+
+
+def load_model(args_ckpt, device, base=16):
+    sd = torch.load(args_ckpt, map_location=device, weights_only=True)
+    model = build_model(base=base, device=device)
+    model.load_state_dict(sd["model"])
+    model.eval()
+    return model, sd

src/unet3d.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv3d(in_ch, out_ch, 3, padding=1),
+            nn.BatchNorm3d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(out_ch, out_ch, 3, padding=1),
+            nn.BatchNorm3d(out_ch),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Unet3D(nn.Module):
+    def __init__(self, in_ch=1, out_ch=2, base=16, depth=4):
+        super().__init__()
+        chs = [base * 2 ** i for i in range(depth + 1)]
+        self.encs = nn.ModuleList([DoubleConv(in_ch if i == 0 else chs[i - 1], chs[i]) for i in range(depth)])
+        self.pools = nn.ModuleList([nn.MaxPool3d(2) for _ in range(depth - 1)])
+        self.decs = nn.ModuleList([
+            DoubleConv(chs[depth - 1 - k] + chs[depth - 2 - k], chs[depth - 2 - k])
+            for k in range(depth - 1)
+        ])
+        self.head = nn.Conv3d(chs[0], out_ch, 1)
+        self.depth = depth
+        self.chs = chs
+
+    def forward(self, x):
+        skips = []
+        h = x
+        for i, enc in enumerate(self.encs):
+            h = enc(h)
+            skips.append(h)
+            if i < len(self.pools):
+                h = self.pools[i](h)
+        for k, dec in enumerate(self.decs):
+            h = F.interpolate(h, scale_factor=2, mode="nearest")
+            h = torch.cat([h, skips[self.depth - 2 - k]], dim=1)
+            h = dec(h)
+        return self.head(h)