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.
230 lines
No EOL
9.4 KiB
Python
230 lines
No EOL
9.4 KiB
Python
"""First-stage curation + TV-slice QA figures for labeled rows.
|
|
|
|
Runs on the labeled source rows (data/manifests/{ntuh,m6_labeled}.jsonl, 1mm
|
|
processed) — the first place rows exist after fetching from the sources,
|
|
before any split (05). For each labeled volume:
|
|
|
|
1. classify with the exclusion rules in src/exclusions.py (total TV > 1e6
|
|
mm3, single target < 100 mm3, curated non-brain list),
|
|
2. render one PNG with the three orthogonal slices (axial | coronal |
|
|
sagittal) that carry the largest TV (target-volume) area, with the TV
|
|
contour overlaid; when the label has several connected targets
|
|
(26-connectivity), only the largest one is shown,
|
|
3. write the figure to data/tv_slices/inclusion/ (kept) or
|
|
data/tv_slices/exclusion/ (excluded, reason in the title).
|
|
|
|
Slices are picked per body axis (S/P/L) from the procmeta direction, so the
|
|
panel order follows radiological conventions (same mapping as
|
|
scripts/make_screenshots.py); volumes that are not body-aligned fall back to
|
|
the image-axis slices with a note in the title. Filenames are
|
|
<total_vol_mm3:07d>_<key>.png, so each folder lists small TV -> large TV.
|
|
|
|
The split files (05) are the dataset of record; this stage only curates and
|
|
documents. 05 applies the same src/exclusions rules after split assignment,
|
|
so the two always agree on kept/excluded.
|
|
|
|
Usage:
|
|
python scripts/make_tv_slices.py # all source rows, 32 workers
|
|
python scripts/make_tv_slices.py --from-splits # rows/tag from split files
|
|
python scripts/make_tv_slices.py --keys k1 k2 # specific rows
|
|
python scripts/make_tv_slices.py --limit 50 # first N rows
|
|
python scripts/make_tv_slices.py --force # overwrite existing PNGs
|
|
"""
|
|
import os
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
import glob
|
|
import json
|
|
import argparse
|
|
from collections import Counter
|
|
|
|
import numpy as np
|
|
import SimpleITK as sitk
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
from scipy import ndimage
|
|
|
|
from src.common import DATA, d, load_jsonl, proc_row, labeled_source_rows
|
|
from src.exclusions import label_stats, exclusion_reason
|
|
from scripts.make_screenshots import vol_to_body
|
|
|
|
|
|
def out_name(vol_total, key):
|
|
"""Flat-folder filename: zero-padded TOTAL TV volume (mm3, all targets)
|
|
first so a plain listing sorts small -> large, then the row key."""
|
|
return f"{int(vol_total):07d}_{key}.png"
|
|
|
|
|
|
def largest_cc(mask, structure=None):
|
|
"""Largest connected component of a boolean mask (26-conn by default)."""
|
|
if structure is None:
|
|
structure = np.ones((3, 3, 3), dtype=int)
|
|
lab, n = ndimage.label(mask, structure=structure)
|
|
if n == 0:
|
|
return mask, 0
|
|
sizes = ndimage.sum(mask, lab, range(1, n + 1))
|
|
return lab == (int(np.argmax(sizes)) + 1), int(n)
|
|
|
|
|
|
def max_area_slices(cc):
|
|
"""Per-axis index of the slice containing the largest CC area."""
|
|
return tuple(int(np.argmax(cc.sum(axis=tuple(j for j in range(cc.ndim) if j != i))))
|
|
for i in range(cc.ndim))
|
|
|
|
|
|
def render(key, tag, mr, cc, n_cc, vol_total, direction, reason, out, dpi=110):
|
|
"""1x3 figure: axial | coronal | sagittal, max-CC-area slice each, with the
|
|
largest-CC contour in red. Writes to `out`. Returns the path."""
|
|
body, ok, _ = vol_to_body(mr, direction)
|
|
if ok:
|
|
c, _, _ = vol_to_body(cc.astype(np.float32), direction)
|
|
si, pi, li = max_area_slices(c)
|
|
panels = [
|
|
(body[si], c[si], "axial", "S", si),
|
|
(body[:, pi, :][::-1], c[:, pi, :][::-1], "coronal", "P", pi),
|
|
(body[:, :, li][::-1], c[:, :, li][::-1], "sagittal", "L", li),
|
|
]
|
|
else: # not body-aligned: image-axis fallback
|
|
zi, yi, xi = max_area_slices(cc)
|
|
panels = [
|
|
(mr[zi], cc[zi], "axial*", "z", zi),
|
|
(mr[:, yi, :][::-1], cc[:, yi, :][::-1], "coronal*", "y", yi),
|
|
(mr[:, :, xi][::-1], cc[:, :, xi][::-1], "sagittal*", "x", xi),
|
|
]
|
|
|
|
lo, hi = np.percentile(mr[mr > 0], [1, 99.5])
|
|
panel_specs = []
|
|
for im, c2, title, axn, idx in panels:
|
|
nr, nc = im.shape
|
|
a = 1.0 # processed volumes are 1mm isotropic
|
|
panel_specs.append((im, c2, f"{title} {axn}={idx}", a, nr, nc))
|
|
|
|
max_in, gap_in, m_in, top_in, bot_in = 4.4, 0.25, 0.3, 1.7, 0.2
|
|
w_in, h_in = [], []
|
|
for _, _, _, a, nrow, ncol in panel_specs:
|
|
pw, ph = ncol, nrow * a
|
|
m = max(pw, ph)
|
|
w_in.append(max_in * pw / m)
|
|
h_in.append(max_in * ph / m)
|
|
plot_h = max_in
|
|
fig_w = sum(w_in) + 2 * gap_in + 2 * m_in
|
|
fig_h = plot_h + top_in + bot_in
|
|
fig = plt.figure(figsize=(fig_w, fig_h))
|
|
left = m_in / fig_w
|
|
for (im, c2, title, a, _, _), w, h in zip(panel_specs, w_in, h_in):
|
|
axi = fig.add_axes([left, (bot_in + (plot_h - h) / 2) / fig_h, w / fig_w, h / fig_h])
|
|
axi.imshow(im, cmap="gray", vmin=lo, vmax=hi, origin="upper", aspect=a)
|
|
if c2.any():
|
|
axi.contour(c2, levels=[0.5], colors="r", linewidths=1.2)
|
|
axi.set_title(title, color="w", fontsize=11)
|
|
axi.set_xticks([])
|
|
axi.set_yticks([])
|
|
for spine in axi.spines.values():
|
|
spine.set_edgecolor("0.35")
|
|
left += (w + gap_in) / fig_w
|
|
tkey = key.split("_", 1)[1] if "_" in key else key
|
|
align = "body-aligned" if ok else "NOT body-aligned (image axes)"
|
|
status = f"EXCLUDED: {reason}" if reason else "kept"
|
|
fig.suptitle(
|
|
f"{tkey} [{tag}] TV total={int(vol_total)} mm3, "
|
|
f"largest={int(cc.sum())} mm3 ({n_cc} targets) {status} {align}",
|
|
color="w", fontsize=12)
|
|
fig.patch.set_facecolor("k")
|
|
fig.savefig(out, facecolor="k", dpi=dpi)
|
|
plt.close(fig)
|
|
return out
|
|
|
|
|
|
def worker(args):
|
|
key, tag, pimg, plabel, out_base, force, dpi = args
|
|
try:
|
|
meta = json.load(open(os.path.join(d("data/procmeta"), key + ".json")))
|
|
stats = label_stats(plabel)
|
|
except Exception:
|
|
return key, "error: unreadable label/procmeta"
|
|
reason = exclusion_reason(key, stats)
|
|
cat = reason.split(" (", 1)[0] if reason else None
|
|
# reconcile a stale figure for this key left in the other folder by an
|
|
# earlier run under a different exclusion status (e.g. a key newly added to
|
|
# NON_BRAIN_KEYS would otherwise keep its old inclusion/ figure)
|
|
folder = "exclusion" if reason else "inclusion"
|
|
other = "inclusion" if reason else "exclusion"
|
|
for stale in glob.glob(os.path.join(out_base, other, f"*_{key}.png")):
|
|
os.remove(stale)
|
|
if stats["total"] == 0:
|
|
return key, "empty-label", cat
|
|
status = "excluded" if reason else "included"
|
|
out = os.path.join(out_base, folder, out_name(stats["total"], key))
|
|
if os.path.exists(out) and not force:
|
|
return key, f"skip-{status}", cat
|
|
try:
|
|
mr = np.asarray(sitk.GetArrayFromImage(sitk.ReadImage(pimg)), dtype=np.float32)
|
|
lab = sitk.GetArrayFromImage(sitk.ReadImage(plabel))
|
|
cc, n_cc = largest_cc(lab > 0.5)
|
|
try:
|
|
render(key, tag, mr, cc, n_cc, int(stats["total"]),
|
|
np.array(meta["direction"]), reason, out, dpi)
|
|
finally:
|
|
plt.close("all")
|
|
return key, status, cat
|
|
except Exception as e: # noqa
|
|
return key, f"error: {e!r}", cat
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--from-splits", action="store_true",
|
|
help="rows/tags from split_{train,val,test}.jsonl instead of "
|
|
"the source manifests (ntuh + m6_labeled)")
|
|
ap.add_argument("--keys", nargs="*", default=[])
|
|
ap.add_argument("--limit", type=int, default=0)
|
|
ap.add_argument("--workers", type=int, default=32)
|
|
ap.add_argument("--force", action="store_true")
|
|
ap.add_argument("--out", default=d("data/tv_slices"))
|
|
ap.add_argument("--dpi", type=int, default=110)
|
|
args = ap.parse_args()
|
|
os.makedirs(os.path.join(args.out, "inclusion"), exist_ok=True)
|
|
os.makedirs(os.path.join(args.out, "exclusion"), exist_ok=True)
|
|
|
|
if args.from_splits:
|
|
rows = []
|
|
for s in ("train", "val", "test"):
|
|
for r in load_jsonl(os.path.join(DATA, "manifests", f"split_{s}.jsonl")):
|
|
rows.append((r["key"], s, r["pimg"], r.get("plabel")))
|
|
else:
|
|
rows = [(r["key"],
|
|
"ntuh" if r["key"].startswith("ntuh_") else "m6",
|
|
r["pimg"], r.get("plabel")) for r in labeled_source_rows()]
|
|
if args.keys:
|
|
keep = set(args.keys)
|
|
rows = [r for r in rows if r[0] in keep]
|
|
if args.limit:
|
|
rows = rows[: args.limit]
|
|
print(f"rows: {len(rows)}")
|
|
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
tasks = [(k, tag, pimg, plabel, args.out, args.force, args.dpi)
|
|
for k, tag, pimg, plabel in rows]
|
|
stats = {}
|
|
reasons = Counter()
|
|
done = 0
|
|
with ProcessPoolExecutor(max_workers=args.workers) as ex:
|
|
futs = {ex.submit(worker, t): t for t in tasks}
|
|
for fu in as_completed(futs):
|
|
k, status, cat = fu.result()
|
|
stats[status] = stats.get(status, 0) + 1
|
|
if status == "excluded" and cat:
|
|
reasons[cat] += 1
|
|
done += 1
|
|
if done % 100 == 0 or done == len(tasks):
|
|
print(f" {done}/{len(tasks)} {stats}", flush=True)
|
|
print(f"done: {stats}")
|
|
if reasons:
|
|
print("excluded by rule:")
|
|
for r, n in sorted(reasons.items()):
|
|
print(f" {n:4d} {r}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |