"""Benchmark the 5 rigid CT->T1c registration pipelines on N M6 test cases. Pipelines (see src/registration_pipelines.py for the shared transform contract): simpleitk | elastix | antspyx | deepreg | impact Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI (cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto the T1c native grid. Every pipeline is scored with the SAME crop-independent geometric QA used in production (scripts/02_build_m6_dataset.py): - skull_ring_gap_mm : concentricity of the warped-CT skull ring vs the MR head silhouette (mid/low level + angular std). Aligned = 0 <= mid <= 20 mm and std <= 12 mm. - tv_in_ct_frac : fraction of the (dilated) warped GTV that falls inside the warped-CT head tissue (the slab must straddle it). - tv_nvox : warped GTV volume (>0.5 voxels). - accepted : production rule (nvox>0 and tv_in_ct and concentric). - reg_wall_s : wall-clock registration time (excludes I/O + scoring). Because some toolkits are not pip-installable into the main `longitudinal-313` env (antspyx is kept in an isolated env, see REGISTRATION_BENCHMARK.md; PyPI deepreg is a broken stub), the benchmark is env-portable: run it under whichever conda env has a given toolkit installed, and it records only the pipelines that are importable there. elastix runs in the main env through the ITK-Elastix python API (pip ``itk`` + ``itk-elastix``). ``impact`` runs in the ``vboussot/elastix_impact`` Docker image (needs Docker + the image + a model, see REGISTRATION_BENCHMARK.md). The `--merge` mode then combines the per-pipeline JSON files (accumulated across envs) into one summary + plot. Usage (repo root): # main env (SimpleITK + elastix; impact needs Docker + the image; deepreg unavailable) source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 python scripts/bench_registration.py --pipelines simpleitk elastix impact # isolated env that has antspyx (see REGISTRATION_BENCHMARK.md) source /opt/conda/etc/profile.d/conda.sh && conda activate python scripts/bench_registration.py --pipelines antspyx simpleitk # combine everything produced so far into one table + plot python scripts/bench_registration.py --merge """ import argparse import importlib.util import json import os import sys import time from datetime import datetime from concurrent.futures import ProcessPoolExecutor, as_completed ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) import numpy as np import SimpleITK as sitk from scipy import ndimage from src.common import DATA, d, read_nii_img, 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() # --------------------------------------------------------------------------- # # 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 0.0 <= gap["mid_mm"] <= 20.0 and (gap["std_mm"] or 0) <= 12.0 accepted = nvox > 0 and tv_in_ct_frac >= 0.5 and ok_gap return { "skull_ring_gap_mm": gap, "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 == "deepreg": return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir, checkpoint=args.deepreg_ckpt) if pipeline == "impact": return RP.register_impact(ct, t1c, model=args.impact_model, gpu=args.impact_gpu, iters=args.impact_iters, samples=args.impact_samples) raise KeyError(pipeline) def _run_case(pipeline, case, args): """Worker body: read once + register + score one case. Runs in a forked process. The CT-completeness guard is applied here (on the already-read CT) so each file is read exactly once.""" key = case["key"] ct = read_nii_img(case["ct"]) t1c = read_nii_img(case["img"]) tv = read_nii_img(case["label"]) out = {"key": key, "subject": case["subject"], "date": case["date"]} if not ct_complete(ct): out.update({"skull_ring_gap_mm": None, "tv_nvox": 0, "tv_in_ct_frac": None, "alignment": "ct_incomplete", "accepted": False, "note": "CT head incomplete (data issue, not a pipeline failure)"}) return out t0 = time.time() T, meta = _call(pipeline, ct, t1c, tv, args) out["reg_wall_s"] = round(time.time() - t0, 1) out.update(meta) out.update(evaluate(ct, t1c, tv, T)) 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 def run_pipeline(pipeline, cases, args): avail, note = RP.availability().get(pipeline, (False, "unknown")) if not avail: row = {"pipeline": pipeline, "available": False, "reason": note, "n_cases": len(cases), "cases": []} print(f"[{pipeline}] UNAVAILABLE: {note}") return row print(f"[{pipeline}] available ({note}); registering {len(cases)} cases " f"with {args.workers} workers") results = [] errors = [] with ProcessPoolExecutor(max_workers=args.workers) as ex: futs = {ex.submit(_run_case, pipeline, c, args): c for c in cases} for fu in as_completed(futs): c = futs[fu] try: results.append(fu.result()) r = results[-1] g = r.get("skull_ring_gap_mm") reg = f"reg={r.get('reg_wall_s')}s" if r.get("reg_wall_s") is not None else "incomplete-CT" print(f" {c['key']}: accepted={r['accepted']} " f"gap={g['mid_mm'] if g else 'None'}mm " f"std={g['std_mm'] if g else 'None'}mm " f"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" + "=" * 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} {'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} {'-':>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} {_fmt(mi, '{:.3f}'):>7s}") print("=" * 100) def write_plot(rows, out_png): try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt except Exception as e: # noqa print("plot skipped:", repr(e)) return ok_rows = {r["pipeline"]: r for r in rows if r.get("available")} if not ok_rows: print("no available pipelines to plot") return labels = list(ok_rows) gap_mid = [ok_rows[p].get("skull_ring_gap_mid_mm", {}).get("median") or 0 for p in labels] gap_std = [ok_rows[p].get("skull_ring_gap_std_mm", {}).get("median") or 0 for p in labels] regt = [ok_rows[p].get("reg_wall_s", {}).get("median") or 0 for p in labels] acc = [ok_rows[p].get("accepted_rate") or 0 for p in labels] 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), ("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()), ha="center", va="bottom", fontsize=8) ax.set_title(title, fontsize=9) ax.set_ylabel(ylab) ax.set_xticks(range(len(vals))) ax.set_xticklabels(labels, fontsize=9) plt.suptitle("CT->T1c registration benchmark (M6, cross-modality rigid)", fontsize=11) fig.tight_layout() fig.savefig(out_png, dpi=120) plt.close(fig) print("plot:", out_png) def merge(args): out_dir = os.path.join(args.out_dir) rows = [] for p in RP.ALL_PIPELINES: fp = os.path.join(out_dir, p + ".json") if os.path.exists(fp): with open(fp) as f: rows.append(json.load(f)) if not rows: raise SystemExit(f"no pipeline results found in {out_dir}") print(f"[merge] {len(rows)} pipeline result files in {out_dir}") print_table(rows) summary = { "timestamp": datetime.now().isoformat(timespec="seconds"), "protocol": ("rigid CT->T1c; cross-modality; scored by production skull-ring " "concentricity (aligned: 0<=mid<=20mm, std<=12mm) + GTV-in-CT fraction " "(>=0.5); reg time excludes I/O+scoring; per-pipeline best-of-T vs T^-1 " "orientation gate for elastix/antspyx"), "n_cases": next((r["n_cases"] for r in rows if r.get("available")), None), "case_keys": sorted({c["key"] for r in rows if r.get("available") for c in r.get("cases", [])}), "results": rows, } sf = os.path.join(out_dir, "summary.json") with open(sf, "w") as f: json.dump(summary, f, indent=1) print("summary:", sf) write_plot(rows, os.path.join(out_dir, "bench_registration.png")) # markdown table for the report md = ["| pipeline | available | toolkit | cases | ok | err | accepted | gap mid (mm) | gap std (mm) | s/reg (median) | tv-in-CT | 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']}) |") continue md.append( f"| {p} | yes | {(r.get('toolkit') or '')} | {r['n_cases']} | {r.get('n_registered', 0)} | " f"{r['n_error']} | {r['accepted_rate']:.2f} | " f"{(r.get('skull_ring_gap_mid_mm') or {}).get('median')} | " f"{(r.get('skull_ring_gap_std_mm') or {}).get('median')} | " f"{(r.get('reg_wall_s') or {}).get('median')} | " f"{(r.get('tv_in_ct_frac') or {}).get('mean')} | " 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("--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("--impact-model", default=None, help="IMPACT TorchScript model, path relative to the HF " f"layout (default {RP.IMPACT_DEFAULT_MODEL}) or an absolute host path") 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=400, help="IMPACT per-resolution max iterations (default 400, a converged " "rigid budget; the authors' full recommended is 500)") ap.add_argument("--impact-samples", type=int, default=1500, help="IMPACT random spatial samples per iteration (default 1500; " "full recommended 2000)") 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 def main(): args = parse_args() if args.merge: merge(args) return args.elastix_dir = args.elastix_dir or d("data/m6reg/elastix") args.deepreg_dir = args.deepreg_dir or d("data/m6reg/deepreg") 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()