From f80f22dabefd5c0dacf7b27156428b7dcea71a37 Mon Sep 17 00:00:00 2001 From: Furen Xiao Date: Mon, 28 Sep 2026 23:13:44 +0800 Subject: [PATCH] feat(pipelines): add SynthMorph Docker registration and drop DeepReg Replace the unavailable DeepReg benchmark path with the freesurfer/synthmorph Docker pipeline for rigid CT-to-T1c registration. Add a shared CT head-box resample helper for centered initialization and update benchmark docs and CLI for the new SynthMorph options. --- REGISTRATION_BENCHMARK.md | 189 +++++++++++++------ scripts/bench_registration.py | 43 +++-- src/registration_pipelines.py | 344 +++++++++++++++++++++++----------- 3 files changed, 386 insertions(+), 190 deletions(-) diff --git a/REGISTRATION_BENCHMARK.md b/REGISTRATION_BENCHMARK.md index 67a74da..a2dd5b5 100644 --- a/REGISTRATION_BENCHMARK.md +++ b/REGISTRATION_BENCHMARK.md @@ -22,21 +22,23 @@ 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. +destination→source convention this `T` maps **T1c → CT**. The intensity-metric +pipelines register on the same CT-space head box (MI computed on head tissue, +not neck/air); `synthmorph` uses the full volumes (it normalizes each one to an +isotropic network space internally). All single-init tools start 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 +## The five pipelines | # | pipeline | toolkit | metric | transform | status in this env | |---|----------|---------|--------|-----------|--------------------| -| 1 | `elastix` | ITK-Elastix 0.25.4 (pip `itk` + `itk-elastix`, elastix 5.3.1) | Mattes MI (10 bins) | `EulerTransform` rigid (5.x replacement for `Rigid3DTransform`), ASGD, 3 resolutions | ✅ ran (main env, in-process python API — no CLI) | -| 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 (main env; stochastic MI subsampling) | -| 4 | `deepreg` | Oxford DeepReg (CycleGAN) | CycleGAN + reg. loss (pretrained) | 3D displacement field | **not installed** (see below) | -| 5 | `impact` | Docker `vboussot/elastix_impact` (Elastix 5.2.0 + IMPACT) | **IMPACT semantic** (TotalSegmentator `TS/M730.pt` features, L2, Jacobian/GPU) | `EulerTransform` rigid, ASGD, 4 scales (no image resampling) | ✅ ran (main env, Docker + GPU) | +| 1 | `elastix` | ITK-Elastix 0.25.4 (pip `itk` + `itk-elastix`, elastix 5.3.1) | Mattes MI (10 bins) | `EulerTransform` rigid (5.x replacement for `Rigid3DTransform`), ASGD, 3 resolutions | ✅ ran (main env, in-process python API — no CLI) | +| 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 (main env; stochastic MI subsampling) | +| 4 | `synthmorph` | Docker `freesurfer/synthmorph` (TensorFlow, pretrained) | **SynthMorph** anatomy-aware contrast-invariant (deep, pretrained) | 6-DOF rigid matrix (centroid init) | ✅ ran (main env, Docker + GPU) | +| 5 | `impact` | Docker `vboussot/elastix_impact` (Elastix 5.2.0 + IMPACT) | **IMPACT semantic** (TotalSegmentator `TS/M730.pt` features, L2, Jacobian/GPU) | `EulerTransform` rigid, ASGD, 4 scales (no image resampling) | ✅ ran (main env, Docker + GPU) | `simpleitk` has two budgets: `--simpleitk-mode benchmark` (single pyramid, centered start — the default, comparable to the other tools) and @@ -88,34 +90,36 @@ Production QA, crop-independent: | **simpleitk** | SimpleITK 2.5.6 | 9 | **4/10 (0.40)** | 12.6 / 16.1 | 15.4 / 16.0 | 8.6 | 0.93 | **0.643 / 0.630** | | **antspyx** | ANTsPy 0.6.3 | 9 | **0/10 (0.00)** | 14.9 / 16.5 | 18.4 / 22.7 | 14.5 | 0.93 | **0.598 / 0.529** | | **elastix** | ITK-Elastix 0.25.4 | 9 | **5/10 (0.50)** | 15.3 / 16.1 | 9.9 / 13.8 | 3.2 | 0.93 | **0.562 / 0.549** | +| **synthmorph** | Docker SynthMorph (TensorFlow) | 9 | **1/10 (0.10)** | 19.8 / 25.1 | 23.0 / 24.1 | 28.2 | 0.93 | **0.215 / 0.309** | | **impact** | Docker Elastix 5.2 + IMPACT | 9 | **3/10 (0.30)** | 12.4 / 19.0 | 15.1 / 16.7 | **562.3** | 0.93 | **0.475 / 0.448** | -| deepreg | — | — | — | — | — | — | — | — | *(antspyx and impact are stochastic (random metric/patch subsampling), so their accepted counts are one draw each — see Interpretation. impact's s/reg is ~9.4 -min/case on GPU, orders of magnitude above the intensity-metric tools.)* +min/case on GPU, orders of magnitude above the intensity-metric tools. +synthmorph is deterministic at fixed threads; its low acceptance is a real +cross-modality/tilt weakness, not run-to-run variance.)* Per-case (mid / std, `OK` = accepted): -| case | simpleitk | antspyx | elastix | impact | -|---|---|---|---|---| -| m6_22GD5VL2_20230724 | incomplete-CT | incomplete-CT | incomplete-CT | incomplete-CT | -| m6_22IRRKU6_20240115 | 25.7 / 29.3 | 24.4 / 27.7 | 26.4 / 30.3 | **8.4 / 12.6** | -| m6_23XIMO4M_20230222 | 22.2 / 23.9 | 23.4 / 24.8 | 21.1 / 23.3 | 21.4 / 24.1 | -| m6_24V62QQM_20230223 | 10.7 / 16.9 | 9.5 / 18.4 | 10.6 / 8.2 **OK** | 10.0 / 8.3 **OK** | -| m6_24V62QQM_20231017 | 11.5 / 7.3 **OK** | 14.9 / 15.5 | 11.0 / 6.1 **OK** | 11.1 / 8.8 **OK** | -| m6_24V62QQM_20240820 | 11.6 / 9.6 **OK** | 14.3 / 17.5 | 15.3 / 14.0 | 54.4 / 15.1 | -| m6_26TYQI3C_20221027 | 12.2 / 3.7 **OK** | 17.7 / 37.6 | 11.4 / 3.2 **OK** | 23.9 / 22.6 | -| m6_2BBN5XA4_20250411 | 18.6 / 8.0 **OK** | 14.3 / 12.7 | 18.1 / 9.9 **OK** | 10.1 / 5.9 **OK** | -| m6_2C4TKI6H_20240506 | 12.6 / 15.4 | 13.1 / 14.3 | 12.0 / 6.8 **OK** | 12.4 / 22.3 | -| m6_2CDNN7BH_20240102 | 19.4 / 30.2 | 17.2 / 35.4 | 18.9 / 22.7 | 19.6 / 30.3 | +| case | simpleitk | antspyx | elastix | synthmorph | impact | +|---|---|---|---|---|---| +| m6_22GD5VL2_20230724 | incomplete-CT | incomplete-CT | incomplete-CT | incomplete-CT | incomplete-CT | +| m6_22IRRKU6_20240115 | 25.7 / 29.3 | 24.4 / 27.7 | 26.4 / 30.3 | 20.0 / 25.2 | **8.4 / 12.6** | +| m6_23XIMO4M_20230222 | 22.2 / 23.9 | 23.4 / 24.8 | 21.1 / 23.3 | 19.8 / 23.2 | 21.4 / 24.1 | +| m6_24V62QQM_20230223 | 10.7 / 16.9 | 9.5 / 18.4 | 10.6 / 8.2 **OK** | 12.0 / 23.0 | 10.0 / 8.3 **OK** | +| m6_24V62QQM_20231017 | 11.5 / 7.3 **OK** | 14.9 / 15.5 | 11.0 / 6.1 **OK** | 13.0 / 12.6 | 11.1 / 8.8 **OK** | +| m6_24V62QQM_20240820 | 11.6 / 9.6 **OK** | 14.3 / 17.5 | 15.3 / 14.0 | 53.8 / 20.2 | 54.4 / 15.1 | +| m6_26TYQI3C_20221027 | 12.2 / 3.7 **OK** | 17.7 / 37.6 | 11.4 / 3.2 **OK** | 11.3 / 5.0 **OK** | 23.9 / 22.6 | +| m6_2BBN5XA4_20250411 | 18.6 / 8.0 **OK** | 14.3 / 12.7 | 18.1 / 9.9 **OK** | 49.3 / 58.9 | 10.1 / 5.9 **OK** | +| m6_2C4TKI6H_20240506 | 12.6 / 15.4 | 13.1 / 14.3 | 12.0 / 6.8 **OK** | 18.7 / 19.3 | 12.4 / 22.3 | +| m6_2CDNN7BH_20240102 | 19.4 / 30.2 | 17.2 / 35.4 | 18.9 / 22.7 | 27.6 / 29.1 | 19.6 / 30.3 | ### Interpretation -- **Translation is recovered well by all four** (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 every pipeline, so the target is in the right - place. +- **Translation is recovered well by all five** (mid gaps mostly 6–20 mm, + within the `≤ 20 mm` band, except the grossly-drifted cases below). + `tv-in-CT` is identical (~0.93): the warped GTV lands inside the CT head for + every pipeline, 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 local minima that a single centered start can lock onto). This is exactly why production `02_build_m6_dataset.py` @@ -144,17 +148,31 @@ Per-case (mid / std, `OK` = accepted): translation (mid 54 mm). Net 3/10 accepted — between ANTs and SimpleITK — at ~9.4 min/case, two orders of magnitude slower than the intensity tools. - **Speed** (median wall, 4 cases in parallel, 24 threads): Elastix ~3.2 s - (warm workers; a worker's first case pays a ~10 s warm-up), SimpleITK ~8.6 s, - ANTs ~14.5 s (CPU-bound), **IMPACT ~562 s** (GPU-bound: every ASGD iteration - runs the TotalSegmentator extractor on the sampled patches with backprop). -- **Mattes MI (warped-CT ~ MR) is an independent, intensity-based confirmation**: - medians 0.643 (SimpleITK) / 0.598 (ANTs) / 0.562 (Elastix) / 0.475 (IMPACT), + (warm workers; a worker's first case pays a ~10 s warm-up), SimpleITK ~8.6 s, + ANTs ~14.5 s (CPU-bound), **SynthMorph ~28 s** (GPU: a single pretrained-model + forward pass per case), **IMPACT ~562 s** (GPU-bound: every ASGD iteration + runs the TotalSegmentator extractor on the sampled patches with backprop). + - **Mattes MI (warped-CT ~ MR) is an independent, intensity-based confirmation**: + medians 0.643 (SimpleITK) / 0.598 (ANTs) / 0.562 (Elastix) / 0.475 (IMPACT) / + 0.215 (SynthMorph), and it diverges sharply on the tilted cases — ANTs `m6_26TYQI3C` drops to ~0.1 nats (rotation ~38 mm off) while SimpleITK holds 0.670. A high MI means the CT skull/tissue intensities fall on the matching MR structures after the inverse transform. IMPACT's lower median MI is consistent with its less - reliable single-init convergence. -- **Stochasticity**: ANTs' Mattes-MI subsamples the image at random, so its +reliable single-init convergence. + - **SynthMorph (deep, pretrained) is the weakest of the five on these + cross-modality CT pairs.** Its rigid model is trained on brain *MRI* + (T1/T2, anatomy-aware but intensity-domain MRI), so it recovers gross + translation well (mid 12–20 mm on most cases) but is repeatedly defeated by + the ~30° inter-session tilt: 6 of 9 cases have `std > 20 mm` (rotation error) + and two drift grossly (`m6_2BBN5XA4` mid 49 / std 59, `m6_24V62QQM_20240820` + mid 54). It accepts only the clean `m6_26TYQI3C` (mid 11.3 / std 5.0, MI + 0.64 — where it is competitive with the intensity tools). Its low median MI + (0.215) reflects the tilt failures. On same-modality MRI it would likely + rank much higher; this CT benchmark is an out-of-distribution stress test for + it. It is deterministic (one fixed pretrained model, no random subsampling), + so unlike ANTs/IMPACT the per-case numbers are stable run to run. + - **Stochasticity**: ANTs' Mattes-MI subsamples the image at random, so its per-case numbers carry run-to-run variance (this table is one draw); SimpleITK's MI (and registration) is deterministic at fixed thread count (RSGD wall time is thread-sensitive, convergence is not); the in-process @@ -167,14 +185,15 @@ Per-case (mid / std, `OK` = accepted): ## Reproducing ```bash -# main env — SimpleITK + Elastix + ANTsPy (all pip-installed here); deepreg -# reports unavailable and is recorded as such +# main env — SimpleITK + Elastix + ANTsPy (all pip-installed here) source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 -python scripts/bench_registration.py --n-cases 10 \ - --workers 4 --threads 24 --simpleitk-mode benchmark +python scripts/bench_registration.py --pipelines simpleitk elastix antspyx \ + --n-cases 10 --workers 4 --threads 24 --simpleitk-mode benchmark -# IMPACT (needs Docker + the pulled image + a model; see the Impact section -# below for the one-time setup) +# Docker pipelines (need the pulled image; see the SynthMorph / IMPACT +# sections below for the one-time setup) +python scripts/bench_registration.py --pipelines synthmorph --n-cases 10 \ + --workers 4 --threads 24 python scripts/bench_registration.py --pipelines impact --n-cases 10 \ --workers 4 --threads 24 @@ -185,7 +204,8 @@ 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), `--impact-model `, +`--synthmorph-model {rigid,affine}`, `--synthmorph-gpu `, +`--synthmorph-extent `, `--impact-model `, `--impact-gpu `, `--impact-iters N`, `--impact-samples N`, `--out-dir `, `--qa-images` / `--qa-out ` (visual QA, below). @@ -216,27 +236,55 @@ visuals agree. Generate for the available pipelines with: source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 python scripts/bench_registration.py \ --pipelines simpleitk elastix antspyx --n-cases 10 --qa-images +python scripts/bench_registration.py \ + --pipelines synthmorph --n-cases 10 --qa-images python scripts/bench_registration.py \ --pipelines impact --n-cases 10 --qa-images ``` (9 images each; the incomplete-CT case is skipped.) -## Enabling the blocked pipeline +## SynthMorph pipeline (Docker) — one-time setup -**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: +The `synthmorph` pipeline runs the [`freesurfer/synthmorph`](https://hub.docker.com/r/freesurfer/synthmorph) +image: [SynthMorph](https://synthmorph.io) (Hoffmann et al., IEEE TMI 2022), an +anatomy-aware, acquisition-agnostic **deep** registration (TensorFlow +`VxmAffineFeatureDetector` with the pretrained `synthmorph.rigid.1.h5` weights). +No skull-stripping or intensity preprocessing is needed: each volume is +normalized to an isotropic 1-mm LIA network space internally. It is the +second Docker-based pipeline and the only deep-learning one. + +Prerequisites (check with `RP.availability()`; a missing piece raises +`PipelineUnavailable` with the exact command): ```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 +# 1. Docker (with NVIDIA container toolkit so --gpus works) +docker --version + +# 2. Pull the image (~11 GB, one-time) +docker pull freesurfer/synthmorph ``` -The adapter runs `deepreg test` to predict a voxel displacement field, converts -it to a SimpleITK `DisplacementFieldTransform`, and scores it like the others. +No extra model assets: the rigid weights ship inside the image +(`/freesurfer/models`). + +### How the adapter drives the image + +1. **Shared init** — the `centroid_offset` (CT→T1c head translation) is built + *inside* the container as an initial LTA (`-i`), because surfa requires the + init transform's embedded src/dst geometries to match the input images + exactly. The container also writes the LTA as the host uid so the scratch + removes cleanly. +2. **Rigid matrix model** — `register -m rigid -g` on GPU when one is visible + (auto-spread over the host GPUs by pid; `--synthmorph-gpu -1` forces CPU). +3. **Read-back** — the `-t` LTA (a `LINEAR_RAS_TO_RAS` 4×4 in shared world + space) is parsed into a SimpleITK `AffineTransform` and run through the + usual `fix_orientation` gate. +4. **Image-build workarounds** — this image's `-i` code path has a + float64/float32 bug (the dense 256³ mesh is built as `double` and the final + matrix product mismatches dtypes), so the adapter injects a small wrapper + that casts the transform operands to float32, and sets + `TF_FORCE_GPU_ALLOW_GROWTH` so concurrent workers can share a GPU. ## IMPACT pipeline (Docker) — one-time setup @@ -299,15 +347,36 @@ to add the shared centered start: 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. -- `deepreg` could not be executed here, so its row is "not installed", not - measured; its code path is implemented and unit-verified for import/convention. -- `fix_orientation` (elastix/ants/deepreg/impact) chooses between a transform +- `synthmorph` is deterministic (one fixed pretrained model, no random + subsampling), so its per-case numbers are stable run to run — unlike ANTs / + IMPACT. Its low acceptance (1/10) is a genuine cross-modality / tilt + weakness, not variance. +- `fix_orientation` (elastix/ants/synthmorph/impact) 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. In - practice it never fired (all `auto_inverted: False`): elastix/ants/impact - returned the T1c→CT direction as expected. + the result independent of each toolkit's transform-direction convention. It + rarely fired: elastix/ants/impact returned the T1c→CT direction as expected + (`auto_inverted: False`) and synthmorph did so on 8/9 cases (its one +`True` was the grossly-drifted `m6_2BBN5XA4`, where the gate flipped a bad + pose). +- **SynthMorph specifics** (Docker `freesurfer/synthmorph`, TensorFlow, GPU): + - **Cross-modality / tilt weakness.** The `rigid` model is trained on brain + *MRI*, so on these CT↔T1c pairs it recovers translation but is repeatedly + defeated by the ~30° inter-session tilt (6/9 cases have `std > 20 mm`). + It is a legitimate data point for how far a pretrained MRI deep model + generalizes out-of-domain to CT — not a bug. + - **Shared init is built in-container.** surfa rejects an `-i` transform + whose src/dst geometries don't match the inputs (tol 1e-3), so the + adapter has the container write the centroid-translation LTA from the + mounted images themselves. + - **Image-build dtype/GPU workarounds.** This image's `-i` path builds the + dense 256³ mesh as float64 (2× memory → GPU OOM under contention) and + mismatches dtypes in the final matrix product; the injected wrapper casts + transform operands to float32, and `TF_FORCE_GPU_ALLOW_GROWTH` lets + concurrent workers share a GPU. + - **~28 s/case** (GPU: one pretrained forward pass), between ANTs (~15 s) + and IMPACT (~562 s). - **Elastix specifics** (ITK-Elastix 0.25.4, elastix 5.3.1, in-process python - API — no CLI binary involved): + API — no CLI binary involved): - Elastix 5.x renamed components vs the classic parameter files: the rigid 3D transform is `EulerTransform` (6-DOF, `Rigid3DTransform` is gone from the component set), the metric key is `Metric` (not `ImageMetric`), the diff --git a/scripts/bench_registration.py b/scripts/bench_registration.py index 3f1d837..3b2313b 100644 --- a/scripts/bench_registration.py +++ b/scripts/bench_registration.py @@ -1,7 +1,7 @@ """Benchmark the 5 rigid CT->T1c registration pipelines on N M6 test cases. Pipelines (see src/registration_pipelines.py for the shared transform contract): - simpleitk | elastix | antspyx | deepreg | impact + simpleitk | elastix | antspyx | synthmorph | impact Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI (cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto @@ -19,22 +19,22 @@ geometric QA used in production (scripts/02_build_m6_dataset.py): Because some toolkits are not pip-installable into the main `longitudinal-313` env (antspyx is kept in an isolated env, see -REGISTRATION_BENCHMARK.md; PyPI deepreg is a broken stub), the benchmark is -env-portable: run it under whichever conda env has a given toolkit installed, -and it records only the pipelines that are importable there. elastix runs in the main env through the ITK-Elastix -python API (pip ``itk`` + ``itk-elastix``). ``impact`` runs in the -``vboussot/elastix_impact`` Docker image (needs Docker + the image + a model, -see REGISTRATION_BENCHMARK.md). The `--merge` mode then combines -the per-pipeline JSON files (accumulated across envs) into one summary + plot. +REGISTRATION_BENCHMARK.md), the benchmark is env-portable: run it under +whichever conda env has a given toolkit installed, and it records only the +pipelines that are importable there. elastix runs in the main env through the +ITK-Elastix python API (pip ``itk`` + ``itk-elastix``). ``synthmorph`` and +``impact`` run in Docker images (``freesurfer/synthmorph`` and +``vboussot/elastix_impact`` respectively; needs Docker + the pulled image, see +REGISTRATION_BENCHMARK.md). The `--merge` mode then combines the +per-pipeline JSON files (accumulated across envs) into one summary + plot. Usage (repo root): - # main env (SimpleITK + elastix + antspyx; impact needs Docker + the image; - # deepreg unavailable) + # main env (SimpleITK + elastix + antspyx) source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 python scripts/bench_registration.py --pipelines simpleitk elastix antspyx - # impact (GPU, Docker) — one-time setup in REGISTRATION_BENCHMARK.md - python scripts/bench_registration.py --pipelines impact + # synthmorph / impact (GPU, Docker) — one-time setup in REGISTRATION_BENCHMARK.md + python scripts/bench_registration.py --pipelines synthmorph impact # any pipeline that is not importable in the current env is recorded as # unavailable, so the script is env-portable; --merge combines the @@ -225,9 +225,10 @@ def _call(pipeline, ct, t1c, tv, args): return RP.register_elastix(ct, t1c, workdir=args.elastix_dir, threads=args.threads) if pipeline == "antspyx": return RP.register_antspyx(ct, t1c) - if pipeline == "deepreg": - return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir, - checkpoint=args.deepreg_ckpt) + if pipeline == "synthmorph": + return RP.register_synthmorph(ct, t1c, model=args.synthmorph_model, + gpu=args.synthmorph_gpu, threads=args.threads, + extent=args.synthmorph_extent) if pipeline == "impact": return RP.register_impact(ct, t1c, model=args.impact_model, gpu=args.impact_gpu, iters=args.impact_iters, samples=args.impact_samples) @@ -476,8 +477,15 @@ def parse_args(): "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("--synthmorph-model", choices=["rigid", "affine"], default="rigid", + help="SynthMorph matrix model (default rigid; deformable/joint are " + "outside the rigid-transform contract)") + ap.add_argument("--synthmorph-gpu", type=int, default=None, + help="GPU device index for SynthMorph (default: auto-spread over " + "the host GPUs by pid; -1 forces CPU)") + ap.add_argument("--synthmorph-extent", type=int, default=256, + help="SynthMorph isotropic registration-space extent in mm " + "(192 or 256; default 256)") ap.add_argument("--impact-model", default=None, help="IMPACT TorchScript model, path relative to the HF " f"layout (default {RP.IMPACT_DEFAULT_MODEL}) or an absolute host path") @@ -506,7 +514,6 @@ def main(): merge(args) return args.elastix_dir = args.elastix_dir or d("data/m6reg/elastix") - args.deepreg_dir = args.deepreg_dir or d("data/m6reg/deepreg") if args.qa_images: args.qa_out = os.path.abspath(args.qa_out or os.path.join(args.out_dir, "qa")) os.makedirs(args.qa_out, exist_ok=True) diff --git a/src/registration_pipelines.py b/src/registration_pipelines.py index 344187a..e6163c5 100644 --- a/src/registration_pipelines.py +++ b/src/registration_pipelines.py @@ -29,15 +29,17 @@ Pipelines - ``register_elastix`` : ITK-Elastix Python API (pip ``itk`` + ``itk-elastix``, elastix 5.x: ``EulerTransform`` rigid + Mattes MI, ASGD). - ``register_antspyx`` : ANTs / ANTsPy ``ants.registration`` (Rigid). -- ``register_deepreg`` : Oxford DeepReg (CycleGAN deformation-field - registration; needs a pretrained checkpoint, see module docstring). +- ``register_synthmorph``: FreeSurfer SynthMorph (TensorFlow, pretrained + anatomy-aware contrast-invariant model) inside the ``freesurfer/synthmorph`` + Docker image. Rigid 6-DOF matrix model, GPU when available. Needs Docker + + the pulled image. - ``register_impact`` : IMPACT (semantic-similarity metric) inside the ``vboussot/elastix_impact`` Docker image (Elastix 5.2 + a TorchScript TotalSegmentator feature extractor as a native Elastix metric). Rigid 6-DOF, GPU (Jacobian mode). Needs Docker + the image + a model asset. -A toolkit that is not installed (or, for DeepReg, missing its checkpoint, or, -for IMPACT, missing its Docker image / model) +A toolkit that is not installed (or, for the Docker-based pipelines, missing +its image) 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 @@ -50,11 +52,10 @@ 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 / impact, not to the trusted production -SimpleITK path. +applied to elastix / antspyx / synthmorph / impact, not to the trusted +production SimpleITK path. """ import os -import json import shutil import subprocess import importlib.util @@ -161,6 +162,22 @@ def _ct_head_box_mm(ct): return box_lo, box_hi +def _ct_head_box_resample(ct, t1c): + """Resample the CT into the shared CT-space head box (mirrors the other + pipelines), or None when no head box is found. Used only to get the CT + *head* tissue centroid for the shared centered start.""" + lo, hi = _ct_head_box_mm(ct) + if lo is None: + return None + 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)) + return sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) + + # --------------------------------------------------------------------------- # # 2. SimpleITK (production baseline -- reused verbatim) # --------------------------------------------------------------------------- # @@ -484,108 +501,211 @@ def sys_pyver(): # --------------------------------------------------------------------------- # -# 4. DeepReg (Oxford) +# 4. SynthMorph (Docker image freesurfer/synthmorph) # --------------------------------------------------------------------------- # -def register_deepreg(ct, t1c, workdir=None, checkpoint=None, config=None): - """Oxford DeepReg deformation-field registration (CT -> T1c). +SYNTHMORPH_IMAGE = "freesurfer/synthmorph" - 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 +# The image's `-i` (initial-matrix) code path has a float64/float32 bug in this +# build: it computes `net_to_mov = net_to_mov @ tf.linalg.inv(init)` where `init` +# is float64, so `net_to_mov` becomes a TF *float64* tensor. Downstream, (a) +# voxelmorph's compose() only casts non-tensor inputs to float32, so the final +# matrix product fails ("expected a double tensor but is a float tensor"), and +# (b) affine_to_dense_shift() builds the 256^3 dense mesh as float64 (2x memory), +# which OOMs a GPU under contention. This wrapper monkey-patches tf.linalg.inv +# to return float32 (the root cause) and compose() to cast every operand to +# float32 (a safety net), then drives the mri_synthmorph CLI unchanged. It is +# written into the per-case scratch dir and run inside the container. +_SYNTHMORPH_WRAPPER = '''\ +import os +os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" # let concurrent workers share a GPU +import sys +sys.path.insert(0, "/freesurfer") +import tensorflow as tf +import voxelmorph as vxm - -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) +def _c32(t): + if not tf.is_tensor(t): + return tf.constant(t, dtype=tf.float32) + if t.dtype != tf.float32: + return tf.cast(t, tf.float32) return t +_orig_inv = tf.linalg.inv +def inv(x, *a, **k): + return _orig_inv(_c32(x), *a, **k) +tf.linalg.inv = inv + +_orig_compose = vxm.utils.compose +def compose(transforms, *a, **k): + return _orig_compose(tuple(_c32(t) for t in transforms), *a, **k) +vxm.utils.compose = compose + +sys.argv = ["mri_synthmorph"] + sys.argv[1:] +import runpy +runpy.run_path("/freesurfer/mri_synthmorph", run_name="__main__") +''' + +# surfa refuses an init transform whose embedded src/dst geometries do not +# match the input images (tol 1e-3), and reproducing that geometry on the host +# is fragile. So the container builds the -i initial LTA itself from the +# mounted images: a pure CT->T1c world translation (the shared centroid start) +# with source=moving, target=fixed. The offset is passed via SYNTH_INIT_T. +_SYNTHMORPH_BUILD_INIT = '''\ +import os +import numpy as np +import surfa as sf +mov = sf.load_volume("/data/moving.nii.gz") +fix = sf.load_volume("/data/fixed.nii.gz") +t = [float(x) for x in os.environ["SYNTH_INIT_T"].split()] +M = np.eye(4) +M[:3, 3] = t +sf.Affine(M, source=mov, target=fix, space="world").save("/data/init.lta") +''' + + +def _synthmorph_docker_ok(): + """Return (ok, note): is the Docker CLI + the synthmorph image available?""" + if shutil.which("docker") is None: + return False, "docker CLI not found on PATH" + try: + r = subprocess.run(["docker", "image", "inspect", SYNTHMORPH_IMAGE], + capture_output=True, text=True, timeout=30) + if r.returncode != 0: + return False, (f"docker image {SYNTHMORPH_IMAGE} not pulled " + f"(run: docker pull {SYNTHMORPH_IMAGE})") + except Exception as e: # noqa + return False, f"docker image check failed ({type(e).__name__}: {e})" + return True, SYNTHMORPH_IMAGE + + +def _lta_to_sitk(path): + """Parse a FreeSurfer LTA (LINEAR_RAS_TO_RAS) into a SimpleITK AffineTransform. + + The 4x4 matrix block lives in the shared RAS world space. Both inputs are + RAS-oriented by ``read_nii_img`` (DICOMOrient), so RAS == SimpleITK's + physical frame and the matrix applies directly to world coordinates.""" + lines = [l for l in open(path).read().splitlines() + if l and not l.startswith("#")] + i = lines.index("1 4 4") + M = np.array([[float(x) for x in lines[i + 1 + r].split()] for r in range(4)]) + T = sitk.AffineTransform(3) + T.SetMatrix(M[:3, :3].ravel()) + T.SetTranslation(M[:3, 3]) + return T + + +def register_synthmorph(ct, t1c, model="rigid", gpu=None, threads=None, + extent=256, keep_scratch=False): + """SynthMorph CT->T1c (moving=CT, fixed=T1c) in the Docker image + ``freesurfer/synthmorph`` (TensorFlow ``VxmAffineFeatureDetector``, + pretrained ``synthmorph.rigid.1.h5`` weights). + + SynthMorph (Hoffmann et al., IEEE TMI 2022) is an anatomy-aware, + acquisition-agnostic deep registration: each volume is normalized to an + isotropic 1-mm LIA network space and the (pretrained) network estimates the + transform, so no skull-stripping or intensity preprocessing is needed. To + stay on the benchmark's equal footing the pipeline + - starts from the *shared* centroid-aligned translation (``centroid_offset``, + built inside the container as an initial LTA -- see + ``_SYNTHMORPH_BUILD_INIT`` -- because surfa requires the init transform's + geometries to match the inputs exactly), + - runs the ``rigid`` (6-DOF) matrix model, on GPU when one is visible + (auto-spread over the host GPUs by pid; ``gpu=-1`` forces CPU), + - reads back the ``-t`` LTA. In RAS world space that matrix is the + T1c->CT transform the unified contract requires (its ``-T`` sibling is + the inverse); :func:`fix_orientation` still verifies the direction. + + Preconditions (else :class:`PipelineUnavailable`): the ``docker`` CLI and + the ``freesurfer/synthmorph`` image. Only the matrix models (``rigid``, + ``affine``) are exposed: the deformable/joint models emit displacement + fields, outside the rigid-transform contract. + """ + import time + + ok, note = _synthmorph_docker_ok() + if not ok: + raise PipelineUnavailable(f"synthmorph: {note}") + if model not in ("rigid", "affine"): + raise PipelineUnavailable( + f"synthmorph: only the matrix models (rigid/affine) fit the unified " + f"rigid-transform contract, got {model!r}") + + # GPU: auto-spread across the host GPUs by pid, or forced value + if gpu is None: + ngpu = _n_gpus() + gpu = (os.getpid() % ngpu) if ngpu else -1 + + # shared centered start: the CT->T1c head-centroid translation. As with the + # other single-init pipelines, resample the CT to the head box so the tissue + # centroid is the CT *head* (not head+neck). + box = _ct_head_box_resample(ct, t1c) + off = centroid_offset(ct, t1c, ct_box=box) or (0.0, 0.0, 0.0) + + # per-case scratch on a local (non-NFS) tmp fs + tmp = tempfile.mkdtemp(prefix="synthmorph_reg_") + try: + sitk.WriteImage(sitk.Cast(ct, sitk.sitkFloat32), + os.path.join(tmp, "moving.nii.gz"), True) # CT + sitk.WriteImage(sitk.Cast(t1c, sitk.sitkFloat32), + os.path.join(tmp, "fixed.nii.gz"), True) # T1c + with open(os.path.join(tmp, "build_init.py"), "w") as f: + f.write(_SYNTHMORPH_BUILD_INIT) + with open(os.path.join(tmp, "run_synthmorph.py"), "w") as f: + f.write(_SYNTHMORPH_WRAPPER) + + # run as the host uid/gid so the container-written LTA is host-owned + # (clean scratch removal) and TF still reaches the GPU. + cmd = ["docker", "run", "--rm", + "--user", f"{os.getuid()}:{os.getgid()}"] + if gpu >= 0: + cmd += ["--gpus", "all", "-e", f"CUDA_VISIBLE_DEVICES={gpu}"] + gpu_flag = " -g" if gpu >= 0 else "" + thread_flag = f" -j {threads}" if threads else "" + cmd += [ + "-e", f"SYNTH_INIT_T={off[0]:.6f} {off[1]:.6f} {off[2]:.6f}", + "-v", f"{tmp}:/data", + "--entrypoint", "sh", + SYNTHMORPH_IMAGE, + "-c", ("python3 /data/build_init.py && " + "python3 /data/run_synthmorph.py register " + f"-m {model}{gpu_flag}{thread_flag} -e {extent} " + "-i /data/init.lta -t /data/t.lta " + "/data/moving.nii.gz /data/fixed.nii.gz"), + ] + t0 = time.time() + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=7200) + dt = time.time() - t0 + if proc.returncode != 0: + tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:]) + raise RuntimeError(f"synthmorph: docker failed (rc={proc.returncode}):\n{tail}") + if not os.path.exists(os.path.join(tmp, "t.lta")): + tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:]) + raise RuntimeError(f"synthmorph: no /data/t.lta produced:\n{tail}") + + T = _lta_to_sitk(os.path.join(tmp, "t.lta")) + T, inverted = fix_orientation(ct, t1c, T) + finally: + if not keep_scratch: + shutil.rmtree(tmp, ignore_errors=True) + + meta = { + "pipeline": "synthmorph", + "toolkit": f"Docker {SYNTHMORPH_IMAGE} (TensorFlow, pretrained)", + "metric": "SynthMorph anatomy-aware contrast-invariant (deep, pretrained)", + "transform": (f"SynthMorph {model} (centroid init; " + f"device={'GPU' + str(gpu) if gpu >= 0 else 'CPU'}, " + f"extent {extent})"), + "convention": "T1c->CT (verified)", "auto_inverted": inverted, + "fixed": "T1c", "moving": "CT", + "seconds_reg": round(dt, 1), + "docker_image": SYNTHMORPH_IMAGE, "model": model, "gpu": gpu, + "extent": extent, + } + if keep_scratch: + meta["scratch"] = tmp + return T, meta + # --------------------------------------------------------------------------- # # 5. IMPACT (Docker image vboussot/elastix_impact) @@ -618,7 +738,7 @@ def _impact_docker_ok(): _NGPUS_CACHE = None -def _impact_n_gpus(): +def _n_gpus(): global _NGPUS_CACHE if _NGPUS_CACHE is None: n = 0 @@ -809,7 +929,7 @@ def register_impact(ct, t1c, model=None, gpu=None, iters=500, samples=2000, # GPU: auto-spread across the host GPUs by pid, or forced value if gpu is None: - ngpu = _impact_n_gpus() + ngpu = _n_gpus() gpu = (os.getpid() % ngpu) if ngpu else -1 # per-level feature-extraction scales (recommended 6/3/1.5/1 mm), scaled so @@ -916,7 +1036,7 @@ REGISTERS = { "simpleitk": register_simpleitk, "elastix": register_elastix, "antspyx": register_antspyx, - "deepreg": register_deepreg, + "synthmorph": register_synthmorph, "impact": register_impact, } ALL_PIPELINES = list(REGISTERS) @@ -943,12 +1063,12 @@ def availability(): 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") + elif name == "synthmorph": + ok, note = _synthmorph_docker_ok() + if not ok: + out[name] = (False, note) + else: + out[name] = (True, f"Docker {SYNTHMORPH_IMAGE}") elif name == "impact": ok, note = _impact_docker_ok() if not ok: