439 lines
No EOL
21 KiB
Python
439 lines
No EOL
21 KiB
Python
"""Benchmark the 6 DEFORMABLE CT->T1c registration pipelines on N M6 cases.
|
|
|
|
Companion to ``scripts/bench_registration.py`` (rigid). Pipelines:
|
|
simpleitk | elastix | antspyx | synthmorph | impact | niftyreg
|
|
(each = its toolkit's native deformable refinement initialized from the
|
|
toolkit's own rigid result; see src/registration_pipelines_deformable.py
|
|
and REGISTRATION_BENCHMARK_DEFORMABLE.md).
|
|
|
|
Headline scores (identical for every pipeline):
|
|
- mi_ct_mr : achieved Mattes MI (nats), warped-CT ~ MR over the MR head
|
|
(higher = better). Also computed for the pipeline's own
|
|
rigid base (mi_base) -> d_mi = the deformable gain.
|
|
- impact_loss : IMPACT semantic loss (TS/M730 features), warped-CT ~ MR
|
|
over the full volumes (lower = better). Also mi/IL for
|
|
the base (il_base).
|
|
- jac_* : Jacobian-determinant statistics of the FINAL deformable
|
|
transform over the MR head (finite differences on the
|
|
T1c grid): p1/p50/mean + fraction < 0 (folded), < 0.3,
|
|
< 0.5. A healthy diffeomorphic refinement has J ~ 1.
|
|
- jac_gtv_* : same Jacobian statistics restricted to the warped-GTV
|
|
region (does the deformation fold the tumor?).
|
|
- disp_base_* : |T_final(x) - T_base(x)| over the MR head (p50/p95/max,
|
|
mm) -- the deformable increment.
|
|
- tv_vol_ratio : warped-GTV volume final / base (a diffeomorphism
|
|
preserves volume, ~1).
|
|
- reg_wall_s : wall-clock of the deformable stage (excludes I/O, the
|
|
rigid base and scoring).
|
|
Production QA (skull-ring gap, accepted) is still computed per case and kept
|
|
in the JSON.
|
|
|
|
Usage (repo root):
|
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
|
python scripts/bench_registration_deformable.py --pipelines "simpleitk elastix antspyx niftyreg"
|
|
python scripts/bench_registration_deformable.py --pipelines "impact synthmorph"
|
|
python scripts/bench_registration_deformable.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 src.common import DATA, d, read_nii_img, head_mask_from_image
|
|
from src import registration_pipelines as RP
|
|
from src import registration_pipelines_deformable as RPD
|
|
|
|
_spec = importlib.util.spec_from_file_location("br", os.path.join(ROOT, "scripts", "bench_registration.py"))
|
|
BR = importlib.util.module_from_spec(_spec)
|
|
_spec.loader.exec_module(BR)
|
|
|
|
JAC_STEP_MM = 1.0 # finite-difference half-step for the Jacobian estimate
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# deformable-specific scoring
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def _dense_H(T, ref):
|
|
"""Absolute mapping H(x) = T(x) at every voxel of ``ref`` (z,y,x,3)."""
|
|
return RPD.densify_sitk(T, ref)[2]
|
|
|
|
|
|
def jacobian_of_T(T, ref, step_mm=JAC_STEP_MM):
|
|
"""det dT/dx of an arbitrary contract transform over ``ref``'s grid.
|
|
|
|
Central finite differences with world-axis steps of ~``step_mm`` on the
|
|
densified absolute mapping (linear resampling of coordinate volumes);
|
|
border voxels (no both-side neighbour) are NaN. Returns J (z,y,x)."""
|
|
H = RPD.densify_sitk(T, ref)[2]
|
|
sp = np.array(ref.GetSpacing()) # (x, y, z); array axis i has spacing sp[2-i]
|
|
off = np.array([int(np.clip(np.round(step_mm / sp[2 - i]), 1, None))
|
|
for i in range(3)])
|
|
H = np.ascontiguousarray(H, dtype='float64')
|
|
dH = []
|
|
for i in range(3):
|
|
d = np.full(H.shape, np.nan)
|
|
sl_p = [slice(None)] * 3; sl_m = [slice(None)] * 3; sl_mid = [slice(None)] * 3
|
|
sl_p[i] = slice(2 * off[i], None)
|
|
sl_m[i] = slice(None, -2 * off[i])
|
|
sl_mid[i] = slice(off[i], -off[i])
|
|
d[tuple(sl_mid)] = (H[tuple(sl_p)] - H[tuple(sl_m)]) / (2.0 * off[i] * sp[2 - i])
|
|
dH.append(d)
|
|
# array axis i <-> world axis (2-i): axis 2 = x, axis 1 = y, axis 0 = z.
|
|
# J = det[dH/dx, dH/dy, dH/dz]
|
|
A = dH[2]; B = dH[1]; C = dH[0]
|
|
J = (A[..., 0] * (B[..., 1] * C[..., 2] - B[..., 2] * C[..., 1])
|
|
- A[..., 1] * (B[..., 0] * C[..., 2] - B[..., 2] * C[..., 0])
|
|
+ A[..., 2] * (B[..., 0] * C[..., 1] - B[..., 1] * C[..., 0]))
|
|
return J
|
|
|
|
|
|
def _stats(J, mask):
|
|
v = J[mask]
|
|
v = v[np.isfinite(v)]
|
|
if v.size < 50:
|
|
return None
|
|
return {
|
|
"p1": round(float(np.percentile(v, 1)), 4),
|
|
"p50": round(float(np.median(v)), 4),
|
|
"mean": round(float(v.mean()), 4),
|
|
"frac_neg": round(float((v < 0).mean()), 5),
|
|
"frac_lt03": round(float((v < 0.3).mean()), 5),
|
|
"frac_lt05": round(float((v < 0.5).mean()), 5),
|
|
"nvox": int(v.size),
|
|
}
|
|
|
|
|
|
def deformable_metrics(ct, t1c, tv, T_final, T_base):
|
|
"""Jacobian / displacement-increment / GTV-preservation metrics of a
|
|
deformable result (final transform vs its rigid base)."""
|
|
ha = sitk.GetArrayFromImage(t1c).astype('float32')
|
|
hm = head_mask_from_image(ha)
|
|
J = jacobian_of_T(T_final, t1c)
|
|
jac_head = _stats(J, hm)
|
|
# warped GTV mask (final)
|
|
wtv = sitk.GetArrayFromImage(
|
|
sitk.Resample(tv, t1c, T_final, sitk.sitkNearestNeighbor, 0.0))
|
|
gtv = wtv > 0.5
|
|
if gtv.sum() > 50:
|
|
gtv = gtv & hm
|
|
jac_gtv = _stats(J, gtv)
|
|
else:
|
|
jac_gtv = None
|
|
# displacement increment |T_final - T_base| over the head
|
|
Hf = _dense_H(T_final, t1c)
|
|
Hb = _dense_H(T_base, t1c)
|
|
inc = np.linalg.norm(Hf - Hb, axis=-1)
|
|
inc_stats = {
|
|
"p50": round(float(np.median(inc[hm])), 2),
|
|
"p95": round(float(np.percentile(inc[hm], 95)), 2),
|
|
"max": round(float(inc[hm].max()), 2),
|
|
}
|
|
# GTV volume preservation (final vs base)
|
|
wtv_b = sitk.GetArrayFromImage(
|
|
sitk.Resample(tv, t1c, T_base, sitk.sitkNearestNeighbor, 0.0))
|
|
nv_f = int((wtv > 0.5).sum()); nv_b = int((wtv_b > 0.5).sum())
|
|
tv_vol_ratio = round(nv_f / nv_b, 3) if nv_b > 0 else None
|
|
return {"jac_head": jac_head, "jac_gtv": jac_gtv,
|
|
"disp_base_mm": inc_stats, "tv_vol_ratio": tv_vol_ratio,
|
|
"gtv_nvox_final": nv_f, "gtv_nvox_base": nv_b}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# case worker
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def _call(pipeline, ct, t1c, args):
|
|
if pipeline == "simpleitk":
|
|
return RPD.register_simpleitk_deformable(ct, t1c, threads=args.threads)
|
|
if pipeline == "elastix":
|
|
return RPD.register_elastix_deformable(ct, t1c, workdir=args.elastix_dir,
|
|
threads=args.threads)
|
|
if pipeline == "antspyx":
|
|
return RPD.register_antspyx_deformable(ct, t1c, threads=args.threads)
|
|
if pipeline == "niftyreg":
|
|
return RPD.register_niftyreg_deformable(ct, t1c, threads=args.threads,
|
|
gpu=args.niftyreg_gpu,
|
|
backend=args.niftyreg_backend)
|
|
if pipeline == "synthmorph":
|
|
return RPD.register_synthmorph_deformable(ct, t1c, gpu=args.synthmorph_gpu,
|
|
extent=args.synthmorph_extent)
|
|
if pipeline == "impact":
|
|
return RPD.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,
|
|
threads=args.threads)
|
|
raise KeyError(pipeline)
|
|
|
|
|
|
def _run_case(pipeline, case, args):
|
|
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 BR.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
|
|
t0 = time.time()
|
|
T_final, T_base, meta = _call(pipeline, ct, t1c, args)
|
|
out["reg_wall_s"] = round(time.time() - t0, 1)
|
|
out.update(meta)
|
|
# final + base alignment scores
|
|
out.update(BR.evaluate(ct, t1c, tv, T_final))
|
|
out["base"] = BR.evaluate(ct, t1c, tv, T_base)
|
|
out["d_mi"] = (round(out["mi_ct_mr"] - out["base"]["mi_ct_mr"], 4)
|
|
if out["mi_ct_mr"] is not None and out["base"]["mi_ct_mr"] is not None
|
|
else None)
|
|
out["d_il"] = (round(out["impact_loss"] - out["base"]["impact_loss"], 4)
|
|
if out["impact_loss"] is not None and out["base"]["impact_loss"] is not None
|
|
else None)
|
|
out.update(deformable_metrics(ct, t1c, tv, T_final, T_base))
|
|
return out
|
|
|
|
|
|
def run_pipeline(pipeline, cases, args):
|
|
avail, note = RPD.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')} (base {r.get('base', {}).get('mi_ct_mr')}) "
|
|
f"il={r.get('impact_loss')} "
|
|
f"jac_p1={(r.get('jac_head') or {}).get('p1')} {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"])
|
|
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.get("accepted")]
|
|
|
|
def _med(vals, nd=4):
|
|
vals = [v for v in vals if v is not None]
|
|
return round(float(np.median(vals)), nd) if vals else None
|
|
|
|
row = {
|
|
"pipeline": pipeline, "available": True,
|
|
"toolkit": (registered or [{}])[0].get("toolkit"),
|
|
"metric": (registered or [{}])[0].get("metric"),
|
|
"transform": (registered or [{}])[0].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": {
|
|
"median": _med([r["reg_wall_s"] for r in registered], 1),
|
|
"mean": round(float(np.mean([r["reg_wall_s"] for r in registered])), 1) if registered else None,
|
|
"max": round(float(np.max([r["reg_wall_s"] for r in registered])), 1) if registered else None,
|
|
},
|
|
"mi_ct_mr": {"median": _med([r.get("mi_ct_mr") for r in registered]),
|
|
"mean": round(float(np.mean([r["mi_ct_mr"] for r in registered if r.get("mi_ct_mr") is not None])), 4) if any(r.get("mi_ct_mr") is not None for r in registered) else None},
|
|
"mi_base": {"median": _med([r.get("base", {}).get("mi_ct_mr") for r in registered])},
|
|
"d_mi": {"median": _med([r.get("d_mi") for r in registered])},
|
|
"impact_loss": {"median": _med([r.get("impact_loss") for r in registered]),
|
|
"mean": round(float(np.mean([r["impact_loss"] for r in registered if r.get("impact_loss") is not None])), 4) if any(r.get("impact_loss") is not None for r in registered) else None},
|
|
"il_base": {"median": _med([r.get("base", {}).get("impact_loss") for r in registered])},
|
|
"d_il": {"median": _med([r.get("d_il") for r in registered])},
|
|
"jac_p1": {"median": _med([(r.get("jac_head") or {}).get("p1") for r in registered])},
|
|
"jac_frac_neg": {"median": _med([(r.get("jac_head") or {}).get("frac_neg") for r in registered], 5)},
|
|
"jac_frac_lt05": {"median": _med([(r.get("jac_head") or {}).get("frac_lt05") for r in registered], 5)},
|
|
"jac_gtv_p1": {"median": _med([(r.get("jac_gtv") or {}).get("p1") for r in registered])},
|
|
"disp_base_p95_mm": {"median": _med([(r.get("disp_base_mm") or {}).get("p95") for r in registered], 2)},
|
|
"tv_vol_ratio": {"median": _med([r.get("tv_vol_ratio") for r in registered])},
|
|
"cases": results, "errors": errors,
|
|
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
|
}
|
|
return row
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# reporting
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def _fmt(v, spec="{:g}", default="-"):
|
|
return default if v is None else spec.format(v)
|
|
|
|
|
|
def print_table(rows):
|
|
order = [p for p in RPD.ALL_PIPELINES if any(r["pipeline"] == p for r in rows)]
|
|
print("\n" + "=" * 118)
|
|
print(f"{'pipeline':10s} {'avail':5s} {'n':>3s} {'ok':>3s} {'err':>3s} {'s/reg':>7s} "
|
|
f"{'MI fin':>7s} {'MI base':>8s} {'dMI':>7s} {'IL fin':>7s} {'IL base':>8s} "
|
|
f"{'jacP1':>7s} {'J<0':>7s} {'J<.5':>6s} {'incr p95':>9s}")
|
|
print("-" * 118)
|
|
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':5s} {r['n_cases']:>3d} {0:>3d} {'-':>3s} [{r['reason'][:60]}]")
|
|
continue
|
|
print(f"{p:10s} {'yes':5s} {r['n_cases']:>3d} {r.get('n_registered', 0):>3d} "
|
|
f"{r.get('n_error', 0):>3d} {_fmt((r.get('reg_wall_s') or {}).get('median'), '{:.0f}'):>7s} "
|
|
f"{_fmt((r.get('mi_ct_mr') or {}).get('median'), '{:.3f}'):>7s} "
|
|
f"{_fmt((r.get('mi_base') or {}).get('median'), '{:.3f}'):>8s} "
|
|
f"{_fmt((r.get('d_mi') or {}).get('median'), '{:+.3f}'):>7s} "
|
|
f"{_fmt((r.get('impact_loss') or {}).get('median'), '{:.3f}'):>7s} "
|
|
f"{_fmt((r.get('il_base') or {}).get('median'), '{:.3f}'):>8s} "
|
|
f"{_fmt((r.get('jac_p1') or {}).get('median'), '{:.3f}'):>7s} "
|
|
f"{_fmt((r.get('jac_frac_neg') or {}).get('median'), '{:.4f}'):>7s} "
|
|
f"{_fmt((r.get('jac_frac_lt05') or {}).get('median'), '{:.3f}'):>6s} "
|
|
f"{_fmt((r.get('disp_base_p95_mm') or {}).get('median'), '{:.1f}'):>9s}")
|
|
print("=" * 118)
|
|
|
|
|
|
def merge(args):
|
|
out_dir = args.out_dir
|
|
rows = []
|
|
for p in RPD.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": ("deformable CT->T1c (each toolkit's native deformable refinement "
|
|
"initialized from its own rigid result); alignment scored by "
|
|
"achieved Mattes MI (nats, warped-CT ~ MR over the MR head, higher "
|
|
"= better) and the IMPACT semantic loss (TS/M730, full volumes, "
|
|
"lower = better), plus Jacobian-determinant statistics of the "
|
|
"final transform (finite differences on the T1c grid), the "
|
|
"deformable displacement increment vs the rigid base, and GTV "
|
|
"volume preservation; same contract / shared rules as the rigid "
|
|
"benchmark (full volumes, T1c fixed / CT moving, no mask)"),
|
|
"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)
|
|
md = ["| pipeline | toolkit | cases | ok | s/reg (med) | MI (med) | MI base (med) | ΔMI (med) | ImpactLoss (med) | IL base (med) | ΔIL (med) | jac p1 (med) | frac J<0 (med) | frac J<0.5 (med) | incr p95 (med, mm) | tv vol ratio (med) |",
|
|
"|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|"]
|
|
for r in rows:
|
|
if not r.get("available"):
|
|
md.append(f"| {r['pipeline']} | - | {r['n_cases']} | 0 | - | - | - | - | - | - | - | - | - | - | - | - | ({r['reason']}) |")
|
|
continue
|
|
md.append(
|
|
f"| {r['pipeline']} | {(r.get('toolkit') or '')} | {r['n_cases']} | {r.get('n_registered', 0)} | "
|
|
f"{_fmt((r.get('reg_wall_s') or {}).get('median'), '{:.0f}')} | "
|
|
f"{_fmt((r.get('mi_ct_mr') or {}).get('median'), '{:.3f}')} | "
|
|
f"{_fmt((r.get('mi_base') or {}).get('median'), '{:.3f}')} | "
|
|
f"{_fmt((r.get('d_mi') or {}).get('median'), '{:+.3f}')} | "
|
|
f"{_fmt((r.get('impact_loss') or {}).get('median'), '{:.3f}')} | "
|
|
f"{_fmt((r.get('il_base') or {}).get('median'), '{:.3f}')} | "
|
|
f"{_fmt((r.get('d_il') or {}).get('median'), '{:+.3f}')} | "
|
|
f"{_fmt((r.get('jac_p1') or {}).get('median'), '{:.3f}')} | "
|
|
f"{_fmt((r.get('jac_frac_neg') or {}).get('median'), '{:.4f}')} | "
|
|
f"{_fmt((r.get('jac_frac_lt05') or {}).get('median'), '{:.3f}')} | "
|
|
f"{_fmt((r.get('disp_base_p95_mm') or {}).get('median'), '{:.1f}')} | "
|
|
f"{_fmt((r.get('tv_vol_ratio') or {}).get('median'), '{:.3f}')} |")
|
|
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(RPD.ALL_PIPELINES))
|
|
ap.add_argument("--n-cases", type=int, default=10)
|
|
ap.add_argument("--sample", choices=["first", "random"], default="first")
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ap.add_argument("--keys", default=None)
|
|
ap.add_argument("--workers", type=int, default=4)
|
|
ap.add_argument("--threads", type=int, default=24)
|
|
ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_registration_deformable"))
|
|
ap.add_argument("--elastix-dir", default=None)
|
|
ap.add_argument("--synthmorph-gpu", type=int, default=None)
|
|
ap.add_argument("--synthmorph-extent", type=int, default=256)
|
|
ap.add_argument("--impact-gpu", type=int, default=None)
|
|
ap.add_argument("--impact-iters", type=int, default=None)
|
|
ap.add_argument("--impact-samples", type=int, default=None)
|
|
ap.add_argument("--niftyreg-backend", choices=["auto", "conda", "docker"], default="auto")
|
|
ap.add_argument("--niftyreg-gpu", type=int, default=None)
|
|
ap.add_argument("--merge", action="store_true")
|
|
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")
|
|
pipelines = [p.lower() for p in args.pipelines.split()]
|
|
unknown = set(pipelines) - set(RPD.ALL_PIPELINES)
|
|
if unknown:
|
|
raise SystemExit(f"unknown pipeline(s): {unknown}")
|
|
if args.keys:
|
|
m6 = RP._load_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 = BR.select_cases(args)
|
|
print(f"\n[bench-def] 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:
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
fp = os.path.join(args.out_dir, r["pipeline"] + ".json")
|
|
with open(fp, "w") as f:
|
|
json.dump(r, f, indent=1)
|
|
print(f"[{r['pipeline']}] wrote {fp}")
|
|
print_table(rows)
|
|
avail = RP.availability()
|
|
with open(os.path.join(args.out_dir, "_availability.json"), "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("[bench-def] done. Re-run with --merge to combine.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |