Furen Xiao
5e525c07c8
This commit introduces a new script, `setup_impact_konfai.sh`, which facilitates a one-time setup for the IMPACT registration pipeline (impact-reg-konfai). The script builds a self-contained shim that includes: - A glibc 2.35 loader and libraries extracted from the niftyreg:2.1.1-ubuntu Docker image. - A self-contained LibTorch 2.8.0 (cu128) library. - A wrapper script that runs the stock elastix-impact binary through the glibc 2.35 loader with the appropriate library paths. The setup ensures compatibility with hosts that do not meet the required glibc and LibTorch versions, allowing the pipeline to function correctly. The script also verifies the installation and provides usage instructions.
597 lines
No EOL
29 KiB
Python
597 lines
No EOL
29 KiB
Python
"""Benchmark the 6 rigid CT->T1c registration pipelines on N M6 test cases.
|
|
|
|
Pipelines (see src/registration_pipelines.py for the shared transform contract):
|
|
simpleitk | elastix | antspyx | synthmorph | impact | niftyreg
|
|
|
|
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. The headline alignment score is the achieved Mattes MI
|
|
(mi_ct_mr, nats) between the warped-CT and the original MR over the MR head --
|
|
higher = better in-sync = better aligned. The production geometric QA
|
|
(scripts/02_build_m6_dataset.py) is still recorded per case in the JSON:
|
|
|
|
- mi_ct_mr : Mattes MI (nats), warped-CT ~ MR over the MR head
|
|
(the reported alignment score).
|
|
- 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).
|
|
- skull_ring_gap_mm : concentricity of the warped-CT skull ring vs the MR
|
|
head silhouette (mid/low level + angular std).
|
|
- accepted : production rule (nvox>0 and tv_in_ct and concentric).
|
|
- reg_wall_s : wall-clock registration time (excludes I/O + scoring).
|
|
|
|
`--clip-ct LO HI` (e.g. `--clip-ct 0 80`, brain window) clips the CT to
|
|
[LO, HI] before registration AND before scoring, so a clipped vs unclipped
|
|
A/B compares the same input to both the toolkit and the MI score.
|
|
|
|
Because some toolkits are not pip-installable into every
|
|
`longitudinal-313` env (antspyx needs a cp310-cp313 wheel, see
|
|
REGISTRATION_BENCHMARK.md), 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. elastix runs in the main env through the
|
|
ITK-Elastix python API (pip ``itk`` + ``itk-elastix``); niftyreg runs natively
|
|
(``conda install -c conda-forge niftyreg`` provides ``reg_aladin``);
|
|
``impact`` runs the ``impact-reg-konfai`` package (pip; GPU; one-time elastix
|
|
shim setup on hosts too old for the stock binary — see
|
|
REGISTRATION_BENCHMARK.md); ``synthmorph`` runs in a Docker image
|
|
(``freesurfer/synthmorph``; needs Docker + the pulled image). The `--merge`
|
|
mode then combines the per-pipeline JSON files (accumulated across envs) into
|
|
one summary + plot.
|
|
|
|
Usage (repo root):
|
|
# main env (SimpleITK + elastix + antspyx + niftyreg + impact)
|
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
|
python scripts/bench_registration.py --pipelines simpleitk elastix antspyx niftyreg impact
|
|
|
|
# synthmorph (GPU, Docker) — one-time setup in REGISTRATION_BENCHMARK.md
|
|
python scripts/bench_registration.py --pipelines synthmorph
|
|
|
|
# any pipeline that is not importable in the current env is recorded as
|
|
# unavailable, so the script is env-portable; --merge combines the
|
|
# per-pipeline JSONs accumulated across envs 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, head_mask_from_image
|
|
from src import registration_pipelines as RP
|
|
|
|
|
|
def _mi_masked(fixed, moving, mask, bins=32):
|
|
"""Mattes MI (nats) between two co-registered volumes restricted to ``mask``,
|
|
each image percentile-binned independently (SimpleITK Mattes-MI style).
|
|
Higher = the warped-CT and MR intensities are better in-sync = better
|
|
alignment. ``fixed`` is the CT inverse-transformed into MR coordinates."""
|
|
f = np.asarray(fixed)[mask]
|
|
m = np.asarray(moving)[mask]
|
|
if f.size < 1000:
|
|
return None
|
|
|
|
def q(a):
|
|
lo, hi = np.percentile(a, (1, 99))
|
|
return np.clip(np.round((a - lo) * (bins - 1) / max(hi - lo, 1e-9)).astype(int),
|
|
0, bins - 1)
|
|
|
|
A, B = q(f), q(m)
|
|
h = np.bincount(A.ravel() * bins + B.ravel(), minlength=bins * bins)
|
|
p = h.reshape(bins, bins) / max(h.sum(), 1)
|
|
px, py = p.sum(axis=1, keepdims=True), p.sum(axis=0, keepdims=True)
|
|
with np.errstate(divide="ignore", invalid="ignore"):
|
|
ll = np.log(p / (px * py))
|
|
return float(np.where((p > 0) & (px > 0) & (py > 0), p * ll, 0.0).sum())
|
|
|
|
|
|
def _m6():
|
|
return RP._load_m6()
|
|
|
|
|
|
def _clip_image(img, lo, hi):
|
|
"""Clip image intensities to [lo, hi] (geometry preserved)."""
|
|
a = np.clip(sitk.GetArrayFromImage(img).astype("float32"), lo, hi)
|
|
out = sitk.GetImageFromArray(a)
|
|
out.CopyInformation(img)
|
|
return out
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 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)
|
|
# achieved-metric alignment score: Mattes MI between the CT inverse-transformed
|
|
# into MR coordinates (warped_ct) and the original MR (ha), over the MR head.
|
|
ha_f = ha.astype("float32", copy=False)
|
|
mi = _mi_masked(warped_ct, ha_f, head_mask_from_image(ha_f))
|
|
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 abs(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,
|
|
"mi_ct_mr": round(mi, 4) if mi is not None else None,
|
|
"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 make_qa_png(warped_ct, mr, out_png, title, block=24, ct_brain_window=(0.0, 80.0)):
|
|
"""checkreg-style visual QA of (inverse-warped CT) vs (original MR):
|
|
rows = the three orthogonal mid-slices through the head center
|
|
(mid-axial, mid-coronal, mid-sagittal); columns = left checkerboard blend,
|
|
right additive color overlay (CT=red, MR=green -> yellow where aligned,
|
|
red/green fringes where not). The CT is shown in a brain window
|
|
(default L40/W80 = 0..80 HU). Writes ``out_png``."""
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
mr = np.asarray(mr, dtype="float32")
|
|
wc = np.asarray(warped_ct, dtype="float32")
|
|
pos = mr > 0
|
|
thr = np.percentile(mr[pos], 2) if pos.any() else 0.0
|
|
head = mr > max(thr, 1e-6)
|
|
cen = np.argwhere(head).mean(0) if head.any() else np.array(mr.shape) / 2.0
|
|
ia, ic, isg = (int(round(cen[0])), int(round(cen[1])), int(round(cen[2])))
|
|
# orthogonal mid-slices (arrays are [z, y, x]); CT and MR sliced identically
|
|
# so the two stay registered in each view. coronal/sagittal flip the z-axis
|
|
# (superior up) to match the axial view.
|
|
def _flip0(a):
|
|
return np.ascontiguousarray(a[::-1, :])
|
|
|
|
views = [
|
|
("Mid-axial", np.ascontiguousarray(mr[ia]), np.ascontiguousarray(wc[ia])),
|
|
("Mid-coronal", _flip0(mr[:, ic, :]), _flip0(wc[:, ic, :])),
|
|
("Mid-sagittal", _flip0(mr[:, :, isg]), _flip0(wc[:, :, isg])),
|
|
]
|
|
mr_lo, mr_hi = np.percentile(mr[head], (2, 98)) if head.any() else (0.0, 1.0)
|
|
ct_lo, ct_hi = ct_brain_window # CT brain window (0..80 HU default)
|
|
|
|
def w(a, lo, hi):
|
|
return np.clip((a - lo) / max(hi - lo, 1e-9), 0.0, 1.0)
|
|
|
|
n = len(views)
|
|
fig, axes = plt.subplots(nrows=n, ncols=2, figsize=(7.6, 2.15 * n))
|
|
for r, (vname, mrs_raw, wcs_raw) in enumerate(views):
|
|
mrs, wcs = w(mrs_raw, mr_lo, mr_hi), w(wcs_raw, ct_lo, ct_hi)
|
|
yy, xx = np.indices(mrs.shape)
|
|
cb = mrs.copy()
|
|
cb[((yy // block) + (xx // block)) % 2 == 0] = wcs[((yy // block) + (xx // block)) % 2 == 0]
|
|
axes[r, 0].imshow(cb, cmap="gray", vmin=0, vmax=1); axes[r, 0].axis("off")
|
|
axes[r, 0].text(0.01, 0.02, vname, transform=axes[r, 0].transAxes, color="white",
|
|
bbox=dict(facecolor="black", alpha=0.55, pad=1.2), fontsize=8)
|
|
ov = np.zeros((mrs.shape[0], mrs.shape[1], 3))
|
|
ov[..., 0] = wcs; ov[..., 1] = mrs
|
|
ov[..., 2] = 0.28 * np.minimum(wcs, mrs)
|
|
axes[r, 1].imshow(ov); axes[r, 1].axis("off")
|
|
axes[r, 1].text(0.01, 0.02, vname, transform=axes[r, 1].transAxes, color="white",
|
|
bbox=dict(facecolor="black", alpha=0.55, pad=1.2), fontsize=8)
|
|
axes[0, 0].set_title("Checkerboard blend (CT | MR)", fontsize=9)
|
|
axes[0, 1].set_title("Color overlay (CT=red, MR=green)", fontsize=9)
|
|
fig.suptitle(title, fontsize=10)
|
|
fig.tight_layout(rect=(0, 0, 1, 0.97))
|
|
fig.savefig(out_png, dpi=110)
|
|
plt.close(fig)
|
|
|
|
|
|
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 == "synthmorph":
|
|
return RP.register_synthmorph(ct, t1c, model=args.synthmorph_model,
|
|
gpu=args.synthmorph_gpu, threads=args.threads,
|
|
extent=args.synthmorph_extent)
|
|
if pipeline == "impact":
|
|
return RP.register_impact(ct, t1c, preset=args.impact_preset, gpu=args.impact_gpu,
|
|
iters=args.impact_iters, samples=args.impact_samples)
|
|
if pipeline == "niftyreg":
|
|
return RP.register_niftyreg(ct, t1c, threads=args.threads,
|
|
backend=args.niftyreg_backend,
|
|
gpu=args.niftyreg_gpu)
|
|
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
|
|
if args.clip_ct is not None:
|
|
# clip BEFORE registration AND scoring so the A/B compares like for like
|
|
lo, hi = args.clip_ct
|
|
ct = _clip_image(ct, lo, hi)
|
|
out["clip_ct"] = [lo, hi]
|
|
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))
|
|
if getattr(args, "qa_images", False):
|
|
title = f"{pipeline} | {key} | MI={out.get('mi_ct_mr')}"
|
|
qp = os.path.join(args.qa_out, f"{pipeline}__{key}.png")
|
|
try:
|
|
wc = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
|
|
make_qa_png(wc, sitk.GetArrayFromImage(t1c), qp, title)
|
|
out["qa_png"] = qp
|
|
except Exception as e: # noqa
|
|
out["qa_png_error"] = repr(e)
|
|
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]
|
|
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']}: mi={r.get('mi_ct_mr')} "
|
|
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]
|
|
mifr = [r["mi_ct_mr"] for r in registered if r.get("mi_ct_mr") 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,
|
|
"mi_ct_mr": {"mean": round(float(np.mean(mifr)), 4),
|
|
"median": round(float(np.median(mifr)), 4)} if mifr 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" + "=" * 92)
|
|
hdr = (f"{'pipeline':10s} {'avail':6s} {'toolkit':22s} {'n':>3s} {'ok':>3s} {'err':>4s} "
|
|
f"{'s/reg':>7s} {'tvInCT':>7s} {'MI':>7s}")
|
|
print(hdr)
|
|
print("-" * 92)
|
|
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} [{r['reason']}]")
|
|
continue
|
|
rt = (r.get("reg_wall_s") or {}).get("median")
|
|
tv = (r.get("tv_in_ct_frac") or {}).get("mean")
|
|
mi = (r.get("mi_ct_mr") or {}).get("median")
|
|
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} {_fmt(rt, '{:.1f}'):>7s} "
|
|
f"{_fmt(tv, '{:.2f}'):>7s} {_fmt(mi, '{:.3f}'):>7s}")
|
|
print("=" * 92)
|
|
|
|
|
|
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)
|
|
regt = [(ok_rows[p].get("reg_wall_s") or {}).get("median") or 0 for p in labels]
|
|
mi = [(ok_rows[p].get("mi_ct_mr") or {}).get("median") or 0 for p in labels]
|
|
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
|
|
colors = ["#4878CF", "#6ACC65"]
|
|
for ax, (title, vals, fmt, ylab) in zip(axes, [
|
|
("Registration time (s, median)", regt, "{:.0f}", "s"),
|
|
("Mattes MI: warped-CT ~ MR (nats, median)", mi, "{:.3f}", "nats")]):
|
|
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; alignment scored by achieved Mattes MI "
|
|
"(nats, warped-CT ~ MR over the MR head, higher = better); reg time "
|
|
"excludes I/O+scoring; per-case JSON retains the production skull-ring "
|
|
"gap / acceptance QA; best-of-T vs T^-1 orientation gate for the "
|
|
"elastix-style toolkits"),
|
|
"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 | s/reg (median) | tv-in-CT | MI (nats, median) |",
|
|
"|---|---|---|---|---|---|---|---|---|"]
|
|
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']} | "
|
|
f"{(r.get('reg_wall_s') or {}).get('median')} | "
|
|
f"{(r.get('tv_in_ct_frac') or {}).get('mean')} | "
|
|
f"{(r.get('mi_ct_mr') or {}).get('median')} |")
|
|
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("--synthmorph-model", choices=["rigid", "affine"], default="rigid",
|
|
help="SynthMorph matrix model (default rigid; deformable/joint are "
|
|
"outside the rigid-transform contract)")
|
|
ap.add_argument("--synthmorph-gpu", type=int, default=None,
|
|
help="GPU device index for SynthMorph (default: auto-spread over "
|
|
"the host GPUs by pid; -1 forces CPU)")
|
|
ap.add_argument("--synthmorph-extent", type=int, default=256,
|
|
help="SynthMorph isotropic registration-space extent in mm "
|
|
"(192 or 256; default 256)")
|
|
ap.add_argument("--impact-preset", default=RP.IMPACT_DEFAULT_PRESET,
|
|
help="ImpactReg preset from the published database "
|
|
f"(default: {RP.IMPACT_DEFAULT_PRESET}; e.g. "
|
|
"Elastix_IMPACT_Jacobian, Generic_Rigid, "
|
|
"Generic_Rigid_BSpline, FireANTs_SyN, FireANTs_IMPACT, "
|
|
"ConvexAdam_Coarse/Composite; list live with: "
|
|
"impact-reg-konfai register --help)")
|
|
ap.add_argument("--impact-gpu", type=int, default=None,
|
|
help="GPU device index for IMPACT (default: auto-spread over "
|
|
"the host GPUs by pid; -1 forces CPU)")
|
|
ap.add_argument("--impact-iters", type=int, default=None,
|
|
help="preset max_iterations override, one value for all "
|
|
"resolutions (only effective for presets without a "
|
|
"per-resolution model matrix; default: the preset's "
|
|
"own budget, 400/300/200/200 for Elastix_IMPACT_Static)")
|
|
ap.add_argument("--impact-samples", type=int, default=None,
|
|
help="preset spatial_samples override (NumberOfSpatialSamples; "
|
|
"default: the preset's own value, 2000)")
|
|
ap.add_argument("--niftyreg-backend", choices=["auto", "conda", "docker"],
|
|
default="auto",
|
|
help="niftyreg reg_aladin build: docker = official v2.1.1 in the "
|
|
"niftyreg:2.1.1-ubuntu image (host glibc is too old for the "
|
|
"v2.x binaries), conda = conda-forge 1.5.69 on PATH, "
|
|
"auto (default) = docker when the image is present else conda")
|
|
ap.add_argument("--niftyreg-gpu", type=int, default=None,
|
|
help="niftyreg CUDA device index (docker backend; default: "
|
|
"auto-spread over the host GPUs by pid, -1 forces CPU "
|
|
"-platf 0; on the GPU the run uses 'docker run --gpus all "
|
|
"-platf 1' with linear interpolation, the v2.1.1 CUDA "
|
|
"limit)")
|
|
ap.add_argument("--clip-ct", nargs=2, type=float, default=None, metavar=("LO", "HI"),
|
|
help="clip the CT (HU) to [LO, HI] BEFORE registration and scoring "
|
|
"(e.g. --clip-ct 0 80 = brain window, skull/air removed; "
|
|
"default: no clipping)")
|
|
ap.add_argument("--qa-images", action="store_true",
|
|
help="also write a checkreg-style PNG (checkerboard blend + color "
|
|
"overlay of the inverse-warped CT vs the original MR) per case")
|
|
ap.add_argument("--qa-out", default=None,
|
|
help="QA image dir (default <out-dir>/qa)")
|
|
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")
|
|
if args.qa_images:
|
|
args.qa_out = os.path.abspath(args.qa_out or os.path.join(args.out_dir, "qa"))
|
|
os.makedirs(args.qa_out, exist_ok=True)
|
|
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() |