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.
This commit is contained in:
parent
6b85e950e5
commit
f80f22dabe
3 changed files with 386 additions and 190 deletions
|
|
@ -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 <rel|abs>`,
|
||||
`--synthmorph-model {rigid,affine}`, `--synthmorph-gpu <int|-1>`,
|
||||
`--synthmorph-extent <int>`, `--impact-model <rel|abs>`,
|
||||
`--impact-gpu <int|-1>`, `--impact-iters N`, `--impact-samples N`,
|
||||
`--out-dir <dir>`, `--qa-images` / `--qa-out <dir>` (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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue