feat(benchmark): add Mattes MI metric and visual QA for registration evaluation

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 |
+| 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.3 | 0.93 |
+| **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.3 / 15.4 | 18.2 / 21.7 | 13.4 | 0.93 |
+| **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 | — | — | — | — | — | — | — |
+| elastix | — | — | — | — | — | — | — | — |
-| deepreg | — | — | — | — | — | — | — |
+| 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 <dir>`.
+`--deepreg-ckpt <.ckpt>` (DeepReg), `--out-dir <dir>`,
+`--qa-images` / `--qa-out <dir>` (visual QA, below).
+
+## Visual QA images (checkerboard + color overlay)
+
+`--qa-images` writes a `checkreg`-style PNG per (pipeline, case) to
+`<out-dir>/qa/<pipeline>__<case>.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
 

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))
+    mi = [ok_rows[p].get("mi_ct_mr", {}).get("median") or 0 for p in labels]
-    colors = ["#4878CF", "#D64F5C", "#6ACC65", "#9467BD"]
+    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 <out-dir>/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: