longitudinal/scripts/05_build_splits.py
Furen Xiao 2453fb48ba feat(dataset): add row-level exclusion curation for labeled volumes
Introduce src/exclusions.py (rules: total TV > 1e6 mm3, single target
< 100 mm3, non-brain NTUH keys) and scripts/make_tv_slices.py, a
first-stage curation script that classifies every labeled row right
after the source fetch and writes TV-slice QA figures to
data/tv_slices/{inclusion,exclusion}/.

05_build_splits.py re-applies the same rules after split assignment:
kept rows retain their split, excluded subjects still count in the
subject-level shuffle boundaries, and the audit trail is written to
data/manifests/excluded.jsonl. Label stats run in a
ProcessPoolExecutor (32 workers).

Move proc_row into src/common.py and add labeled_source_rows() so the
curation stage and the split builder share one processed-row view.

Document the 143-row exclusion audit and the interim 10-epoch
benchmark results in BENCHMARKING.md; update the README data layout
and script tables.
2026-10-05 16:53:56 +08:00

168 lines
No EOL
7.3 KiB
Python

"""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)
Each row carries both processed fields (pimg/plabel, data/proc — 1mm cropped/
normalized, used by Pipelines A and C) and native fields (img/label, the raw
source NIfTIs — used by Pipeline B, which feeds nnU-Net its own
plan_and_preprocess). Row membership still requires the processed volume to
exist, so the patient-level splits stay identical across pipelines.
Before splitting, every labeled row is screened with the row-level exclusion
rules (src/exclusions.py: total TV > 1e6 mm3, single target < 100 mm3, and
the non-brain ntuh rows in NON_BRAIN_KEYS). Excluded rows are written to
data/manifests/excluded.jsonl with their reason and never enter a split.
"""
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, DATA, d, load_jsonl, save_jsonl, proc_row
from src.exclusions import label_stats, exclusion_reason
def _label_stats_row(r):
"""(key, label stats) for the ProcessPoolExecutor in main()."""
try:
return r["key"], label_stats(r["plabel"])
except Exception:
return r["key"], {"total": 0, "n_cc": 0}
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(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)}")
# row-level curation (src/exclusions.py): label stats + exclusion reason
# per row. Exclusions are applied AFTER the split assignment, so kept rows
# retain the split of the 2026-10-04 rebuild and excluded.jsonl is the
# audit trail (excluded subjects still count in the shuffle boundaries).
from collections import Counter
from concurrent.futures import ProcessPoolExecutor
tv_vols, reason_of = {}, {}
with ProcessPoolExecutor(max_workers=32) as ex:
for k, stats in ex.map(_label_stats_row, labeled_rows, chunksize=16):
tv_vols[k] = stats["total"]
reason_of[k] = exclusion_reason(k, stats)
n_excl = sum(1 for v in reason_of.values() if v)
def _cat(reason):
return (reason.split(" image", 1)[0] + " image"
if "non-brain" in reason else
reason.split(" TV ", 1)[0] + " TV < 100 mm3"
if "single target" in reason else
"total TV > 1e6 mm3")
rc = Counter(_cat(v) for v in reason_of.values() if v)
print(f"excluded rows: {n_excl}")
for reason, n in sorted(rc.items()):
print(f" excluded {n:4d} {reason}")
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, excluded = {"train": [], "val": [], "test": []}, []
for r in labeled_rows:
r["w"] = 1.0
if r["subject"] in test_subj:
sk = "test"
elif r["subject"] in val_subj:
sk = "val"
else:
sk = "train"
if reason_of[r["key"]]:
excluded.append({**r, "split": sk, "exclude_reason": reason_of[r["key"]]})
else:
splits[sk].append(r)
for k, v in splits.items():
save_jsonl(v, os.path.join(DATA, "manifests", f"split_{k}.jsonl"))
print(f" split_{k}: {len(v)} volumes / {len(set(r['subject'] for r in v))} subjects")
save_jsonl(excluded, os.path.join(DATA, "manifests", "excluded.jsonl"))
kept_rows = [r for v in splits.values() for r in v]
labeled_subjects = train_subj | val_subj | test_subj
# unlabeled pool
pool = []
for r in load_jsonl(os.path.join(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(DATA, "manifests/lee_t1c_selected.jsonl")
lee_rows = {}
lee_manifest = []
if os.path.exists(lee_sel_path):
lee_manifest = load_jsonl(lee_sel_path)
raw_path = os.path.join(DATA, "manifests", "lee_t1c_raw.jsonl")
if os.path.exists(raw_path):
# last-resort rescues: thick series the scan filtered out, kept
# when their timepoint otherwise yielded no T1c volume
have = {r["key"] for r in lee_manifest}
for r in load_jsonl(raw_path):
if r.get("thick_dropped") and r["key"] not in have:
lee_manifest.append(r)
have.add(r["key"])
for r in lee_manifest:
nii = os.path.join(d("data/lee_nii"), r["sid"], f"{r['date']}_s{r['ser']}.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,
"img": nii, "label": None, "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(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 of the kept dataset rows (1mm voxels == mm3; totals from
# the curation pass)
vols = []
for r in kept_rows:
vols.append({"key": r["key"], "subject": r["subject"],
"vol_mm3": tv_vols[r["key"]]})
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(DATA, "vols.jsonl"))
with open(os.path.join(DATA, "vols.json"), "w") as f:
json.dump(stats, f, indent=1)
print("tumor volumes (mm3):", stats)
if __name__ == "__main__":
main()