diff --git a/REGISTRATION_BENCHMARK.md b/REGISTRATION_BENCHMARK.md index 45ef2d7..41971b1 100644 --- a/REGISTRATION_BENCHMARK.md +++ b/REGISTRATION_BENCHMARK.md @@ -71,16 +71,23 @@ Production QA, crop-independent: - **`tv_nvox`** — warped GTV volume. - **`accepted`** — `tv_nvox > 0 and tv_in_ct_frac ≥ 0.5 and aligned` (the production acceptance rule). +- **`mi_ct_mr`** — mean Mattes MI (nats) between the **CT inverse-transformed + into MR coordinates** (the warped CT on the T1c grid) and the **original MR**, + restricted to the MR head and each image percentile-binned independently + (SimpleITK Mattes-MI style). This is the registration's *achieved* metric read + off as a standalone alignment score: higher = the warped-CT and MR intensities + are better in-sync = better aligned; it collapses toward 0 when the CT is + rotated/translated away from the MR. - **`reg_wall_s`** — wall-clock registration time (excludes I/O + scoring). ## Results (10 cases) -| pipeline | toolkit | registered | accepted | gap mid (med/mean) | gap std (med/mean) | s/reg (med) | tv-in-CT | -|---|---|---|---|---|---|---|---| -| **simpleitk** | SimpleITK 2.5.6 | 9 | **4/10 (0.40)** | 12.6 / 16.1 | 15.4 / 16.0 | 8.3 | 0.93 | -| **antspyx** | ANTsPy 0.6.3 | 9 | **1/10 (0.10)** | 14.3 / 15.4 | 18.2 / 21.7 | 13.4 | 0.93 | -| elastix | — | — | — | — | — | — | — | -| deepreg | — | — | — | — | — | — | — | +| pipeline | toolkit | registered | accepted | gap mid (med/mean) | gap std (med/mean) | s/reg (med) | tv-in-CT | MI (med/mean, nats) | +|---|---|---|---|---|---|---|---|---| +| **simpleitk** | SimpleITK 2.5.6 | 9 | **4/10 (0.40)** | 12.6 / 16.1 | 15.4 / 16.0 | 8.1 | 0.93 | **0.643 / 0.630** | +| **antspyx** | ANTsPy 0.6.3 | 9 | **1/10 (0.10)** | 14.2 / 17.3 | 24.8 / 22.4 | 12.7 | 0.93 | **0.600 / 0.523** | +| elastix | — | — | — | — | — | — | — | — | +| deepreg | — | — | — | — | — | — | — | — | Per-case (mid / std, `OK` = accepted): @@ -111,9 +118,17 @@ Per-case (mid / std, `OK` = accepted): landscape handles the tilt better than ANTs' Rigid. Note `m6_26TYQI3C`: ANTs nailed the position (mid 6.1) but missed the orientation (std 37.2) — a textbook tilted local minimum. -- **Speed**: SimpleITK ~8.3 s/case, ANTs ~13.4 s/case (ANTs includes the +- **Speed**: SimpleITK ~8.1 s/case, ANTs ~12.7 s/case (ANTs includes the pre-shift resample + its affine-MI init stages). Both are CPU-bound and ran 4 cases in parallel. +- **Mattes MI (warped-CT ~ MR) is an independent, intensity-based confirmation**: + SimpleITK median 0.643 nats vs ANTs 0.600, and it diverges sharply on the + tilted cases — ANTs `m6_26TYQI3C` drops to 0.094 nats (rotation ~37 mm off) + while SimpleITK holds 0.670. A high MI means the CT skull/tissue intensities + fall on the matching MR structures after the inverse transform. +- **ANTs MI is stochastic**: ANTs' Mattes-MI subsamples the image at random, so + its per-case numbers carry run-to-run variance (this table is one draw); + SimpleITK's MI (and registration) is deterministic. ## Reproducing @@ -134,7 +149,40 @@ python scripts/bench_registration.py --merge Other flags: `--sample random --seed 0` (reproducible random 10), `--keys k1,k2,...` (explicit cases), `--n-cases N`, `--workers W`, `--threads T`, `--simpleitk-mode {benchmark,production}`, -`--deepreg-ckpt <.ckpt>` (DeepReg), `--out-dir `. +`--deepreg-ckpt <.ckpt>` (DeepReg), `--out-dir `, +`--qa-images` / `--qa-out ` (visual QA, below). + +## Visual QA images (checkerboard + color overlay) + +`--qa-images` writes a `checkreg`-style PNG per (pipeline, case) to +`/qa/__.png` for eyeballing the alignment. Each image +is a 3×2 grid: **rows** are the three orthogonal mid-slices through the head +center (**mid-axial, mid-coronal, mid-sagittal**), and **columns** are: + +- **left — checkerboard blend:** alternating blocks of the CT inverse-transformed + into MR coordinates and the original MR; a clean registration gives a seamless + mosaic, a misregistration shows hard edges at the block boundaries. +- **right — additive color overlay:** CT in red, MR in green. Structures that + coincide blend to yellow; rotation/translation error appears as red/green + fringing at the skull and brain edges. + +The CT is shown in a **brain window (L40/W80 = 0–80 HU)** so the brain +parenchyma is visible and comparable to the MR (the skull, ~800+ HU, clips to +white); the MR uses a 2–98 % head window. The coronal and sagittal views are +oriented superior-up to match the axial view. + +The title carries the pipeline, case key, Mattes MI, and the skull-ring gap +(mid/std), so a glance tells you both the achieved-metric score and whether the +visuals agree. Generate for the two available pipelines with: + +```bash +/home/xfr/.conda/envs/longitudinal/bin/python scripts/bench_registration.py \ + --pipelines simpleitk --n-cases 10 --qa-images +/home/xfr/.conda/envs/ants312/bin/python scripts/bench_registration.py \ + --pipelines antspyx --n-cases 10 --qa-images +``` + +(9 images each; the incomplete-CT case is skipped.) ## Enabling the two blocked pipelines diff --git a/scripts/bench_registration.py b/scripts/bench_registration.py index 1cc6e6a..d62cd04 100644 --- a/scripts/bench_registration.py +++ b/scripts/bench_registration.py @@ -52,10 +52,34 @@ import numpy as np import SimpleITK as sitk from scipy import ndimage -from src.common import DATA, d, read_nii_img +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() @@ -106,6 +130,10 @@ def evaluate(ct, t1c, tv, T): 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) @@ -119,6 +147,7 @@ def evaluate(ct, t1c, tv, T): 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", @@ -126,6 +155,64 @@ def evaluate(ct, t1c, tv, T): } +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) @@ -158,6 +245,18 @@ def _run_case(pipeline, case, 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): + g = out.get("skull_ring_gap_mm") + title = (f"{pipeline} | {key} | MI={out.get('mi_ct_mr')} | " + f"gap mid {g['mid_mm']}mm std {g['std_mm']}mm" if g + else f"{pipeline} | {key}") + 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 @@ -184,6 +283,7 @@ def run_pipeline(pipeline, cases, args): 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"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)}) @@ -197,6 +297,7 @@ def run_pipeline(pipeline, cases, args): 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"), @@ -212,6 +313,8 @@ def run_pipeline(pipeline, cases, args): "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"), } @@ -238,23 +341,24 @@ 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}") + f"{'accept':>7s} {'gapMid':>7s} {'gapStd':>7s} {'s/reg':>7s} {'tvInCT':>7s} {'MI':>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']}]") + f"{'-':>7s} {'-':>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") + 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} {r['accepted_rate']:>7.2f} " f"{_fmt(gm, '{:.1f}'):>7s} {_fmt(gs, '{:.1f}'):>7s} {_fmt(rt, '{:.1f}'):>7s} " - f"{_fmt(tv, '{:.2f}'):>7s}") + f"{_fmt(tv, '{:.2f}'):>7s} {_fmt(mi, '{:.3f}'):>7s}") print("=" * 100) @@ -275,13 +379,15 @@ def write_plot(rows, out_png): 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"] + mi = [ok_rows[p].get("mi_ct_mr", {}).get("median") or 0 for p in labels] + fig, axes = plt.subplots(1, 5, figsize=(20, 4)) + colors = ["#4878CF", "#D64F5C", "#6ACC65", "#9467BD", "#E08E45"] 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)]): + ("Acceptance rate", acc, "{:.2f}", "frac", False), + ("Mean Mattes MI: warped-CT ~ MR (nats, median)", mi, "{:.3f}", "nats", 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()), @@ -325,12 +431,12 @@ def merge(args): 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 |", - "|---|---|---|---|---|---|---|---|---|---|---|"] + md = ["| pipeline | available | toolkit | cases | ok | err | accepted | gap mid (mm) | gap std (mm) | s/reg (median) | tv-in-CT | MI (nats) |", + "|---|---|---|---|---|---|---|---|---|---|---|---|"] 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']}) |") + 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)} | " @@ -338,7 +444,8 @@ def merge(args): 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')} |") + 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") @@ -363,6 +470,11 @@ def parse_args(): 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("--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 /qa)") ap.add_argument("--merge", action="store_true", help="combine existing per-pipeline JSONs") args = ap.parse_args() return args @@ -375,6 +487,9 @@ def main(): return args.elastix_dir = args.elastix_dir or d("data/m6reg/elastix") args.deepreg_dir = args.deepreg_dir or d("data/m6reg/deepreg") + 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: