Furen Xiao
e8dfe0fe35
Update M6 manifest building to register T1c to CT with rigid Mattes MI on a CT-space head box, including multi-start initial tilts for large inter-session head motion. The inverse transform warps CT and GTV into T1c native space, with acceptance based on concentricity of the warped CT skull ring and MR head silhouette. Rerun preprocessing when a row has a label but the warped label file is missing. Normalize volume orientation in read_nii_img. Add registration benchmark documentation, benchmark script, and shared registration pipeline helpers.
423 lines
No EOL
19 KiB
Python
423 lines
No EOL
19 KiB
Python
"""Benchmark the 4 rigid CT->T1c registration pipelines on N M6 test cases.
|
|
|
|
Pipelines (see src/registration_pipelines.py for the shared transform contract):
|
|
simpleitk | elastix | antspyx | deepreg
|
|
|
|
Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI
|
|
(cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto
|
|
the T1c native grid. Every pipeline is scored with the SAME crop-independent
|
|
geometric QA used in production (scripts/02_build_m6_dataset.py):
|
|
|
|
- skull_ring_gap_mm : concentricity of the warped-CT skull ring vs the MR
|
|
head silhouette (mid/low level + angular std).
|
|
Aligned = 0 <= mid <= 20 mm and std <= 12 mm.
|
|
- tv_in_ct_frac : fraction of the (dilated) warped GTV that falls inside
|
|
the warped-CT head tissue (the slab must straddle it).
|
|
- tv_nvox : warped GTV volume (>0.5 voxels).
|
|
- accepted : production rule (nvox>0 and tv_in_ct and concentric).
|
|
- reg_wall_s : wall-clock registration time (excludes I/O + scoring).
|
|
|
|
Because three of the four toolkits are not pip-installable into the main
|
|
`longitudinal` env (elastix needs conda-forge; antspyx has no cp314 wheel;
|
|
PyPI deepreg is a broken stub), the benchmark is env-portable: run it under
|
|
whichever conda env has a given toolkit installed, and it records only the
|
|
pipelines that are importable there. The `--merge` mode then combines the
|
|
per-pipeline JSON files (accumulated across envs) into one summary + plot.
|
|
|
|
Usage (repo root):
|
|
# main env (SimpleITK baseline; elastix/deepreg report as unavailable)
|
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
|
python scripts/bench_registration.py --pipelines simpleitk elastix deepreg
|
|
|
|
# isolated env that has antspyx (see REGISTRATION_BENCHMARK.md)
|
|
source /opt/conda/etc/profile.d/conda.sh && conda activate <ants312>
|
|
python scripts/bench_registration.py --pipelines antspyx simpleitk
|
|
|
|
# combine everything produced so far into one table + plot
|
|
python scripts/bench_registration.py --merge
|
|
"""
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, ROOT)
|
|
|
|
import numpy as np
|
|
import SimpleITK as sitk
|
|
from scipy import ndimage
|
|
|
|
from src.common import DATA, d, read_nii_img
|
|
from src import registration_pipelines as RP
|
|
|
|
|
|
def _m6():
|
|
return RP._load_m6()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# case selection
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def ct_complete(ct):
|
|
"""Production guard: skip corrupt CTs (upper brain missing)."""
|
|
m6 = _m6()
|
|
vz = m6._ct_vertex_slice(ct)
|
|
if vz is None:
|
|
return False
|
|
sp = ct.GetSpacing()[2]
|
|
per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2))
|
|
iz = vz - int(round(100 / sp))
|
|
return bool(0 <= iz < len(per) and per[iz] > 5000)
|
|
|
|
|
|
def select_cases(args):
|
|
"""Pick N labeled M6 cases deterministically. No CT I/O here -- each file
|
|
is read exactly once, inside the registration worker (NFS is the bottleneck)."""
|
|
m6 = _m6()
|
|
labeled, _unlabeled = m6.discover()
|
|
print(f"[cases] m6 labeled candidates (CT+T1c+TV): {len(labeled)}")
|
|
pool = list(labeled)
|
|
if args.sample == "random":
|
|
idx = np.random.default_rng(args.seed).choice(
|
|
len(pool), min(args.n_cases, len(pool)), replace=False)
|
|
sel = [pool[i] for i in sorted(idx.tolist())]
|
|
else:
|
|
sel = pool[:args.n_cases]
|
|
print(f"[cases] selected {len(sel)} test cases (sample={args.sample} seed={args.seed})")
|
|
for r in sel:
|
|
print(f" - {r['key']}")
|
|
return sel
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# shared evaluation (identical for every pipeline)
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def evaluate(ct, t1c, tv, T):
|
|
"""Warp CT + GTV onto the T1c grid with T and score them (production QA)."""
|
|
m6 = _m6()
|
|
ha = sitk.GetArrayFromImage(t1c)
|
|
warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
|
|
warped_tv = sitk.GetArrayFromImage(sitk.Resample(tv, t1c, T, sitk.sitkNearestNeighbor, 0.0))
|
|
gap = m6._skull_ring_gap(warped_ct, ha)
|
|
posw = warped_ct[warped_ct > 0]
|
|
tissue_w = (warped_ct > max(2.0, float(np.percentile(posw, 2)))) if posw.size \
|
|
else np.zeros_like(warped_ct, bool)
|
|
tissue_w = ndimage.binary_opening(tissue_w, iterations=2)
|
|
nvox = int((warped_tv > 0.5).sum())
|
|
tv_in_ct_frac = 1.0
|
|
if nvox > 0:
|
|
tvd = ndimage.binary_dilation(warped_tv > 0.5, iterations=3)
|
|
tv_in_ct_frac = float((tvd & tissue_w).sum() / max(tvd.sum(), 1))
|
|
ok_gap = bool(gap) and 0.0 <= gap["mid_mm"] <= 20.0 and (gap["std_mm"] or 0) <= 12.0
|
|
accepted = nvox > 0 and tv_in_ct_frac >= 0.5 and ok_gap
|
|
return {
|
|
"skull_ring_gap_mm": gap,
|
|
"tv_nvox": nvox,
|
|
"tv_in_ct_frac": round(tv_in_ct_frac, 3),
|
|
"alignment": "concentric" if ok_gap else "not_concentric",
|
|
"accepted": bool(accepted),
|
|
}
|
|
|
|
|
|
def _call(pipeline, ct, t1c, tv, args):
|
|
if pipeline == "simpleitk":
|
|
return RP.register_simpleitk(ct, t1c, threads=args.threads, mode=args.simpleitk_mode)
|
|
if pipeline == "elastix":
|
|
return RP.register_elastix(ct, t1c, workdir=args.elastix_dir, threads=args.threads)
|
|
if pipeline == "antspyx":
|
|
return RP.register_antspyx(ct, t1c)
|
|
if pipeline == "deepreg":
|
|
return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir,
|
|
checkpoint=args.deepreg_ckpt)
|
|
raise KeyError(pipeline)
|
|
|
|
|
|
def _run_case(pipeline, case, args):
|
|
"""Worker body: read once + register + score one case. Runs in a forked
|
|
process. The CT-completeness guard is applied here (on the already-read CT)
|
|
so each file is read exactly once."""
|
|
key = case["key"]
|
|
ct = read_nii_img(case["ct"])
|
|
t1c = read_nii_img(case["img"])
|
|
tv = read_nii_img(case["label"])
|
|
out = {"key": key, "subject": case["subject"], "date": case["date"]}
|
|
if not ct_complete(ct):
|
|
out.update({"skull_ring_gap_mm": None, "tv_nvox": 0, "tv_in_ct_frac": None,
|
|
"alignment": "ct_incomplete", "accepted": False,
|
|
"note": "CT head incomplete (data issue, not a pipeline failure)"})
|
|
return out
|
|
t0 = time.time()
|
|
T, meta = _call(pipeline, ct, t1c, tv, args)
|
|
out["reg_wall_s"] = round(time.time() - t0, 1)
|
|
out.update(meta)
|
|
out.update(evaluate(ct, t1c, tv, T))
|
|
return out
|
|
|
|
|
|
def run_pipeline(pipeline, cases, args):
|
|
avail, note = RP.availability().get(pipeline, (False, "unknown"))
|
|
if not avail:
|
|
row = {"pipeline": pipeline, "available": False, "reason": note,
|
|
"n_cases": len(cases), "cases": []}
|
|
print(f"[{pipeline}] UNAVAILABLE: {note}")
|
|
return row
|
|
print(f"[{pipeline}] available ({note}); registering {len(cases)} cases "
|
|
f"with {args.workers} workers")
|
|
results = []
|
|
errors = []
|
|
with ProcessPoolExecutor(max_workers=args.workers) as ex:
|
|
futs = {ex.submit(_run_case, pipeline, c, args): c for c in cases}
|
|
for fu in as_completed(futs):
|
|
c = futs[fu]
|
|
try:
|
|
results.append(fu.result())
|
|
r = results[-1]
|
|
g = r.get("skull_ring_gap_mm")
|
|
reg = f"reg={r.get('reg_wall_s')}s" if r.get("reg_wall_s") is not None else "incomplete-CT"
|
|
print(f" {c['key']}: accepted={r['accepted']} "
|
|
f"gap={g['mid_mm'] if g else 'None'}mm "
|
|
f"std={g['std_mm'] if g else 'None'}mm "
|
|
f"tv_in_ct={r.get('tv_in_ct_frac')} {reg}", flush=True)
|
|
except Exception as e: # noqa
|
|
errors.append({"key": c["key"], "error": repr(e)})
|
|
print(f" {c['key']}: ERROR {e!r}", flush=True)
|
|
results.sort(key=lambda x: x["key"])
|
|
# aggregate
|
|
registered = [r for r in results if r.get("reg_wall_s") is not None]
|
|
incomplete = [r for r in results if r.get("alignment") == "ct_incomplete"]
|
|
accepted = [r for r in registered if r["accepted"]]
|
|
gaps_mid = [r["skull_ring_gap_mm"]["mid_mm"] for r in registered if r.get("skull_ring_gap_mm")]
|
|
gaps_std = [r["skull_ring_gap_mm"]["std_mm"] for r in registered if r.get("skull_ring_gap_mm")]
|
|
regt = [r["reg_wall_s"] for r in registered]
|
|
tvfr = [r["tv_in_ct_frac"] for r in registered if r.get("tv_in_ct_frac") is not None]
|
|
first = (registered or results)[0] if (registered or results) else {}
|
|
row = {
|
|
"pipeline": pipeline, "available": True, "toolkit": first.get("toolkit"),
|
|
"metric": first.get("metric"), "transform": first.get("transform"),
|
|
"n_cases": len(cases), "n_registered": len(registered),
|
|
"n_ct_incomplete": len(incomplete), "n_error": len(errors),
|
|
"n_accepted": len(accepted),
|
|
"accepted_rate": round(len(accepted) / max(len(cases), 1), 3),
|
|
"reg_wall_s": {"mean": round(float(np.mean(regt)), 1), "median": round(float(np.median(regt)), 1),
|
|
"max": round(float(np.max(regt)), 1)} if regt else None,
|
|
"skull_ring_gap_mid_mm": {"mean": round(float(np.mean(gaps_mid)), 1),
|
|
"median": round(float(np.median(gaps_mid)), 1)} if gaps_mid else None,
|
|
"skull_ring_gap_std_mm": {"mean": round(float(np.mean(gaps_std)), 1),
|
|
"median": round(float(np.median(gaps_std)), 1)} if gaps_std else None,
|
|
"tv_in_ct_frac": {"mean": round(float(np.mean(tvfr)), 3)} if tvfr else None,
|
|
"cases": results, "errors": errors,
|
|
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
|
}
|
|
return row
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# reporting / merge
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def _write(row, out_dir):
|
|
p = os.path.join(out_dir, row["pipeline"] + ".json")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
with open(p, "w") as f:
|
|
json.dump(row, f, indent=1)
|
|
return p
|
|
|
|
|
|
def _fmt(v, spec="{}", default="-"):
|
|
return default if v is None else spec.format(v)
|
|
|
|
|
|
def print_table(rows):
|
|
order = [p for p in RP.ALL_PIPELINES if any(r["pipeline"] == p for r in rows)]
|
|
print("\n" + "=" * 100)
|
|
hdr = (f"{'pipeline':10s} {'avail':6s} {'toolkit':22s} {'n':>3s} {'ok':>3s} {'err':>4s} "
|
|
f"{'accept':>7s} {'gapMid':>7s} {'gapStd':>7s} {'s/reg':>7s} {'tvInCT':>7s}")
|
|
print(hdr)
|
|
print("-" * 100)
|
|
for p in order:
|
|
r = next(x for x in rows if x["pipeline"] == p)
|
|
if not r.get("available"):
|
|
print(f"{p:10s} {('no'):6s} {'-':22s} {r['n_cases']:>3d} {0:>3d} {'-':>4s} "
|
|
f"{'-':>7s} {'-':>7s} {'-':>7s} {'-':>7s} {'-':>7s} [{r['reason']}]")
|
|
continue
|
|
gm = (r.get("skull_ring_gap_mid_mm") or {}).get("median")
|
|
gs = (r.get("skull_ring_gap_std_mm") or {}).get("median")
|
|
rt = (r.get("reg_wall_s") or {}).get("median")
|
|
tv = (r.get("tv_in_ct_frac") or {}).get("mean")
|
|
print(f"{p:10s} {'yes':6s} {(r.get('toolkit') or '')[:22]:22s} {r['n_cases']:>3d} "
|
|
f"{r.get('n_registered', 0):>3d} {r['n_error']:>4d} {r['accepted_rate']:>7.2f} "
|
|
f"{_fmt(gm, '{:.1f}'):>7s} {_fmt(gs, '{:.1f}'):>7s} {_fmt(rt, '{:.1f}'):>7s} "
|
|
f"{_fmt(tv, '{:.2f}'):>7s}")
|
|
print("=" * 100)
|
|
|
|
|
|
def write_plot(rows, out_png):
|
|
try:
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
except Exception as e: # noqa
|
|
print("plot skipped:", repr(e))
|
|
return
|
|
ok_rows = {r["pipeline"]: r for r in rows if r.get("available")}
|
|
if not ok_rows:
|
|
print("no available pipelines to plot")
|
|
return
|
|
labels = list(ok_rows)
|
|
gap_mid = [ok_rows[p].get("skull_ring_gap_mid_mm", {}).get("median") or 0 for p in labels]
|
|
gap_std = [ok_rows[p].get("skull_ring_gap_std_mm", {}).get("median") or 0 for p in labels]
|
|
regt = [ok_rows[p].get("reg_wall_s", {}).get("median") or 0 for p in labels]
|
|
acc = [ok_rows[p].get("accepted_rate") or 0 for p in labels]
|
|
fig, axes = plt.subplots(1, 4, figsize=(18, 4))
|
|
colors = ["#4878CF", "#D64F5C", "#6ACC65", "#9467BD"]
|
|
for ax, (title, vals, fmt, ylab, inv) in zip(axes, [
|
|
("Skull-ring gap mid (mm, median)", gap_mid, "{:.1f}", "mm", True),
|
|
("Skull-ring gap std (mm, median)", gap_std, "{:.1f}", "mm", True),
|
|
("Registration time (s, median)", regt, "{:.0f}", "s", True),
|
|
("Acceptance rate", acc, "{:.2f}", "frac", False)]):
|
|
bars = ax.bar(range(len(vals)), vals, color=colors[:len(vals)])
|
|
for b, v in zip(bars, vals):
|
|
ax.annotate(fmt.format(v), (b.get_x() + b.get_width() / 2, b.get_height()),
|
|
ha="center", va="bottom", fontsize=8)
|
|
ax.set_title(title, fontsize=9)
|
|
ax.set_ylabel(ylab)
|
|
ax.set_xticks(range(len(vals)))
|
|
ax.set_xticklabels(labels, fontsize=9)
|
|
plt.suptitle("CT->T1c registration benchmark (M6, cross-modality rigid)", fontsize=11)
|
|
fig.tight_layout()
|
|
fig.savefig(out_png, dpi=120)
|
|
plt.close(fig)
|
|
print("plot:", out_png)
|
|
|
|
|
|
def merge(args):
|
|
out_dir = os.path.join(args.out_dir)
|
|
rows = []
|
|
for p in RP.ALL_PIPELINES:
|
|
fp = os.path.join(out_dir, p + ".json")
|
|
if os.path.exists(fp):
|
|
with open(fp) as f:
|
|
rows.append(json.load(f))
|
|
if not rows:
|
|
raise SystemExit(f"no pipeline results found in {out_dir}")
|
|
print(f"[merge] {len(rows)} pipeline result files in {out_dir}")
|
|
print_table(rows)
|
|
summary = {
|
|
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
|
"protocol": ("rigid CT->T1c; cross-modality; scored by production skull-ring "
|
|
"concentricity (aligned: 0<=mid<=20mm, std<=12mm) + GTV-in-CT fraction "
|
|
"(>=0.5); reg time excludes I/O+scoring; per-pipeline best-of-T vs T^-1 "
|
|
"orientation gate for elastix/antspyx"),
|
|
"n_cases": next((r["n_cases"] for r in rows if r.get("available")), None),
|
|
"case_keys": sorted({c["key"] for r in rows if r.get("available") for c in r.get("cases", [])}),
|
|
"results": rows,
|
|
}
|
|
sf = os.path.join(out_dir, "summary.json")
|
|
with open(sf, "w") as f:
|
|
json.dump(summary, f, indent=1)
|
|
print("summary:", sf)
|
|
write_plot(rows, os.path.join(out_dir, "bench_registration.png"))
|
|
# markdown table for the report
|
|
md = ["| pipeline | available | toolkit | cases | ok | err | accepted | gap mid (mm) | gap std (mm) | s/reg (median) | tv-in-CT |",
|
|
"|---|---|---|---|---|---|---|---|---|---|---|"]
|
|
for p in [x["pipeline"] for x in rows]:
|
|
r = next(x for x in rows if x["pipeline"] == p)
|
|
if not r.get("available"):
|
|
md.append(f"| {p} | NO | - | {r['n_cases']} | 0 | - | - | - | - | - | - | ({r['reason']}) |")
|
|
continue
|
|
md.append(
|
|
f"| {p} | yes | {(r.get('toolkit') or '')} | {r['n_cases']} | {r.get('n_registered', 0)} | "
|
|
f"{r['n_error']} | {r['accepted_rate']:.2f} | "
|
|
f"{(r.get('skull_ring_gap_mid_mm') or {}).get('median')} | "
|
|
f"{(r.get('skull_ring_gap_std_mm') or {}).get('median')} | "
|
|
f"{(r.get('reg_wall_s') or {}).get('median')} | "
|
|
f"{(r.get('tv_in_ct_frac') or {}).get('mean')} |")
|
|
mf = os.path.join(out_dir, "summary.md")
|
|
with open(mf, "w") as f:
|
|
f.write("\n".join(md) + "\n")
|
|
print("markdown:", mf)
|
|
|
|
|
|
def parse_args():
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--pipelines", default=" ".join(RP.ALL_PIPELINES),
|
|
help="space-separated subset of: " + " ".join(RP.ALL_PIPELINES))
|
|
ap.add_argument("--n-cases", type=int, default=10, help="number of M6 test cases (default 10)")
|
|
ap.add_argument("--sample", choices=["first", "random"], default="first",
|
|
help="first N by discover order, or a seeded random sample of N")
|
|
ap.add_argument("--seed", type=int, default=0, help="seed for --sample random")
|
|
ap.add_argument("--keys", default=None, help="comma-separated explicit case keys (overrides selection)")
|
|
ap.add_argument("--workers", type=int, default=4, help="cases registered in parallel")
|
|
ap.add_argument("--threads", type=int, default=12, help="per-registration CPU threads")
|
|
ap.add_argument("--simpleitk-mode", choices=["benchmark", "production"], default="benchmark",
|
|
help="benchmark = single pyramid (comparable budget to other toolkits); "
|
|
"production = full multi-start tilt grid (slow)")
|
|
ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_registration"))
|
|
ap.add_argument("--elastix-dir", default=None, help="elastix scratch dir (default data/m6reg/elastix)")
|
|
ap.add_argument("--deepreg-dir", default=None, help="deepreg scratch dir")
|
|
ap.add_argument("--deepreg-ckpt", default=None, help="path to a DeepReg 3D pretrained checkpoint")
|
|
ap.add_argument("--merge", action="store_true", help="combine existing per-pipeline JSONs")
|
|
args = ap.parse_args()
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
if args.merge:
|
|
merge(args)
|
|
return
|
|
args.elastix_dir = args.elastix_dir or d("data/m6reg/elastix")
|
|
args.deepreg_dir = args.deepreg_dir or d("data/m6reg/deepreg")
|
|
pipelines = [p.lower() for p in args.pipelines.split()]
|
|
unknown = set(pipelines) - set(RP.ALL_PIPELINES)
|
|
if unknown:
|
|
raise SystemExit(f"unknown pipeline(s): {unknown}")
|
|
|
|
if args.keys:
|
|
m6 = _m6()
|
|
labeled, _ = m6.discover()
|
|
bykey = {r["key"]: r for r in labeled}
|
|
missing = [k for k in args.keys.split(",") if k not in bykey]
|
|
if missing:
|
|
raise SystemExit(f"keys not found in m6 labeled: {missing}")
|
|
cases = [bykey[k.strip()] for k in args.keys.split(",")]
|
|
print(f"[cases] using {len(cases)} explicit keys")
|
|
else:
|
|
cases = select_cases(args)
|
|
|
|
print(f"\n[bench] pipelines={pipelines} n_cases={len(cases)} workers={args.workers} "
|
|
f"threads={args.threads} out_dir={args.out_dir}")
|
|
rows = []
|
|
for p in pipelines:
|
|
try:
|
|
rows.append(run_pipeline(p, cases, args))
|
|
except RP.PipelineUnavailable as e:
|
|
rows.append({"pipeline": p, "available": False, "reason": str(e),
|
|
"n_cases": len(cases), "cases": []})
|
|
print(f"[{p}] UNAVAILABLE: {e}")
|
|
for r in rows:
|
|
fp = _write(r, args.out_dir)
|
|
print(f"[{r['pipeline']}] wrote {fp}")
|
|
print_table(rows)
|
|
# run-level availability snapshot (which env saw what)
|
|
avail = RP.availability()
|
|
af = os.path.join(args.out_dir, "_availability.json")
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
with open(af, "w") as f:
|
|
json.dump({"timestamp": datetime.now().isoformat(timespec="seconds"),
|
|
"env_python": sys.version.split()[0],
|
|
"exec": sys.executable, "availability": avail}, f, indent=1)
|
|
print(f"[bench] availability in this env: " +
|
|
", ".join(f"{k}={'yes' if v[0] else 'no'}" for k, v in avail.items()))
|
|
print(f"[bench] done. Re-run with --merge to combine across envs.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |