xfr/longitudinal
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
longitudinal/scripts/bench_registration.py

698 lines
No EOL
34 KiB
Python
Raw Blame History

"""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, higher = better).
- impact_loss : IMPACT semantic loss (vboussot/ImpactLoss; TotalSeg
M730 deep features, weighted L1), warped-CT ~ MR over
the full co-registered volumes (lower = better).
- 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_RIGID.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 ``vboussot/elastix_impact`` Docker image (GPU; the rigid
MIND + TotalSegmentator static config, needs Docker + the pulled image + the
model assets — see REGISTRATION_BENCHMARK_RIGID.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_RIGID.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()
# the ImpactLoss scoring model (the repo's default baseline for MR/CT)
IMPACT_SCORE_MODEL = "TS/M730.pt"
_IMPACT_LOSS_STATE = None
def _impact_loss_scorer():
"""Lazily build the ImpactLoss semantic-loss scorer (once per process) on
the GPU auto-picked for this worker. Returns ``(fn, device)`` or None if
torch / the model are unavailable (the score is then omitted)."""
global _IMPACT_LOSS_STATE
if _IMPACT_LOSS_STATE is None:
try:
import warnings
import torch
from src.impact_loss import IMPACTReg
ngpu = torch.cuda.device_count() if torch.cuda.is_available() else 0
if ngpu:
torch.cuda.set_device(os.getpid() % ngpu)
with warnings.catch_warnings():
# torch.jit.load is marked deprecated by upstream torch 2.14;
# the IMPACT models are TorchScript and jit.load is required
warnings.simplefilter("ignore", FutureWarning)
fn = IMPACTReg(IMPACT_SCORE_MODEL, shape=[0, 0, 0], in_channels=1,
weights=[1.0, 1.0], model_cache=d("data/impact_models"))
_IMPACT_LOSS_STATE = (fn, torch.device("cuda" if ngpu else "cpu"))
except Exception: # noqa
_IMPACT_LOSS_STATE = False
return _IMPACT_LOSS_STATE or None
def _impact_loss_value(warped_ct, mr):
"""IMPACT semantic loss (lower = more in-sync) between the co-registered
pair: the CT inverse-transformed into MR coordinates vs the original MR,
over the full volumes (the IMPACT models standardize intensity inside
the metric, so both go in raw)."""
st = _impact_loss_scorer()
if st is None:
return None
fn, dev = st
import torch
a = torch.as_tensor(np.ascontiguousarray(warped_ct, dtype="float32"),
device=dev).unsqueeze(0).unsqueeze(0)
b = torch.as_tensor(np.ascontiguousarray(mr, dtype="float32"),
device=dev).unsqueeze(0).unsqueeze(0)
with torch.no_grad():
return float(fn(a, b).item())
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 scores: Mattes MI between the CT
# inverse-transformed into MR coordinates (warped_ct) and the original MR
# (ha), over the MR head; IMPACT semantic loss over the full volumes.
ha_f = ha.astype("float32", copy=False)
mi = _mi_masked(warped_ct, ha_f, head_mask_from_image(ha_f))
il = _impact_loss_value(warped_ct, 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,
"impact_loss": round(il, 4) if il 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 args.deformable:
if pipeline == "simpleitk":
return RP.register_simpleitk_deformable(ct, t1c, threads=args.threads)
if pipeline == "elastix":
return RP.register_elastix_deformable(ct, t1c, workdir=args.elastix_dir,
threads=args.threads)
if pipeline == "antspyx":
return RP.register_antspyx_deformable(ct, t1c)
if pipeline == "synthmorph":
return RP.register_synthmorph_deformable(ct, t1c,
gpu=args.synthmorph_gpu,
threads=args.threads,
extent=args.synthmorph_extent)
if pipeline == "impact":
return RP.register_impact_deformable(ct, t1c, gpu=args.impact_gpu,
iters=args.impact_iters or RP.IMPACT_DEFAULT_ITERS,
samples=args.impact_samples or RP.IMPACT_DEFAULT_SAMPLES)
if pipeline == "niftyreg":
return RP.register_niftyreg_deformable(ct, t1c, threads=args.threads,
backend=args.niftyreg_backend,
gpu=args.niftyreg_gpu)
raise KeyError(pipeline)
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, gpu=args.impact_gpu,
iters=args.impact_iters or RP.IMPACT_DEFAULT_ITERS,
samples=args.impact_samples or RP.IMPACT_DEFAULT_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,
"impact_loss": 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"impact_loss={r.get('impact_loss')} {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]
ilv = [r["impact_loss"] for r in registered if r.get("impact_loss") 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,
"impact_loss": {"mean": round(float(np.mean(ilv)), 4),
"median": round(float(np.median(ilv)), 4)} if ilv 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" + "=" * 96)
hdr = (f"{'pipeline':10s} {'avail':6s} {'toolkit':22s} {'n':>3s} {'ok':>3s} {'err':>4s} "
f"{'s/reg':>7s} {'ImpactLoss':>10s} {'MI':>7s}")
print(hdr)
print("-" * 96)
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} {'-':>10s} {'-':>7s} [{r['reason']}]")
continue
rt = (r.get("reg_wall_s") or {}).get("median")
il = (r.get("impact_loss") or {}).get("median")
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(il, '{:.3f}'):>10s} {_fmt(mi, '{:.3f}'):>7s}")
print("=" * 96)
def write_plot(rows, out_png, label="rigid"):
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]
il = [(ok_rows[p].get("impact_loss") 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, 3, figsize=(14, 4))
colors = ["#6ACC65", "#F5A623", "#4878CF"]
for ax, (title, vals, fmt, ylab) in zip(axes, [
("Registration time (s, median)", regt, "{:.0f}", "s"),
("ImpactLoss: warped-CT ~ MR (semantic, median; lower = better)", il, "{:.3f}", "loss"),
("Mattes MI: warped-CT ~ MR (nats, median)", mi, "{:.3f}", "nats")]):
# wide time ranges (deformable: simpleitk ~100x the others) -> log axis
nz = [v for v in vals if v > 0]
if len(nz) >= 2 and max(nz) / min(nz) > 30:
ax.set_yscale("log")
bars = ax.bar(range(len(vals)), vals, color=colors[0])
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(f"CT->T1c registration benchmark (M6, cross-modality {label})", 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)
deformable = (os.path.basename(os.path.normpath(args.out_dir))
== "bench_registration_deformable")
if deformable:
protocol = ("deformable CT->T1c; cross-modality; each pipeline runs its "
"toolkit's native deformable workflow (linear init + deformable "
"stage, or a single deformable stage where that is native); the "
"result is rasterised to the T1c-grid displacement field of the "
"unified contract (T1c->CT); alignment scored by achieved Mattes MI "
"(nats, warped-CT ~ MR over the MR head, higher = better) and the "
"IMPACT semantic loss (vboussot/ImpactLoss, TS/M730 features, "
"warped-CT ~ MR over the full volumes, lower = better); reg time "
"excludes I/O+scoring; per-case JSON retains the production "
"skull-ring gap / acceptance QA")
else:
protocol = ("rigid CT->T1c; cross-modality; alignment scored by achieved "
"Mattes MI (nats, warped-CT ~ MR over the MR head, higher = "
"better) and the IMPACT semantic loss (vboussot/ImpactLoss, "
"TS/M730 features, warped-CT ~ MR over the full volumes, lower "
"= 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")
summary = {
"timestamp": datetime.now().isoformat(timespec="seconds"),
"protocol": protocol,
"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"),
label="deformable" if deformable else "rigid")
# markdown table for the report
md = ["| pipeline | available | toolkit | cases | ok | err | s/reg (median) | ImpactLoss (median) | 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('impact_loss') or {}).get('median')} | "
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("--deformable", action="store_true",
help="run the deformable variants of the pipelines "
"(each toolkit's native deformable workflow; results go "
"to results/bench_registration_deformable by default)")
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-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="IMPACT MaximumNumberOfIterations, one value for all "
"levels (default: the reference map's per-level budget "
f"{RP.IMPACT_DEFAULT_ITERS})")
ap.add_argument("--impact-samples", type=int, default=None,
help="IMPACT NumberOfSpatialSamples per iteration (default: "
f"{RP.IMPACT_DEFAULT_SAMPLES})")
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.deformable:
if args.out_dir == os.path.join(ROOT, "results", "bench_registration"):
args.out_dir = os.path.join(ROOT, "results", "bench_registration_deformable")
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()
Reference in a new issue View git blame Copy permalink