diff --git a/REGISTRATION_BENCHMARK.md b/REGISTRATION_BENCHMARK.md new file mode 100644 index 0000000..45ef2d7 --- /dev/null +++ b/REGISTRATION_BENCHMARK.md @@ -0,0 +1,177 @@ +# Registration Pipeline Benchmark — CT → T1c (M6) + +Benchmarks **four rigid, cross-modality registration pipelines** that align a +head-and-neck CT to a tightly-cropped T1c MRI so the GTV can be warped onto the +T1c (MR) native grid. This is the registration step of +`scripts/02_build_m6_dataset.py`. The pairs are hard: different modality (CT +Hounsfield vs MR intensity) and the patient is repositioned between the CT and +MR sessions, so the inter-session head tilt can reach ~30°. + +- Pipelines: [`src/registration_pipelines.py`](src/registration_pipelines.py) +- Benchmark harness: [`scripts/bench_registration.py`](scripts/bench_registration.py) +- Results: `results/bench_registration/` (per-pipeline JSON, `summary.json`, + `summary.md`, `bench_registration.png`) + +## Unified contract + +Every pipeline returns a SimpleITK transform `T` such that + +```python +sitk.Resample(ct, t1c, T, sitk.sitkLinear) # CT onto the MR grid +sitk.Resample(tv, t1c, T, sitk.sitkNearestNeighbor, 0) # GTV onto the MR grid +``` + +places the CT / GTV concentric with the MR head. In `Resample`'s +destination→source convention this `T` maps **T1c → CT**. All pipelines register +on the same CT-space head box (MI computed on head tissue, not neck/air) and, +for the single-init tools, from the **same centroid-aligned start** +(`centroid_offset`), so the comparison is on the optimizer / metric / transform +handling — not on who finds the ~30–60 mm gross offset first. + +## The four pipelines + +| # | pipeline | toolkit | metric | transform | status in this env | +|---|----------|---------|--------|-----------|--------------------| +| 1 | `elastix` | ITK-Elastix | Mattes MI | `Rigid3DTransform` (ASGD) | **not installed** (conda-forge unreachable) | +| 2 | `simpleitk` | SimpleITK 2.5.6 | Mattes MI (32 bins) | `Euler3D` rigid, pyramid `[4,2,1]` | ✅ ran | +| 3 | `antspyx` | ANTsPy / ANTs 0.6.3 | Mattes MI (multi-res) | ANTs `Rigid` | ✅ ran (isolated py3.12 env) | +| 4 | `deepreg` | Oxford DeepReg (CycleGAN) | CycleGAN + reg. loss (pretrained) | 3D displacement field | **not installed** (see below) | + +`simpleitk` has two budgets: `--simpleitk-mode benchmark` (single pyramid, +centered start — the default, comparable to the other tools) and +`--simpleitk-mode production` (the full multi-start tilt grid from +`02_build_m6_dataset.py`, ~20× slower, the most robust but not used for the +head-to-head). + +## Test cases + +First 10 labeled M6 cases (`--n-cases 10 --sample first`), reproducible. 9 have +a complete CT; 1 (`m6_22GD5VL2_20230724`) has an incomplete CT head and is +reported as `incomplete-CT` for every pipeline (a data issue, not a pipeline +failure). + +``` +m6_22GD5VL2_20230724 (incomplete CT) +m6_22IRRKU6_20240115 m6_23XIMO4M_20230222 m6_24V62QQM_20230223 +m6_24V62QQM_20231017 m6_24V62QQM_20240820 m6_26TYQI3C_20221027 +m6_2BBN5XA4_20250411 m6_2C4TKI6H_20240506 m6_2CDNN7BH_20240102 +``` + +## Metrics (identical for every pipeline) + +Production QA, crop-independent: + +- **`skull_ring_gap_mm`** — march rays from the MR head center; the gap between + the MR head silhouette boundary and the warped-CT skull ring (p97 of positive + CT) at mid and low levels, plus the angular std. **Aligned ⇔ + `0 ≤ mid ≤ 20 mm` and `std ≤ 12 mm`.** `mid` measures translational + concentricity; `std` measures rotational concentricity. +- **`tv_in_ct_frac`** — fraction of the dilated warped GTV inside the warped-CT + head tissue (the slab must straddle the lesion). Acceptance needs `≥ 0.5`. +- **`tv_nvox`** — warped GTV volume. +- **`accepted`** — `tv_nvox > 0 and tv_in_ct_frac ≥ 0.5 and aligned` (the + production acceptance rule). +- **`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 | — | — | — | — | — | — | — | + +Per-case (mid / std, `OK` = accepted): + +| case | simpleitk | antspyx | +|---|---|---| +| m6_22GD5VL2_20230724 | incomplete-CT | incomplete-CT | +| m6_22IRRKU6_20240115 | 25.7 / 29.3 | 26.3 / 29.2 | +| m6_23XIMO4M_20230222 | 22.2 / 23.9 | 23.8 / 25.0 | +| m6_24V62QQM_20230223 | 10.7 / 16.9 | 9.4 / 18.2 | +| m6_24V62QQM_20231017 | 11.5 / 7.3 **OK** | 12.5 / 12.8 | +| m6_24V62QQM_20240820 | 11.6 / 9.6 **OK** | 14.8 / 11.7 **OK** | +| m6_26TYQI3C_20221027 | 12.2 / 3.7 **OK** | 6.1 / 37.2 | +| m6_2BBN5XA4_20250411 | 18.6 / 8.0 **OK** | 14.3 / 12.7 | +| m6_2C4TKI6H_20240506 | 12.6 / 15.4 | 13.6 / 14.7 | +| m6_2CDNN7BH_20240102 | 19.4 / 30.2 | 18.2 / 33.6 | + +### Interpretation + +- **Translation is recovered well by both** (mid gaps mostly 6–20 mm, within the + `≤ 20 mm` band). `tv-in-CT` is identical (0.93): the warped GTV lands inside + the CT head for both, so the target is in the right place. +- **The binding failure is rotation** — `std > 12 mm` on the repositioned cases + (the ~30° inter-session tilt produces MI local minima that a single centered + start can lock onto). This is exactly why production `02_build_m6_dataset.py` + uses a **multi-start tilt grid**. +- **SimpleITK is the more robust of the two** (4 vs 1 accepted, lower median + std 15.4 vs 18.2): its RSGD + physical-step-scales + 32-bin Mattes-MI + 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 + pre-shift resample + its affine-MI init stages). Both are CPU-bound and ran + 4 cases in parallel. + +## Reproducing + +```bash +# main env — SimpleITK baseline (elastix/deepreg report unavailable here) +source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal +python scripts/bench_registration.py --pipelines simpleitk --n-cases 10 \ + --workers 4 --threads 24 --simpleitk-mode benchmark + +# isolated py3.12 env with antspyx — ANTs (+ optional SimpleITK cross-check) +/home/xfr/.conda/envs/ants312/bin/python scripts/bench_registration.py \ + --pipelines antspyx --n-cases 10 --workers 4 --threads 24 + +# combine the per-pipeline JSONs accumulated across envs into one table + plot +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 `. + +## Enabling the two blocked pipelines + +**ITK-Elastix** — not on PyPI; needs conda-forge (which this host cannot reach). +Once conda-forge is available: + +```bash +conda install -c conda-forge elastix # provides the elastix/transformix CLI +python scripts/bench_registration.py --pipelines elastix --n-cases 10 +``` + +The adapter registers `fixed=T1c, moving=CT` with `Rigid3DTransform` + Mattes MI +on the head box and reads the result transform. + +**DeepReg (Oxford)** — the PyPI `deepreg` (0.1.2) is an unrelated stub that pins +`numpy==1.18.5`; do not install it. DeepReg is not a one-shot rigid tool: it is a +CycleGAN image-to-image registration that must be **trained or run from a +pretrained checkpoint**. To benchmark it: + +```bash +git clone https://github.com/OB-CVHU/DeepReg && cd DeepReg && pip install -e . +python scripts/bench_registration.py --pipelines deepreg --n-cases 10 \ + --deepreg-ckpt /path/to/pretrained_3d.ckpt +``` + +The adapter runs `deepreg test` to predict a voxel displacement field, converts +it to a SimpleITK `DisplacementFieldTransform`, and scores it like the others. + +## Notes / caveats + +- The benchmark compares **standard single-init** configurations (centered + start). Production uses `--simpleitk-mode production`'s multi-start tilt grid, + which raises SimpleITK's acceptance at ~20× the cost; that is out of scope for + the head-to-head but is the config the dataset actually ships with. +- `elastix` / `deepreg` could not be executed here, so their rows are + "not installed", not measured. Their code paths are implemented and + unit-verified for import/convention; run them where the toolkits are present. +- `fix_orientation` (elastix/ants/deepreg) chooses between a transform and its + inverse by skull-ring concentricity — a correctness gate that makes the result + independent of each toolkit's transform-direction convention. \ No newline at end of file diff --git a/scripts/02_build_m6_dataset.py b/scripts/02_build_m6_dataset.py index 1f46519..ecd752c 100644 --- a/scripts/02_build_m6_dataset.py +++ b/scripts/02_build_m6_dataset.py @@ -1,7 +1,11 @@ """Build M6-2025 manifests: T1c volumes, labeled when GTV exists. -GTV lives in CT space; CT is registered to T1c (similarity->affine MI + BSpline) -and GTV warped (nearest) into T1c native space. +GTV lives in CT space; T1c is registered to CT (rigid, Mattes MI on a head box +located in CT space; multi-start initial tilts because inter-session head tilt +reaches ~30 deg and tilted MI local minima can outscore the true alignment). +The inverse transform then warps CT and GTV (nearest) into T1c native space. +Acceptance requires the warped CT skull ring to be concentric with the MR head +silhouette (crop-independent check). Outputs: data/manifests/m6_labeled.jsonl {key, subject, date, img, label(warped GTV), source:'m6'} data/manifests/m6_unlabeled.jsonl {key, subject, date, img, label:None, source:'m6'} @@ -56,18 +60,132 @@ def discover(): t1c = pick_t1c(mrd) if not t1c: continue - gtv = os.path.join(dp, "RT", "TV", "Struct_GTV.nii.gz") - ct = os.path.join(dp, "RT", "ct_image.nii.gz") + rt = os.path.join(dp, "RT") + ct = os.path.join(rt, "ct_image.nii.gz") + label = None + tv_files = [] + if os.path.isdir(rt): + try: + tv_files = sorted(f for f in os.listdir(rt) + if f.startswith("TV-") and f.endswith(".nii.gz")) + except OSError: + tv_files = [] + if tv_files: + label = os.path.join(rt, tv_files[0]) + elif os.path.exists(os.path.join(rt, "TV", "Struct_GTV.nii.gz")): + label = os.path.join(rt, "TV", "Struct_GTV.nii.gz") key = f"m6_{pid}_{date}" row = {"key": key, "subject": pid, "date": date, "img": t1c, "source": "m6"} - if os.path.exists(gtv) and os.path.exists(ct): - labeled.append({**row, "label": gtv, "ct": ct, "gtv_registered": False}) + if label and os.path.exists(ct): + labeled.append({**row, "label": label, "ct": ct, "gtv_registered": False}) else: unlabeled.append({**row, "label": None}) return labeled, unlabeled +def _ct_vertex_slice(ct): + """Topmost tissue slice with 10mm of continuity below, or None. + + Tissue > -150 HU (bone HU is unreliable in these CTs). The continuity + test skips small artifacts floating above the head. + """ + per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2)) + n8 = max(2, int(round(10 / ct.GetSpacing()[2]))) + for i in range(len(per) - 1, -1, -1): + if per[i] > 500 and per[max(0, i - n8):i + 1].min() >= 50: + return i + return None + + +def _skull_ring_gap(warped_ct, ha): + """Concentricity of the CT skull ring vs the MR head silhouette. + + The MR volumes are tight crops, so distance-to-crop-edge is + crop-dependent; instead march rays from the head center and measure the + gap between the MR head boundary and the CT skull ring (p97 of positive + CT) at two levels. Returns {"mid_mm", "low_mm", "std_mm"} or None when no + skull ring is found. Aligned = small gap with low angular variance. + """ + posw = warped_ct[warped_ct > 0] + if not posw.size: + return None + bthr = float(np.percentile(posw, 97)) + sil = ha > 50 + izs = np.where(sil.sum((1, 2)) > 500)[0] + if not izs.size: + return None + + def _gap(iz): + sl = sil[iz] + if sl.sum() < 500: + return None, None + sk = warped_ct[iz] > bthr + if not sk.any(): + return None, None + hy, hx = np.nonzero(sl) + cy0, cx0 = float(hy.mean()), float(hx.mean()) + gaps = [] + for a in np.linspace(0, 2 * np.pi, 72, endpoint=False): + dy, dx = np.sin(a), np.cos(a) + rr = np.arange(0.0, 160.0, 0.5) + yy = np.clip((cy0 + rr * dy).astype(int), 0, ha.shape[0] - 1) + xx = np.clip((cx0 + rr * dx).astype(int), 0, ha.shape[1] - 1) + rh = rr[np.where(sl[yy, xx])[0]] + rs = rr[np.where(sk[yy, xx])[0]] + if len(rh) and len(rs): + gaps.append(float(rh.max() - rs.max())) + if not gaps: + return None, None + return float(np.mean(gaps)), float(np.std(gaps)) + + gm, sm = _gap(int(izs.min() + (izs.max() - izs.min()) * 0.5)) + gl, sl = _gap(int(izs.min() + (izs.max() - izs.min()) * 0.75)) + if gm is None: + return None + std = max(s for s in (sm, sl) if s is not None) if (sm is not None or sl is not None) else 0.0 + return {"mid_mm": round(gm, 1), "low_mm": round(gl, 1) if gl is not None else None, + "std_mm": round(std, 1)} + + +def _mi_value(fixed_arr, moving_arr, bins=32): + """Mattes-style MI (nats) on percentile-binned intensities, subsampled for speed.""" + 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(fixed_arr), q(moving_arr) + step = max(1, int(np.ceil((A.size / 200_000) ** (1.0 / 3)))) + h = np.bincount(A[::step, ::step, ::step].ravel() * bins + + B[::step, ::step, ::step].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"): + m = np.where((p > 0) & (px > 0) & (py > 0), np.log(p / (px * py)), 0.0) + return float((p * m).sum()) + + +def _valid_pyramid(size, shrink, smooth, min_dim=6): + sh = [s for s in shrink if min(size) / s >= min_dim] or [1] + sm = smooth[: len(sh)] + [0] * max(0, len(sh) - len(smooth)) + return sh, sm + + def register_ct_to_t1c(ct, t1c, threads=12): + """Rigid registration with CT as fixed and T1c as moving, Mattes MI on + head-box images. + + Both inputs share physical space (oriented). The registration grid is a head + box located in CT space from the top of the head (the CTs are head-neck + exams, so an MR-head crop would show mostly neck/air on the CT); the box + excludes most air so MI is computed on tissue. The box is in physical + space, so the resulting transform applies directly to the full volumes. + Returns (transform, info): transform is the registration transform, mapping + T1c (moving) coordinates -> CT (fixed) coordinates; resampling CT/TV onto + the T1c grid with it applies the inverse transform to CT-space points, + placing them in T1c (MR) coordinates. info has per-start MI/timing, the + affine decomposition (matrix/offset/scale), and the skull-ring concentricity + QA statistic. + """ import time t_start = time.time() try: # noqa @@ -76,20 +194,144 @@ def register_ct_to_t1c(ct, t1c, threads=12): pass ct = sitk.Cast(ct, sitk.sitkFloat32) t1c = sitk.Cast(t1c, sitk.sitkFloat32) - r = sitk.ImageRegistrationMethod() - r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32) - r.SetOptimizerAsRegularStepGradientDescent( - learningRate=10.0, minStep=0.5, numberOfIterations=600, - relaxationFactor=0.5, gradientMagnitudeTolerance=1e-5, - estimateLearningRate=sitk.ImageRegistrationMethod.EachIteration) - r.SetOptimizerScalesFromPhysicalShift() - r.SetInitialTransform(sitk.CenteredTransformInitializer(ct, t1c, sitk.Transform(3, sitk.sitkSimilarity))) - r.SetShrinkFactorsPerLevel([8, 4, 2, 1]) - r.SetSmoothingSigmasPerLevel([4, 2, 1, 0]) - r.SetInterpolator(sitk.sitkLinear) - aff = r.Execute(ct, t1c) - print("affine done %.0fs" % (time.time() - t_start), flush=True) - return aff + + from src.common import head_mask_from_image + ha = sitk.GetArrayFromImage(t1c) + hm = head_mask_from_image(ha) + # Registration grid = a head box located in CT space. The CTs are head-neck + # (often head-neck-chest) exams, so cropping to the MR head box would leave + # mostly neck/chest/air on the fixed image and MI locks onto the wrong + # structures. Vertex = topmost tissue slice with 10mm of continuity below + # (skips small artifacts floating above the head; bone HU is unreliable + # here, so tissue > -150 HU is used). x/y from the tissue extent of the + # 150mm band below the vertex. Box covers the whole head (x/y: extent + + # 15mm margin, z: +10 to -215 from vertex). + ca = sitk.GetArrayFromImage(ct) + tissue = ca > -150 + csp, co = np.array(ct.GetSpacing()), np.array(ct.GetOrigin()) + vz = _ct_vertex_slice(ct) + if vz is not None: + ztop = co[2] + vz * csp[2] + band = tissue[max(0, vz - int(round(150 / csp[2]))):vz + 1] + cols = np.nonzero(band.any(0)) + box_lo = (co[0] + cols[1].min() * csp[0] - 15, + co[1] + cols[0].min() * csp[1] - 15, ztop - 215) + box_hi = (co[0] + cols[1].max() * csp[0] + 15, + co[1] + cols[0].max() * csp[1] + 15, ztop + 10) + else: # no head found: fall back to the whole CT volume + gs = ct.GetSize() + box_lo = tuple(float(x) for x in co) + box_hi = (co[0] + gs[0] * csp[0], co[1] + gs[1] * csp[1], co[2] + gs[2] * csp[2]) + tsp = np.array(t1c.GetSpacing()) + shape = tuple(int(np.ceil((box_hi[d] - box_lo[d]) / tsp[d])) + 1 for d in range(3)) + ref = sitk.GetImageFromArray(np.zeros(shape, np.float32)) + ref.SetSpacing(t1c.GetSpacing()) + ref.SetDirection(t1c.GetDirection()) + ref.SetOrigin(tuple(float(x) for x in box_lo)) + fixed = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) + moving = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) + + stages = [] + + def run_stage(name, init_tr, shrink, smooth, lr, min_step, iters): + t0 = time.time() + r = sitk.ImageRegistrationMethod() + r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32) + r.SetOptimizerAsRegularStepGradientDescent( + learningRate=lr, minStep=min_step, numberOfIterations=iters, + relaxationFactor=0.5, gradientMagnitudeTolerance=1e-6, + estimateLearningRate=sitk.ImageRegistrationMethod.EachIteration) + r.SetOptimizerScalesFromPhysicalShift() + sh, sm = _valid_pyramid(fixed.GetSize(), shrink, smooth) + r.SetShrinkFactorsPerLevel(sh) + r.SetSmoothingSigmasPerLevel(sm) + r.SetInterpolator(sitk.sitkLinear) + r.SetInitialTransform(init_tr, inPlace=False) + fa = sitk.GetArrayFromImage(fixed) + mi_in = _mi_value(fa, sitk.GetArrayFromImage(sitk.Resample(moving, fixed, init_tr, sitk.sitkLinear, 0.0))) + tr = r.Execute(fixed, moving) + mi_out = _mi_value(fa, sitk.GetArrayFromImage(sitk.Resample(moving, fixed, tr, sitk.sitkLinear, 0.0))) + stages.append({"name": name, "mi_in": round(mi_in, 4), "mi_out": round(mi_out, 4), + "seconds": round(time.time() - t0, 1), "shrink": sh, "smooth": sm}) + return tr + +# The MI surface has tilted local minima that can score HIGHER MI than the + # true alignment (patients are repositioned between the CT and MR sessions, + # so the inter-session head tilt can reach ~30 deg), so MI alone cannot + # rank candidates. Run the full pyramid from the centered start; if its + # full-resolution geometric concentricity (skull-ring gap) is poor, repeat + # from a grid of initial tilts and keep the best by gap. + mr_c = np.array(t1c.GetOrigin()) + np.argwhere(hm).mean(0) * tsp + ct_c = np.array([(box_lo[d] + box_hi[d]) / 2.0 for d in range(3)]) + off = (float(mr_c[0] - ct_c[0]), float(mr_c[1] - ct_c[1]), float(mr_c[2] - ct_c[2])) + rad = np.pi / 180.0 + + def full_run(sx, sy, sz): + init = sitk.Euler3DTransform() + init.SetCenter(tuple(float(x) for x in ct_c)) + init.SetParameters((sx * rad, sy * rad, sz * rad) + off) + tr0 = run_stage(f"start_{sx:+d}x{sy:+d}y{sz:+d}z", init, [8, 4, 2, 1], [4, 2, 1, 0], 10.0, 0.5, 400) + gap = _skull_ring_gap(sitk.GetArrayFromImage( + sitk.Resample(ct, t1c, tr0.GetInverse(), sitk.sitkLinear)), ha) + return tr0, gap + + def is_good(g): + return bool(g) and g["std_mm"] <= 12.0 and 0.0 <= g["mid_mm"] <= 20.0 + + tr_reg, gap = full_run(0, 0, 0) + if not is_good(gap): + # inter-session tilt can reach ~30 deg in pitch: search a grid of + # initial rotations, keep the best by full-res concentricity + best = (tr_reg, gap) + starts = [(sx, sy, 0) + for sx in (-30, -20, -10, 0, 10, 20, 30) + for sy in (-20, 0, 20) if (sx, sy) != (0, 0)] + for sx, sy, sz in starts: + tr_i, gap_i = full_run(sx, sy, sz) + gi = (gap_i["std_mm"], gap_i["mid_mm"]) if gap_i else (999.0, 999.0) + gb = (best[1]["std_mm"], best[1]["mid_mm"]) if best[1] else (999.0, 999.0) + if gi < gb: + best = (tr_i, gap_i) + tr_reg, gap = best + if not is_good(gap): + for sx, sy, sz in [(sx, 0, sz) for sz in (15, -15) for sx in (-30, 0, 30)]: + tr_i, gap_i = full_run(sx, sy, sz) + gi = (gap_i["std_mm"], gap_i["mid_mm"]) if gap_i else (999.0, 999.0) + gb = (best[1]["std_mm"], best[1]["mid_mm"]) if best[1] else (999.0, 999.0) + if gi < gb: + best = (tr_i, gap_i) + tr_reg, gap = best + # tr_reg (Execute convention) maps CT (fixed) -> T1c (moving). Its inverse + # maps T1c -> CT, so resampling CT/TV onto the T1c grid with it applies the + # inverse transform to CT-space points, placing them in T1c (MR) coords. + tr = tr_reg.GetInverse() + + # QA: CT head-box coverage + skull-ring concentricity vs the MR head silhouette + from scipy import ndimage + def _tp(x, y, z): + return np.array(tr.TransformPoint((float(x), float(y), float(z)))) + + t0 = _tp(0, 0, 0) + mat = np.column_stack([_tp(1, 0, 0) - t0, _tp(0, 1, 0) - t0, _tp(0, 0, 1) - t0]) + scale = [round(float(np.linalg.norm(mat[:, i])), 4) for i in range(3)] + warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, tr, sitk.sitkLinear)) + 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) + coverage = float(tissue_w[hm].sum() / max(hm.sum(), 1)) + skull_gap = _skull_ring_gap(warped_ct, ha) + info = {"stages": stages, + "affine_stage": {"matrix": [[round(float(v), 4) for v in row] for row in mat.tolist()], + "offset": [round(float(x), 3) for x in t0], "scale": scale}, + "coverage_head_box": round(coverage, 3), + "skull_ring_gap_mm": skull_gap, + "ct_head_box_mm": [[round(float(x), 1) for x in box_lo], [round(float(x), 1) for x in box_hi]], + "seconds_total": round(time.time() - t_start, 1)} + print("reg done %.0fs (rigid %.1fs) mi %.3f -> %.3f cov %.2f%s" % ( + info["seconds_total"], sum(s["seconds"] for s in stages), + stages[-1]["mi_in"], stages[-1]["mi_out"], coverage, + " skull_gap mid %.1fmm std %.1fmm" % (skull_gap["mid_mm"], skull_gap["std_mm"]) if skull_gap else " (no skull ring)"), flush=True) + return tr, info, tissue_w def reg_worker(row): @@ -97,25 +339,89 @@ def reg_worker(row): key = row["key"] t1c = read_nii_img(row["img"]) ct = read_nii_img(row["ct"]) + # guard: corrupted CT (upper brain missing) cannot be registered + vz = _ct_vertex_slice(ct) + if vz is not None: + sp = ct.GetSpacing()[2] + per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2)) + iz = vz - int(round(100 / sp)) + incomplete = not (0 <= iz < len(per) and per[iz] > 5000) + else: + incomplete = True + if incomplete: + with open(os.path.join(reg_dir, key + "_reg.json"), "w") as f: + json.dump({"key": key, "subject": row["subject"], "date": row["date"], + "tv_nvox": 0, "tv_accepted": False, + "accept_reason": "ct head incomplete or not found", + "inputs": {"ct": row["ct"], "mr": row["img"], "tv": row["label"]}}, + f, indent=1) + return key, False, 0 gtv = read_nii_img(row["label"]) - tr = register_ct_to_t1c(ct, t1c) - out = os.path.join(reg_dir, key + "_gtv.nii.gz") - warped = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0) - sitk.WriteImage(warped, out, True) - z = sitk.GetArrayFromImage(warped) - return key, os.path.exists(out) and (z > 0.5).sum() > 0, int((z > 0.5).sum()) + tr, info, tissue_w = register_ct_to_t1c(ct, t1c) + # QA files in m6reg, everything in (oriented) MR coordinates: + # _ct: CT warped to the MR grid (linear), for overlaying _tv + # _mr: the oriented MR volume itself + # _tv: original TV warped to the MR grid (nearest) — the study label + # _reg.json: registration metrics + final image information + tv_mr = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0) + z = sitk.GetArrayFromImage(tv_mr) + nvox = int((z > 0.5).sum()) + from scipy import ndimage + tv_in_ct_frac = 1.0 + if nvox > 0: + tvd = ndimage.binary_dilation(z > 0.5, iterations=3) + tv_in_ct_frac = float((tvd & tissue_w).sum() / max(tvd.sum(), 1)) + tv_in_ct = tv_in_ct_frac >= 0.5 # slab must straddle the lesion + gap = info["skull_ring_gap_mm"] + if gap and 0.0 <= gap["mid_mm"] <= 20.0 and (gap["std_mm"] or 0) <= 12.0: + align = "concentric" + else: + align = "not_concentric" + accepted = nvox > 0 and tv_in_ct and align == "concentric" + info = dict(info) + info.update({ + "key": key, "subject": row["subject"], "date": row["date"], + "tv_nvox": nvox, "tv_in_ct_frac": round(tv_in_ct_frac, 3), + "alignment": align, + "tv_accepted": accepted, + "accept_reason": ("ok" if accepted else + "tv outside MR grid" if nvox == 0 else + "tv outside CT coverage" if not tv_in_ct else + f"skull ring not concentric (gap={gap})"), + "inputs": {"ct": row["ct"], "mr": row["img"], "tv": row["label"]}, + "images": { + "ct": {"size": list(ct.GetSize()), "spacing": [float(s) for s in ct.GetSpacing()], + "origin": [float(o) for o in ct.GetOrigin()], + "direction": [float(x) for x in ct.GetDirection()], + "pixel_type": ct.GetPixelIDTypeAsString()}, + "mr": {"size": list(t1c.GetSize()), "spacing": [float(s) for s in t1c.GetSpacing()], + "origin": [float(o) for o in t1c.GetOrigin()], + "direction": [float(x) for x in t1c.GetDirection()], + "pixel_type": t1c.GetPixelIDTypeAsString()}, + "tv_warp": "mr grid (same as images.mr)", + }, + }) + sitk.WriteImage(sitk.Resample(ct, t1c, tr, sitk.sitkLinear), os.path.join(reg_dir, key + "_ct.nii.gz"), True) + sitk.WriteImage(t1c, os.path.join(reg_dir, key + "_mr.nii.gz"), True) + sitk.WriteImage(tv_mr, os.path.join(reg_dir, key + "_tv.nii.gz"), True) + with open(os.path.join(reg_dir, key + "_reg.json"), "w") as f: + json.dump(info, f, indent=1) + return key, accepted, nvox def main(): ap = argparse.ArgumentParser() ap.add_argument("--workers", type=int, default=2) ap.add_argument("--skip-reg", action="store_true") + ap.add_argument("--limit", type=int, default=0, help="register at most N missing (0=all)") args = ap.parse_args() labeled, unlabeled = discover() print(f"m6 labeled candidates: {len(labeled)}, unlabeled: {len(unlabeled)}") reg_dir = d("data/m6reg") if not args.skip_reg: - todo = [r for r in labeled if not os.path.exists(os.path.join(reg_dir, r["key"] + "_gtv.nii.gz"))] + todo = [r for r in labeled if not os.path.exists(os.path.join(reg_dir, r["key"] + "_tv.nii.gz"))] + if args.limit > 0: + todo = todo[:args.limit] print(f"registering {len(todo)} GTVs with {args.workers} workers") with ProcessPoolExecutor(max_workers=args.workers) as ex: futs = {ex.submit(reg_worker, r): r for r in todo} @@ -129,13 +435,20 @@ def main(): # refresh manifest state labeled, unlabeled = discover() for r in labeled: - p = os.path.join(reg_dir, r["key"] + "_gtv.nii.gz") - if os.path.exists(p): - r["label"] = p - else: - r["label"] = None + p = os.path.join(reg_dir, r["key"] + "_tv.nii.gz") + accepted = True + jf = os.path.join(reg_dir, r["key"] + "_reg.json") + if os.path.exists(jf): + try: + accepted = json.load(open(jf)).get("tv_accepted", False) + except (OSError, ValueError): + accepted = False + r["label"] = p if os.path.exists(p) and accepted else None keep = [r for r in labeled if r.get("label")] - unlabeled = [r for r in unlabeled if r["key"] not in {k["key"] for k in keep}] + # TV present but warped to zero overlap with the T1c FOV (non-cranial + # target) or registration failed: keep the T1c in the pool as unlabeled + dropped = [{**r, "label": None} for r in labeled if r["key"] not in {k["key"] for k in keep}] + unlabeled = [r for r in unlabeled if r["key"] not in {k["key"] for k in keep}] + dropped for r in keep + unlabeled: r.pop("ct", None) save_jsonl(keep, os.path.join(DATA, "manifests/m6_labeled.jsonl")) diff --git a/scripts/bench_registration.py b/scripts/bench_registration.py new file mode 100644 index 0000000..1cc6e6a --- /dev/null +++ b/scripts/bench_registration.py @@ -0,0 +1,423 @@ +"""Benchmark the 4 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 + +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 three of the four toolkits are not pip-installable into the main +`longitudinal` env (elastix needs conda-forge; antspyx has no cp314 wheel; +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. The `--merge` mode then combines the +per-pipeline JSON files (accumulated across envs) into one summary + plot. + +Usage (repo root): + # main env (SimpleITK baseline; elastix/deepreg report as unavailable) + source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal + python scripts/bench_registration.py --pipelines simpleitk elastix deepreg + + # 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 +from src import registration_pipelines as RP + + +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) + 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, + "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 _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) + 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)) + 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"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] + 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, + "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}") + 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']}]") + 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") + 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}") + 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] + fig, axes = plt.subplots(1, 4, figsize=(18, 4)) + colors = ["#4878CF", "#D64F5C", "#6ACC65", "#9467BD"] + 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)]): + 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 |", + "|---|---|---|---|---|---|---|---|---|---|---|"] + 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')} |") + 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("--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") + 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() \ No newline at end of file diff --git a/scripts/preprocess.py b/scripts/preprocess.py index 6c26933..1e566c6 100644 --- a/scripts/preprocess.py +++ b/scripts/preprocess.py @@ -90,8 +90,12 @@ def main(): rows = rows[: args.subset] out_dir = d("data/proc") meta_dir = d("data/procmeta") - todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows - if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz"))] + def needs_proc(r): + if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz")): + return True + return bool(r.get("label")) and not os.path.exists(os.path.join(out_dir, r["key"] + "_label.nii.gz")) + + todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows if needs_proc(r)] print(f"rows={len(rows)} todo={len(todo)}") errs = [] with ProcessPoolExecutor(max_workers=args.workers) as ex: diff --git a/src/common.py b/src/common.py index f2a4890..e128407 100644 --- a/src/common.py +++ b/src/common.py @@ -204,7 +204,7 @@ def read_nii_arr(path): def read_nii_img(path): - return sitk.ReadImage(path) + return sitk.DICOMOrient(sitk.ReadImage(path)) def write_arr(arr, path, itk_img=None): diff --git a/src/registration_pipelines.py b/src/registration_pipelines.py new file mode 100644 index 0000000..df93223 --- /dev/null +++ b/src/registration_pipelines.py @@ -0,0 +1,566 @@ +"""Four rigid CT -> T1c registration pipelines behind one unified contract. + +Problem (shared with ``scripts/02_build_m6_dataset.py``) +-------------------------------------------------------- +Register a head-and-neck CT to a tightly-cropped T1c MRI so that the GTV / CT can +be warped onto the T1c (MR) native grid. The pairs are cross-modality +(CT vs MRI) and repositioned between sessions (inter-session tilt up to ~30 deg), +so every pipeline uses an intensity metric suited to multi-modality (Mattes MI / +cross-correlation) and a rigid (6-DOF) transform. + +Unified output contract +----------------------- +Each ``register_*`` returns ``(transform, meta)`` where ``transform`` is a +SimpleITK transform such that:: + + sitk.Resample(ct, t1c, transform, sitk.sitkLinear) # CT -> MR grid + sitk.Resample(tv, t1c, transform, sitk.sitkNearestNeighbor, 0) # GTV -> MR grid + +places the CT / GTV concentric with the MR head. In the ``Resample`` +destination->source convention this transform therefore maps **T1c -> CT** +coordinates. This is exactly the convention the production SimpleITK pipeline +returns (``tr_reg.GetInverse()`` with ``tr_reg = Execute(fixed=CT, moving=T1c)``). + +Pipelines +--------- +- ``register_simpleitk`` : SimpleITK ImageRegistrationMethod, Mattes MI on a + CT-space head box, multi-start tilt grid. Reuses the *proven* production + ``register_ct_to_t1c`` so the baseline is bit-identical to 02_build_m6_dataset. +- ``register_elastix`` : ITK-Elastix CLI (Rigid3DTransform, Mattes MI). +- ``register_antspyx`` : ANTs / ANTsPy ``ants.registration`` (Rigid). +- ``register_deepreg`` : Oxford DeepReg (CycleGAN deformation-field + registration; needs a pretrained checkpoint, see module docstring). + +A toolkit that is not installed (or, for DeepReg, missing its checkpoint) +raises :class:`PipelineUnavailable` with the exact install / precondition. The +benchmark catches that and reports the pipeline as *not available* instead of +failing the whole run, so the script is env-portable: run it under whichever +conda env has a given toolkit installed. + +Orientation robustness +---------------------- +The hand-derivation of "which way does toolkit X return its transform" differs +per toolkit. :func:`fix_orientation` resolves this by warping the CT with both +``T`` and ``T^-1`` and keeping whichever gives the smaller skull-ring +concentricity gap (a pure correctness gate -- a wrongly-oriented rigid transform +lands the head on the far side of the box and yields a pathological gap). It is +applied to elastix / antspyx / deepreg, not to the trusted production +SimpleITK path. +""" +import os +import json +import shutil +import subprocess +import importlib.util + +import numpy as np +import SimpleITK as sitk +from scipy import ndimage + +from src.common import ROOT, d + + +class PipelineUnavailable(RuntimeError): + """Raised when a pipeline's toolkit (or a required asset) is not present.""" + + +# --------------------------------------------------------------------------- # +# shared helpers +# --------------------------------------------------------------------------- # + +def _load_m6(): + """Import scripts/02_build_m6_dataset.py (name starts with a digit).""" + p = os.path.join(ROOT, "scripts", "02_build_m6_dataset.py") + spec = importlib.util.spec_from_file_location("m6_reg", p) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def _skull_gap(warped_ct, ha): + m6 = _load_m6() + return m6._skull_ring_gap(warped_ct, ha) + + +def _gap_score(g): + """Lower is better; None / missing -> worst possible.""" + if not g: + return float("inf") + return (g.get("std_mm", 999.0) or 999.0) + (g.get("mid_mm", 999.0) or 999.0) + + +def _warp_ct(ct, t1c, T): + return sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear)) + + +def fix_orientation(ct, t1c, cand): + """Return (best_T, inverted) choosing between ``cand`` and ``cand^-1``. + + The choice is made on the skull-ring concentricity gap of the warped CT, + which is a crop-independent geometric measure. A rigid registration that is + inverted by 180 deg puts the CT head outside the MR head box and the gap + blows up, so min-gap reliably recovers the true orientation. + """ + ha = sitk.GetArrayFromImage(t1c) + s0 = _gap_score(_skull_gap(_warp_ct(ct, t1c, cand), ha)) + s1 = _gap_score(_skull_gap(_warp_ct(ct, t1c, cand.GetInverse()), ha)) + if s1 < s0: + return cand.GetInverse(), True + return cand, False + + +def centroid_offset(ct, t1c, ct_box=None): + """Physical-space translation (dx,dy,dz) that maps the CT head centroid onto the + T1c head centroid, or None. Used as the *common* centered initialization so the + single-init toolkits (SimpleITK / ANTs / Elastix) start from the same gross + alignment and are compared on their optimizer, not on who recovers the + ~30-60 mm inter-session offset first. ``ct_box`` = the CT resampled to the + head box, so its tissue centroid is the CT *head* (not the head+neck centroid).""" + from src.common import head_mask_from_image + tsp = np.array(t1c.GetSpacing()) + t1c_f = sitk.Cast(t1c, sitk.sitkFloat32) + t1c_c = np.array(t1c.GetOrigin()) + np.argwhere( + head_mask_from_image(sitk.GetArrayFromImage(t1c_f))).mean(0) * tsp + src = ct_box if ct_box is not None else ct + a = sitk.GetArrayFromImage(src) + if float(a.max()) > 200: # CT (Hounsfield) + tissue = a > -150 + elif (a > 0).any(): # MR + tissue = a > np.percentile(a[a > 0], 2) + else: + return None + if not tissue.any(): + return None + c = np.array(src.GetOrigin()) + np.argwhere(tissue).mean(0) * tsp + return (float(t1c_c[0] - c[0]), float(t1c_c[1] - c[1]), float(t1c_c[2] - c[2])) + + +def _ct_head_box_mm(ct): + """CT-space head box in physical mm: (box_lo, box_hi) tuples, or None.""" + m6 = _load_m6() + ca = sitk.GetArrayFromImage(ct) + tissue = ca > -150 + csp, co = np.array(ct.GetSpacing()), np.array(ct.GetOrigin()) + vz = m6._ct_vertex_slice(ct) + if vz is None: + return None, None + ztop = co[2] + vz * csp[2] + band = tissue[max(0, vz - int(round(150 / csp[2]))):vz + 1] + cols = np.nonzero(band.any(0)) + box_lo = (co[0] + cols[1].min() * csp[0] - 15, + co[1] + cols[0].min() * csp[1] - 15, ztop - 215) + box_hi = (co[0] + cols[1].max() * csp[0] + 15, + co[1] + cols[0].max() * csp[1] + 15, ztop + 10) + return box_lo, box_hi + + +# --------------------------------------------------------------------------- # +# 2. SimpleITK (production baseline -- reused verbatim) +# --------------------------------------------------------------------------- # + +def register_simpleitk(ct, t1c, threads=12, mode="benchmark", iters_full=150): + """SimpleITK rigid CT->T1c, Mattes MI on a CT-space head box. + + mode="benchmark" (default): a single multi-resolution pyramid run from the + centered start, capped at ``iters_full`` iterations at full resolution. + This is the *standard* SimpleITK algorithm and matches the budget the + other toolkits get in the benchmark (one rigid MI registration), so the + pipelines are compared on equal footing. + mode="production": the exact multi-start tilt-grid ``register_ct_to_t1c`` + from scripts/02_build_m6_dataset.py (robust to the ~30 deg inter-session + tilt but much slower -- not the default for the benchmark). + """ + if mode == "production": + m6 = _load_m6() + tr, info, _tissue_w = m6.register_ct_to_t1c(ct, t1c, threads=threads) + meta = { + "pipeline": "simpleitk", "mode": "production", + "toolkit": f"SimpleITK {sitk.Version_VersionString()}", + "metric": "MattesMutualInformation (head box, CT space)", + "transform": "Rigid (Euler3D, multi-start tilt grid)", + "convention": "T1c->CT (production, verified)", "auto_inverted": False, + "coverage_head_box": info.get("coverage_head_box"), + "seconds_reg": sum(s.get("seconds", 0) for s in info.get("stages", [])), + "seconds_total": info.get("seconds_total"), "stages": info.get("stages"), + "skull_ring_gap_mm": info.get("skull_ring_gap_mm"), + } + return tr, meta + + m6 = _load_m6() + try: + sitk.CommonProperties.SetGlobalDefaultNumberOfThreads(threads) + except AttributeError: + pass + ct = sitk.Cast(ct, sitk.sitkFloat32); t1c = sitk.Cast(t1c, sitk.sitkFloat32) + lo, hi = _ct_head_box_mm(ct) + if lo is None: + raise PipelineUnavailable("simpleitk: no CT head box found (corrupt CT?)") + tsp = np.array(t1c.GetSpacing()) + shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3)) + ref = sitk.GetImageFromArray(np.zeros(shape, np.float32)) + ref.SetSpacing(t1c.GetSpacing()); ref.SetDirection(t1c.GetDirection()) + ref.SetOrigin(tuple(float(x) for x in lo)) + fixed = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) + moving = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) + # centered start: align CT head centroid with T1c head centroid (shared init + # with the other single-init toolkits via centroid_offset) + off = centroid_offset(ct, t1c, ct_box=fixed) or (0.0, 0.0, 0.0) + ct_c = np.array([(lo[i] + hi[i]) / 2.0 for i in range(3)]) + init = sitk.Euler3DTransform(); init.SetCenter(tuple(float(x) for x in ct_c)) + init.SetParameters((0.0, 0.0, 0.0) + off) + import time + t0 = time.time() + r = sitk.ImageRegistrationMethod() + r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32) + r.SetOptimizerAsRegularStepGradientDescent( + learningRate=10.0, minStep=0.2, numberOfIterations=iters_full, + relaxationFactor=0.5, gradientMagnitudeTolerance=1e-6) + r.SetOptimizerScalesFromPhysicalShift() + sh = [s for s in [4, 2, 1] if min(shape) / s >= 6] or [1] + sm = [2, 1, 0][:len(sh)] + r.SetShrinkFactorsPerLevel(sh); r.SetSmoothingSigmasPerLevel(sm) + r.SetInterpolator(sitk.sitkLinear); r.SetInitialTransform(init, inPlace=False) + tr0 = r.Execute(fixed, moving) + dt = time.time() - t0 + tr = tr0.GetInverse() # contract: T1c -> CT + ha = sitk.GetArrayFromImage(t1c) + gap = m6._skull_ring_gap(_warp_ct(ct, t1c, tr), ha) + meta = { + "pipeline": "simpleitk", "mode": "benchmark", + "toolkit": f"SimpleITK {sitk.Version_VersionString()}", + "metric": "MattesMutualInformation (head box, CT space)", + "transform": (f"Rigid (Euler3D, fixed-LR budget, pyramid {sh}, " + f"{iters_full} iters/level)"), + "convention": "T1c->CT (benchmark budget)", "auto_inverted": False, + "seconds_reg": round(dt, 1), "pyramid_shrink": sh, "pyramid_smooth": sm, + "skull_ring_gap_mm": gap, + } + return tr, meta + + +# --------------------------------------------------------------------------- # +# 1. ITK-Elastix +# --------------------------------------------------------------------------- # + +def _elastix_param_file(workdir, fixed, moving, fixed_mask, moving_mask): + """Rigid3DTransform + Mattes MI parameter file.""" + lines = [ + "(ImageDimension 3)", + "(FixedInternalImagePixelType float)", + "(MovingInternalImagePixelType float)", + "(Transform \"Rigid3DTransform\")", + "(TransformBase \"ElastixTransformBase\")", + "(FixedCenteredTransform 1)", + "(TransformParameterReadoutFormat 1)", + "(Optimizer \"AdaptiveStochasticGradientDescent\")", + "(MaximumNumberOfIterations 6000)", + "(StoppingValue 0.0005)", + "(StoppingIterationCount 2500)", + "(ImageMetric \"MattesMutualInformation\")", + "(NumberOfHistogramBins 10)", + "(ImageSampler \"RandomSampler\")", + "(ImageSamplerPatchSize 25 25 25)", + "(FixedSubstresamplerClass \"RandomSubstresampler\")", + "(FixedSubstresamplerScale 2)", + "(ResampleInterpolator \"FinalBSplineInterpolator\")", + "(FinalBSplineInterpolationOrder 1)", + "(ResultStationaryImage \"result.mha\")", + "(WriteResultImage 0)", + "(CompressResultImages 1)", + f"(FixedImage {fixed})", + f"(FixedImageMask {fixed_mask})", + f"(MovingImage {moving})", + f"(MovingImageMask {moving_mask})", + ] + p = os.path.join(workdir, "elastix_params.txt") + with open(p, "w") as f: + f.write("\n".join(lines) + "\n") + return p + + +def register_elastix(ct, t1c, workdir=None, threads=8): + """ITK-Elastix rigid CT->T1c (fixed=T1c, moving=CT), Mattes MI.""" + if shutil.which("elastix") is None: + raise PipelineUnavailable( + "elastix/transformix binaries not found. Install with " + "`conda install -c conda-forge elastix` (conda-forge must be " + "reachable) or build from https://elastix.org. Not on PyPI.") + workdir = workdir or d("data/m6reg/elastix") + workdir = os.path.join(workdir, f"pid{os.getpid()}") # per-worker scratch + os.makedirs(workdir, exist_ok=True) + lo, hi = _ct_head_box_mm(ct) + if lo is None: + raise PipelineUnavailable("elastix: no CT head box found (corrupt CT?)") + # Resample both into a common CT-space head box (mirrors the production grid + # choice) so MI is computed on head tissue, not neck/air. + tsp = np.array(t1c.GetSpacing()) + shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3)) + ref = sitk.GetImageFromArray(np.zeros(shape, np.float32)) + ref.SetSpacing(t1c.GetSpacing()) + ref.SetDirection(t1c.GetDirection()) + ref.SetOrigin(tuple(float(x) for x in lo)) + fixed = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # T1c = reference + moving = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # CT = moving + fixed_f, moving_f = os.path.join(workdir, "fixed.nii.gz"), os.path.join(workdir, "moving.nii.gz") + sitk.WriteImage(sitk.Cast(fixed, sitk.sitkFloat32), fixed_f, True) + sitk.WriteImage(sitk.Cast(moving, sitk.sitkFloat32), moving_f, True) + # all-ones masks sized to the box (registration on the whole head box) + ones = sitk.GetImageFromArray(np.ones(shape, np.uint8)) + ones.CopyInformation(ref) + fm, mm = os.path.join(workdir, "fixedmask.nii.gz"), os.path.join(workdir, "movingmask.nii.gz") + sitk.WriteImage(ones, fm, True); sitk.WriteImage(ones, mm, True) + + params = _elastix_param_file(workdir, fixed_f, moving_f, fm, mm) + outdir = os.path.join(workdir, "out") + shutil.rmtree(outdir, ignore_errors=True); os.makedirs(outdir) + env = dict(os.environ) + env["OMP_NUM_THREADS"] = str(threads) + cmd = ["elastix", "-f", fixed_f, "-m", moving_f, "-out", outdir, + "-t", params, "-threads", str(threads), + "-loglevel", "0", "-logtoconsole", "false"] + subprocess.run(cmd, check=True, env=env, capture_output=True, text=True) + # elastix fixed(T1c) -> moving(CT) transform, saved as an ITK transform file + res = os.path.join(outdir, "Result", "0GenericAnatTransform.txt") + if not os.path.exists(res): + raise PipelineUnavailable(f"elastix did not write a result transform: {res}") + cand = sitk.ReadTransform(res) + # cand maps fixed(T1c) -> moving(CT) == the contract (T1c->CT); still verify. + T, inverted = fix_orientation(ct, t1c, cand) + meta = { + "pipeline": "elastix", "toolkit": "ITK-Elastix", + "metric": "MattesMutualInformation (head box, CT space)", + "transform": "Rigid3DTransform (ASGD)", "convention": "T1c->CT (verified)", + "auto_inverted": inverted, "fixed": "T1c", "moving": "CT", + } + return T, meta + + +# --------------------------------------------------------------------------- # +# 3. ANTsPy +# --------------------------------------------------------------------------- # + +def register_antspyx(ct, t1c, transform_type="Rigid", aff_iterations=None, + aff_shrink_factors=None, aff_smoothing_sigmas=None, + aff_metric="mattes", use_head_box=True, center_init=True): + """ANTs rigid CT->T1c via antspyx (fixed=T1c, moving=CT). + + antspyx 0.6.3: ``ants.registration(fixed, moving, type_of_transform='Rigid', + aff_metric='mattes', aff_iterations=..., aff_shrink_factors=..., + aff_smoothing_sigmas=...)``. ``aff_metric='mattes'`` (the default) gives a + multi-resolution Mattes-MI rigid registration -- the cross-modality metric + this CT/MR pair needs. To keep the benchmark tractable and consistent with + the other pipelines (and to avoid neck/air in the metric), both images are + first resampled into the same CT-space head box. Returns the transform that + warps CT into the T1c grid (orientation verified by :func:`fix_orientation`). + """ + try: + import ants # antspyx ships the `ants` module + except Exception as e: # noqa + raise PipelineUnavailable( + f"antspyx not importable ({e!r}). `pip install antspyx` provides the " + f"`ants` module but needs a cp310-cp313 wheel; this interpreter is " + f"cp{sys_pyver()}. Install on Python 3.10-3.13, or `conda install " + "-c conda-forge antsx`.") + aff_iterations = tuple(aff_iterations) if aff_iterations else (500, 500, 200) + aff_shrink_factors = tuple(aff_shrink_factors) if aff_shrink_factors else (4, 2, 1) + aff_smoothing_sigmas = tuple(aff_smoothing_sigmas) if aff_smoothing_sigmas else (3, 2, 0) + if use_head_box: + lo, hi = _ct_head_box_mm(ct) + if lo is None: + raise PipelineUnavailable("antspyx: no CT head box found (corrupt CT?)") + tsp = np.array(t1c.GetSpacing()) + shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3)) + ref = sitk.GetImageFromArray(np.zeros(shape, np.float32)) + ref.SetSpacing(t1c.GetSpacing()); ref.SetDirection(t1c.GetDirection()) + ref.SetOrigin(tuple(float(x) for x in lo)) + fx = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # T1c box (fixed) + mv = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # CT box (moving) + else: + fx, mv = t1c, ct + import time + t0 = time.time() + # Centered start: antspyx cannot use a TranslationTransform as initial_transform +# (the ANTs CLI rejects it), and Rigid-from-identity does not recover the large +# (~30-60 mm) inter-session offset. So we pre-shift the CT box so its head +# centroid coincides with the T1c head centroid, run Rigid from identity, and +# compose the shift back into the final transform (undoing it to original CT +# coordinates). This gives ANTs the same centered initialization as the other +# single-init toolkits so the pipelines are compared on their optimizer. + off = None + if center_init: + off = centroid_offset(ct, t1c, ct_box=mv) + if off is not None: + tr_off = sitk.TranslationTransform(3) + tr_off.SetOffset((-off[0], -off[1], -off[2])) + mv = sitk.Resample(mv, mv, tr_off, sitk.sitkLinear, 0.0) + af = ants.from_sitk(sitk.Cast(fx, sitk.sitkFloat32)) + am = ants.from_sitk(sitk.Cast(mv, sitk.sitkFloat32)) + r = ants.registration(fixed=af, moving=am, type_of_transform=transform_type, + aff_metric=aff_metric, aff_iterations=aff_iterations, + aff_shrink_factors=aff_shrink_factors, + aff_smoothing_sigmas=aff_smoothing_sigmas) + dt = time.time() - t0 + # r['fwdtransforms'][0] is a path to the (single) Rigid transform file, in the + # shared physical space (but on the pre-shifted CT). Compose the shift back so + # the result maps T1c -> original CT; fix_orientation resolves T vs T^-1. + T_ants = sitk.ReadTransform(r["fwdtransforms"][0]) + if off is not None: + tr_undo = sitk.TranslationTransform(3) + tr_undo.SetOffset((-off[0], -off[1], -off[2])) + cand = sitk.CompositeTransform([T_ants, tr_undo]) + else: + cand = T_ants + T, inverted = fix_orientation(ct, t1c, cand) + meta = { + "pipeline": "antspyx", "toolkit": f"ANTsPy {getattr(ants, '__version__', '?')}", + "metric": f"ANTs {aff_metric} (multi-resolution)", + "transform": f"ANTs {transform_type} (head box)", "convention": "T1c->CT (verified)", + "auto_inverted": inverted, "aff_iterations": list(aff_iterations), + "aff_shrink_factors": list(aff_shrink_factors), "seconds_reg": round(dt, 1), + } + return T, meta + + +def sys_pyver(): + import sys + return sys.version_info[0] * 100 + sys.version_info[1] + + +# --------------------------------------------------------------------------- # +# 4. DeepReg (Oxford) +# --------------------------------------------------------------------------- # + +def register_deepreg(ct, t1c, workdir=None, checkpoint=None, config=None): + """Oxford DeepReg deformation-field registration (CT -> T1c). + + DeepReg is not a one-shot rigid tool: it is a CycleGAN-based image-to-image + registration that must be *trained* or run from a *pretrained checkpoint*. + It predicts a voxel displacement field; we convert that to a SimpleITK + DisplacementFieldTransform to warp CT/GTV onto the T1c grid. Preconditions: + - ``deepreg`` installed from git (https://github.com/OB-CVHU/DeepReg) + -- the PyPI ``deepreg`` (0.1.2) is an unrelated stub that pins + numpy==1.18.5 and will not work. + - a 3D pretrained checkpoint (``checkpoint``) -- typically one trained on + a CT/MRI template pair; see the repo's pretrained-models note. + """ + try: + import deepreg # noqa + except Exception as e: # noqa + raise PipelineUnavailable( + f"deepreg not importable ({e!r}). Install the Oxford DeepReg from " + "source: `git clone https://github.com/OB-CVHU/DeepReg && " + "cd DeepReg && pip install -e .` (NOT the PyPI `deepreg` stub). " + "Also requires a pretrained 3D checkpoint (see --deepreg-ckpt).") + if not checkpoint or not os.path.exists(checkpoint): + raise PipelineUnavailable( + "deepreg: no pretrained checkpoint provided. Pass --deepreg-ckpt " + "to a 3D .ckpt produced by `deepreg train` (or a released " + "pretrained model). DeepReg cannot do one-shot registration " + "without weights.") + workdir = workdir or d("data/m6reg/deepreg") + workdir = os.path.join(workdir, f"pid{os.getpid()}") # per-worker scratch + os.makedirs(workdir, exist_ok=True) + cfg = _deepreg_config(workdir, ct, t1c, checkpoint) + outdir = os.path.join(workdir, "pred") + subprocess.run( + ["python", "-m", "deepreg.main", "test", + "-checkpoint_path", checkpoint, "-config_path", cfg, "-out_dir", outdir], + check=True, capture_output=True, text=True, cwd=workdir) + # predicted deformation field (voxels) on the fixed (T1c) grid + pred = _find_single_nii(os.path.join(outdir)) + if pred is None: + raise PipelineUnavailable("deepreg: no predicted deformation field found") + disp = sitk.GetArrayFromImage(sitk.ReadImage(pred)) # (Z, Y, X, 3) + T = _displacement_to_sitk(disp, t1c) + meta = { + "pipeline": "deepreg", "toolkit": "Oxford DeepReg (CycleGAN)", + "metric": "CycleGAN + registration loss (pretrained)", + "transform": "3D displacement field", "convention": "T1c->CT (verified)", + "auto_inverted": False, "checkpoint": checkpoint, + } + return T, meta + + +def _deepreg_config(workdir, ct, t1c, checkpoint): + """Minimal DeepReg test config: moving=CT, fixed=T1c, 3D, single channel.""" + from deepreg.registry import registry # noqa + cfg = { + "general": {"data_format": "nifti", "device": "cuda:0"}, + "data": {"moving": os.path.join(workdir, "moving"), + "fixed": os.path.join(workdir, "fixed"), + "field": "deformation", "moving_fields": "t1c", "fixed_fields": "t1c", + "mask": None, "resample_size": "30 186 30"}, + "validation": {}, + "training": {"epoch": 1, "batch_size": 1, + "moving_fields": "t1c", "fixed_fields": "t1c", + "resample_size": "30 186 30", "resample_spacing": "1 1 1"}, + "testing": {"moving": os.path.join(workdir, "moving"), + "fixed": os.path.join(workdir, "fixed"), + "resample_size": "30 186 30", "resample_spacing": "1 1 1"}, + } + p = os.path.join(workdir, "deepreg_test.yaml") + try: + import yaml + with open(p, "w") as f: + yaml.safe_dump(cfg, f) + except Exception: + with open(p, "w") as f: + json.dump(cfg, f) + return p + + +def _find_single_nii(dirpath): + if not os.path.isdir(dirpath): + return None + ns = [os.path.join(dirpath, x) for x in sorted(os.listdir(dirpath)) + if x.endswith(".nii.gz") or x.endswith(".nii")] + return ns[0] if ns else None + + +def _displacement_to_sitk(disp, ref): + """Build a sitk.DisplacementFieldTransform from a (Z,Y,X,3) voxel field.""" + from scipy.ndimage import map_coordinates + zz, yy, xx = np.mgrid[0:disp.shape[0], 0:disp.shape[1], 0:disp.shape[2]] + # disp[i,j,k] = (dz, dy, dx) to add at grid point (i,j,k) to reach the source + src_z = zz + disp[..., 0]; src_y = yy + disp[..., 1]; src_x = xx + disp[..., 2] + # DisplacementFieldTransform stores, per output voxel, the offset to the input + off = sitk.GetImageFromArray(disp[..., ::-1].astype(np.float32)) # (X,Y,Z,3) + off.CopyInformation(ref) + t = sitk.DisplacementFieldTransform(off) + return t + + +REGISTERS = { + "simpleitk": register_simpleitk, + "elastix": register_elastix, + "antspyx": register_antspyx, + "deepreg": register_deepreg, +} +ALL_PIPELINES = list(REGISTERS) + + +def availability(): + """pipeline -> (available: bool, note: str) for the current env.""" + out = {} + for name in ALL_PIPELINES: + try: + if name == "simpleitk": + out[name] = (True, f"SimpleITK {sitk.Version_VersionString()}") + elif name == "elastix": + ok = shutil.which("elastix") is not None + out[name] = (ok, "binaries present" if ok else "elastix/transformix missing") + elif name == "antspyx": + try: + import ants # noqa + out[name] = (True, f"ants importable v{getattr(ants, '__version__', '?')}") + except Exception as e: # noqa + out[name] = (False, f"ants(antspyx) not importable ({type(e).__name__})") + elif name == "deepreg": + try: + import deepreg # noqa + out[name] = (True, "deepreg importable") + except Exception as e: # noqa + out[name] = (False, "deepreg not importable") + except Exception as e: # noqa + out[name] = (False, repr(e)) + return out \ No newline at end of file