Furen Xiao
2453fb48ba
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.
168 lines
No EOL
7.3 KiB
Python
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() |