diff --git a/AGENTS.md b/AGENTS.md index d90fb4c..cebb331 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,7 +10,7 @@ Activate with: source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 ``` -Key packages: torch 2.14 (+cu126), torchvision, monai 1.6 (pip, Pipeline C), numpy, scipy, pandas, scikit-learn, scikit-image, matplotlib, nibabel, SimpleITK. Registration toolkits (benchmark, see REGISTRATION_BENCHMARK.md): `itk` + `itk-elastix`, `antspyx`, niftyreg 1.5.69 (`conda install -c conda-forge niftyreg`, provides `reg_aladin`), plus Docker images `freesurfer/synthmorph` and `vboussot/elastix_impact`. +Key packages: torch 2.14 (+cu126), torchvision, monai 1.6 (pip, Pipeline C), numpy, scipy, pandas, scikit-learn, scikit-image, matplotlib, nibabel, SimpleITK. Registration toolkits (benchmark, see REGISTRATION_BENCHMARK.md): `itk` + `itk-elastix`, `antspyx`, niftyreg 1.5.69 (`conda install -c conda-forge niftyreg`, provides `reg_aladin`), `impact-reg-konfai` (pip; IMPACT — on this host the elastix binary needs the one-time shim from `scripts/setup_impact_konfai.sh`), plus the Docker image `freesurfer/synthmorph`. ## Project diff --git a/README.md b/README.md index e02f9a8..279f5b6 100644 --- a/README.md +++ b/README.md @@ -44,8 +44,6 @@ data/ pseudo/roundK/ in-house pseudo-labels (rows.jsonl, masks, summary.json) pseudo_nnu/roundK/ nnU-Net pseudo-labels (rows.jsonl, masks, summary.json) pseudo_monai/roundK/ MONAI pseudo-labels (rows.jsonl, masks, summary.json) - impact_models/ IMPACT TorchScript feature-extractor cache (registration - `impact` pipeline, see REGISTRATION_BENCHMARK.md) src/ U-Net, dataset, losses, training/eval helpers scripts/ Pipeline A (in-house 3D U-Net) scripts_nnu/ Pipeline B (nnU-Net) diff --git a/REGISTRATION_BENCHMARK.md b/REGISTRATION_BENCHMARK.md index 565f6a1..f8555b8 100644 --- a/REGISTRATION_BENCHMARK.md +++ b/REGISTRATION_BENCHMARK.md @@ -1,8 +1,9 @@ # Registration Pipeline Benchmark — CT → T1c (M6) -Benchmarks **six 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 +Benchmarks **six cross-modality registration pipelines** (five rigid; +`impact` is the native deformable B-spline preset of the ImpactReg package) +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°. @@ -28,8 +29,12 @@ every pipeline: - **No mask, no crop / pad**: registration runs on the full input volumes exactly as given (no `FixedMask` / `MovingMask`, no head-box resampling, no padding). -- **CT is always the fixed image** (moving = T1c); when both inputs are MR - the first argument is kept as fixed. +- **The T1c (MR) is always the fixed image** (moving = CT) for every + pipeline. `impact`'s ImpactReg **MR_CT** convention coincides with this + rule: its output displacement field (defined on the fixed grid) is the + contract transform T1c→CT directly — see the IMPACT section. (A CT clipped + to a non-negative window reads as MR, so the T1c slot decides; both-CT + pairs keep the given order.) All pipelines also use each toolkit's **native automatic transform initialization** (no manual centroid offset): elastix / impact use @@ -37,13 +42,11 @@ initialization** (no manual centroid offset): elastix / impact use center-of-mass alignment, simpleitk starts from identity, niftyreg (reg_aladin) initializes from the image centres, synthmorph needs none. One consequence of the no-mask rule: the IMPACT metric only evaluates samples that -map inside the moving buffer, so with the CT (head-neck, ~300×300×251 mm) as -the fixed image and the tight T1c head crop (~220×220×150 mm) as moving, most -of the CT buffer maps outside the T1c buffer (~14% of samples in-buffer at the -auto-init). Elastix's sample-validity gate (build default -`RequiredRatioOfValidSamples 0.25`) used to abort the M6 pairs on that ("too -many samples map outside moving image buffer"); the parameter map now sets -`RequiredRatioOfValidSamples 0.1` so the FOV-mismatched runs proceed — see the +map inside the moving buffer. With the ImpactReg MR_CT direction the samples +are drawn from the tight T1c head crop (fixed, ~220×220×150 mm) and map mostly +inside the large CT head-neck FOV (moving, ~300×300×251 mm), so the gate +rarely fires; the pipeline still sets `RequiredRatioOfValidSamples 0.1` (vs +the build default 0.25) as a safety override for the FOV mismatch — see the IMPACT section. The comparison is on each toolkit's native optimizer / metric / global-search capability. @@ -55,25 +58,33 @@ initialization** (no manual centroid offset): elastix / impact use elastix's `AutomaticTransformInitialization` (CenterOfGravity), antspyx uses its default center-of-mass alignment, simpleitk starts from identity, niftyreg initializes from the image centres, and synthmorph needs none (global deep -method). **All use fixed = CT, moving = T1c** (shared rule above; when both -inputs are MR the first argument is fixed). `niftyreg` and `synthmorph` -clip the CT to [0, 80] HU before registration (per the CT-clipping A/B -below; `clip_ct=None` disables it). +method). **All use fixed = T1c, moving = CT** (shared rule above; +`impact`'s ImpactReg MR_CT convention coincides with it). `niftyreg` and +`synthmorph` clip the CT (the moving image) to [0, 80] HU before +registration (per the CT-clipping A/B below; `clip_ct=None` disables it). | # | pipeline | fixed / moving | init | metric | transform | status in this env | |---|----------|----------------|------|--------|-----------|--------------------| -| 1 | `antspyx` | CT / T1c | auto COM | Mattes MI (multi-res) | ANTs `Rigid` | ✅ ran (main env; stochastic MI subsampling) | -| 2 | `elastix` | CT / T1c | auto CenterOfGravity | Mattes MI (32 bins) | `EulerTransform` rigid (5.x), ASGD, 3 resolutions | ✅ ran (main env, in-process python API — no CLI) | -| 3 | `impact` | **CT / T1c** | auto CenterOfGravity | **IMPACT semantic** (MRSegmentator `MRSeg/MRSeg.pt`, L2, Jacobian/GPU) | `EulerTransform` rigid, ASGD, 4 scales | ✅ ran (Docker + GPU; sample gate 0.1 — see IMPACT section) | -| 4 | `niftyreg` | CT / T1c | image centres | NCC block matching + trimmed least squares (symmetric) | `reg_aladin -rigOnly`, 3-level pyramid | ✅ ran (Docker `niftyreg:2.1.1-ubuntu`, GPU `-platf 1`) | -| 5 | `simpleitk` | CT / T1c | identity | Mattes MI (32 bins) | `Euler3D` rigid, pyramid `[4,2,1]` | ✅ ran | -| 6 | `synthmorph` | CT / T1c | none (global) | **SynthMorph** anatomy-aware contrast-invariant (deep, pretrained) | 6-DOF rigid matrix | ✅ ran (main env, Docker + GPU) | +| 1 | `antspyx` | T1c / CT | auto COM | Mattes MI (multi-res) | ANTs `Rigid` | ✅ ran (main env; stochastic MI subsampling) | +| 2 | `elastix` | T1c / CT | auto CenterOfGravity | Mattes MI (32 bins) | `EulerTransform` rigid (5.x), ASGD, 3 resolutions | ✅ ran (main env, in-process python API — no CLI) | +| 3 | `impact` | **T1c / CT** (MR_CT) | auto CenterOfGravity (override) | **IMPACT semantic** (MIND `R1D2_3D` L1 + TotalSegmentator `M730` deepest-layer Dice, per-level weighted, Static) + Mattes MI (w 0.4) + bending energy | B-spline DVF (preset native), ASGD, 4 scales | ✅ ran (impact-reg-konfai + GPU; `Elastix_IMPACT_Static` preset — see IMPACT section) | +| 4 | `niftyreg` | T1c / CT | image centres | NCC block matching + trimmed least squares (symmetric) | `reg_aladin -rigOnly`, 3-level pyramid | ✅ ran (Docker `niftyreg:2.1.1-ubuntu`, GPU `-platf 1`) | +| 5 | `simpleitk` | T1c / CT | identity | Mattes MI (32 bins) | `Euler3D` rigid, pyramid `[4,2,1]` | ✅ ran | +| 6 | `synthmorph` | T1c / CT | none (global) | **SynthMorph** anatomy-aware contrast-invariant (deep, pretrained) | 6-DOF rigid matrix | ✅ ran (main env, Docker + GPU) | -Benchmark budget (best-results settings): `elastix`, `impact` and `simpleitk` -run with **MaximumNumberOfIterations = 2000** (per resolution / per pyramid -level), with the sampler and metric left at their defaults -(**NumberOfSpatialSamples 5000**, **NumberOfHistogramBins 32**). `antspyx`, -`niftyreg` and `synthmorph` keep their native defaults. +Benchmark budget (best-results settings): `elastix` and `simpleitk` run with +**MaximumNumberOfIterations = 2000** (per resolution / per pyramid level), +with the sampler and metric left at their defaults (**NumberOfSpatialSamples +5000**, **NumberOfHistogramBins 32**). `impact` runs the published ImpactReg +preset **`Elastix_IMPACT_Static`** (the `impact-reg-konfai` package) — an +elastix B-spline (4 scales, FinalGridSpacing 14 mm, per-level +**MaximumNumberOfIterations 400/300/200/200**, **NumberOfSpatialSamples +2000**) driven by the IMPACT metric: a per-level **MIND (L1) + +TotalSegmentator M730 (deepest layer, Dice)** combination, layer weights +0.2/0.8 → 0.7/0.3, Static mode, + Mattes MI weight 0.4 + bending energy — +with CenterOfGravity auto-init and `RequiredRatioOfValidSamples 0.1` applied +as raw-elastix overrides for the M6 pairs. `antspyx`, `niftyreg` and +`synthmorph` keep their native defaults. `simpleitk` has two budgets: `--simpleitk-mode benchmark` (single pyramid, identity start — the default, comparable to the other tools) and @@ -132,69 +143,92 @@ longer reported in the tables/plots): Alignment is scored by the achieved **Mattes MI** (nats, warped-CT ~ MR over the MR head; higher = better) with the **unclipped** CT, so all rows are comparable on the same score scale. Full volumes, no mask / crop / pad, -fixed = CT. `niftyreg` and `synthmorph` clip the CT to [0, 80] HU *before -registration* (per the CT-clipping A/B below; `clip_ct=None` disables it); -`antspyx`, `elastix` and `simpleitk` register the CT as-is. `niftyreg` is the -v2.1.1 Docker build on the GPU (`-platf 1`); `impact` cannot run on these -pairs under the no-mask rule (see the IMPACT section). +fixed = T1c / moving = CT (shared rule; `impact`'s MR_CT convention +coincides with it). `niftyreg` and `synthmorph` clip the CT (the moving +image) to [0, 80] HU *before registration* (per the CT-clipping A/B below; +`clip_ct=None` disables it); `antspyx`, `elastix` and `simpleitk` register +the CT as-is. `niftyreg` is the v2.1.1 Docker build on the GPU (`-platf 1`); +`impact` runs the `impact-reg-konfai` package on the GPU (see the IMPACT +section). | pipeline | toolkit | registered | s/reg (med) | tv-in-CT | MI (med/mean, nats) | |---|---|---|---|---|---| -| **antspyx** | ANTsPy 0.6.3 | 9 | 36.6 | 0.93 | **0.653 / 0.644** | -| **elastix** | ITK-Elastix 0.25.4 | 9 | 6.0 | 0.93 | **0.517 / 0.520** | -| **impact** | Docker Elastix 5.2 + IMPACT (MRSegmentator) | 8 | 2534 | 0.93 | **0.591 / 0.575** | -| **niftyreg** | NiftyReg 2.1.1 (reg_aladin, Docker + CUDA) | 9 | 20.8 | 0.93 | **0.292 / 0.338** | -| **simpleitk** | SimpleITK 2.5.6 | 9 | 10.6 | 0.93 | **0.651 / 0.641** | -| **synthmorph** | Docker SynthMorph (TensorFlow) | 9 | 24.5 | 0.93 | **0.277 / 0.328** | +| **antspyx** | ANTsPy 0.63 | 9 | 15.5 | 0.93 | **0.652 / 0.650** | +| **elastix** | ITK-Elastix 0.25.4 | 9 | 9.8 | 0.93 | **0.646 / 0.638** | +| **impact** | impact-reg-konfai (Elastix 5.2 + IMPACT, `Elastix_IMPACT_Static` preset, MIND+M730, deformable) | 9 | 216.4 | 0.93 | **0.578 / 0.558** | +| **niftyreg** | NiftyReg 2.1.1 (reg_aladin, Docker + CUDA) | 9 | 16.2 | 0.93 | **0.292 / 0.339** | +| **simpleitk** | SimpleITK 2.5.6 | 9 | 9.6 | 0.93 | **0.647 / 0.642** | +| **synthmorph** | Docker SynthMorph (TensorFlow) | 9 | 27.9 | 0.93 | **0.277 / 0.328** | -Per-case MI (nats, unclipped score; **bold** = best for the case): +Per-case MI (nats, unclipped score; **bold** = best for the case). All rows +are re-runs under the **fixed = T1c / moving = CT** convention (the shared +rule); the CT-fixed-era numbers are in the git history. The `impact` column +is the `impact-reg-konfai` deformable `Elastix_IMPACT_Static` preset (the +retired rigid Docker config is archived in +`results/bench_registration/impact_docker_legacy.json`). | case | antspyx | elastix | impact | niftyreg | simpleitk | synthmorph | |---|---|---|---|---|---|---| | m6_22GD5VL2_20230724 | incomplete-CT | incomplete-CT | incomplete-CT | incomplete-CT | incomplete-CT | incomplete-CT | -| m6_22IRRKU6_20240115 | **0.649** | 0.472 | 0.635 | 0.292 | 0.642 | 0.277 | -| m6_23XIMO4M_20230222 | **0.655** | 0.416 | 0.478 | 0.632 | 0.652 | 0.616 | -| m6_24V62QQM_20230223 | **0.704** | 0.701 | 0.697 | 0.260 | 0.663 | 0.256 | -| m6_24V62QQM_20231017 | 0.601 | 0.590 | 0.556 | 0.493 | **0.602** | 0.484 | -| m6_24V62QQM_20240820 | **0.653** | 0.644 | 0.531 | 0.118 | 0.652 | 0.111 | -| m6_26TYQI3C_20221027 | 0.659 | 0.653 | gate-abort | 0.609 | **0.695** | 0.590 | -| m6_2BBN5XA4_20250411 | **0.577** | 0.287 | 0.450 | 0.069 | **0.577** | 0.061 | -| m6_2C4TKI6H_20240506 | **0.645** | 0.401 | 0.629 | 0.233 | 0.634 | 0.231 | -| m6_2CDNN7BH_20240102 | **0.655** | 0.516 | 0.625 | 0.340 | 0.651 | 0.323 | +| m6_22IRRKU6_20240115 | **0.648** | 0.644 | 0.542 | 0.292 | 0.647 | 0.277 | +| m6_23XIMO4M_20230222 | **0.654** | 0.652 | 0.543 | 0.632 | **0.654** | 0.616 | +| m6_24V62QQM_20230223 | 0.703 | 0.698 | 0.667 | 0.260 | **0.704** | 0.256 | +| m6_24V62QQM_20231017 | **0.601** | 0.590 | 0.587 | 0.493 | 0.597 | 0.484 | +| m6_24V62QQM_20240820 | **0.652** | 0.646 | 0.612 | 0.118 | 0.638 | 0.111 | +| m6_26TYQI3C_20221027 | **0.718** | 0.650 | 0.717 | 0.609 | 0.670 | 0.590 | +| m6_2BBN5XA4_20250411 | 0.577 | 0.577 | 0.258 | 0.069 | **0.579** | 0.061 | +| m6_2C4TKI6H_20240506 | **0.645** | 0.638 | 0.578 | 0.233 | 0.636 | 0.231 | +| m6_2CDNN7BH_20240102 | **0.655** | 0.650 | 0.523 | 0.340 | 0.652 | 0.323 | ### Interpretation -- **SimpleITK and ANTsPy tie for best** (median MI 0.651 / 0.652): on the - well-positioned cases they cluster at 0.60–0.70 and both survive the harder - repositioned cases (0.58–0.65 on `m6_2BBN5XA4` / `m6_22IRRKU6`). SimpleITK's - per-case MI is **unchanged** from the previous 150-iteration budget — RSGD - converges well before the 2000-iteration cap, so the cap is now headroom - rather than a constraint. -- **The grossly repositioned cases define the ranking.** On - `m6_2BBN5XA4` and `m6_24V62QQM_20240820` (~30° inter-session tilt) the tools - split: ANTs/SimpleITK hold 0.58–0.65, elastix drops to 0.29–0.64, niftyreg - 0.07–0.12, synthmorph 0.06–0.10 — single-start optimizers locking onto local - minima. This is exactly why production `02_build_m6_dataset.py` uses a - **multi-start tilt grid**. -- **Elastix is second at the new budget** (2000 iters per resolution, 32 bins, - 5000 samples): median 0.516, below the 0.641 recorded at the old 500/500/250 - + 10-bin setting. The long ASGD run nails the clean cases (0.64–0.70) but its - single CenterOfGravity start still locks onto local minima on the - repositioned ones (`m6_2BBN5XA4` 0.287, `m6_2C4TKI6H` 0.401, `m6_23XIMO4M` - 0.416). -- **IMPACT (semantic metric) now runs on these pairs** — the no-mask abort was - elastix's sample-validity gate (build default `RequiredRatioOfValidSamples` - 0.25) rejecting the ~14% in-buffer fraction of the CT head-neck FOV vs the - tight T1c crop; the parameter map sets it to 0.1 (see the IMPACT section). - With MRSegmentator at the 2000/5000 budget it is top-three in MI on 7 of the - 8 registered cases (median 0.591): 0.635 on the hardest `m6_22IRRKU6` - (masked-era 0.624), 0.697 on `m6_24V62QQM_20230223` (near-best), 0.629 on - `m6_2C4TKI6H` — and it shares the production-acceptance of the grossly - repositioned `m6_2BBN5XA4` (std 10.7 mm) with antspyx/simpleitk, the only - non-intensity tool to do so. One run diverges mid-registration - (`m6_26TYQI3C`: the in-buffer fraction falls below 10% at iteration 256 and - the gate aborts — a stochastic-draw outcome). It remains by far the slowest - (~42 min/case at 2000/5000). +- **ANTsPy leads, with SimpleITK and Elastix within 0.006 of it** (median MI + 0.652 / 0.647 / 0.646): on the well-positioned cases the three cluster at + 0.60–0.70 and all now survive the harder repositioned cases (0.58–0.65 on + `m6_2BBN5XA4` / `m6_22IRRKU6`). +- **Switching to fixed = T1c / moving = CT lifted Elastix sharply** + (0.517 → 0.646 median, and it is now #3 rather than #4). The Mattes-MI + pipelines sample spatial pairs from the fixed buffer and discard the ones + that map outside the moving image; with the small T1c head crop as fixed and + the large CT head+neck FOV as moving, (nearly) all fixed samples are + in-buffer, versus only a minority when the big CT was fixed. Elastix's + repositioned cases jumped (`m6_2BBN5XA4` 0.287→0.577, `m6_2C4TKI6H` + 0.401→0.638, `m6_23XIMO4M` 0.416→0.652). The other pipelines are essentially + unchanged by the swap (antspyx 0.653→0.652, simpleitk 0.651→0.647, niftyreg + 0.292→0.292, synthmorph 0.277→0.277) — they were not sample-coverage + limited. +- **The grossly repositioned cases still separate the leaders from the + tail.** On `m6_2BBN5XA4` / `m6_24V62QQM_20240820` (~30° inter-session tilt) + ANTs/SimpleITK/Elastix hold 0.58–0.65, while niftyreg (0.07–0.12) and + synthmorph (0.06–0.11) — the single-start block-matching and pretrained-MRI + methods — lock onto local minima. This is exactly why production + `02_build_m6_dataset.py` uses a **multi-start tilt grid**. +- **IMPACT now runs the published `Elastix_IMPACT_Static` preset** + (`impact-reg-konfai` package; the old Docker backend is retired): per-level + **MIND + TotalSegmentator M730** combination (L1 + deepest-layer Dice, + weights 0.2/0.8 → 0.7/0.3, Static mode, whole-image features), the preset's + native deformable B-spline (400/300/200/200 iterations, 2000 samples, + + Mattes MI 0.4 + bending energy), registered in the MR_CT direction + (fixed=T1c) with CenterOfGravity auto-init and the 0.1 sample-validity gate + as overrides. It is the best of the IMPACT variants on M6 and now **second + overall at median MI 0.578** (behind SimpleITK/ANTsPy 0.65; ahead of + elastix 0.52), a clear step up from the old rigid Docker config (median + 0.467, archived in `impact_docker_legacy.json`): the deformable refinement + + MR_CT direction lifts the repositioned cases hardest (`m6_2C4TKI6H` + 0.334→0.578, `m6_23XIMO4M` 0.362→0.543, `m6_22IRRKU6` 0.467→0.542, + `m6_2CDNN7BH` 0.332→0.523; `m6_26TYQI3C` 0.717 — now the best tool on that + case) at a small cost on two well-aligned ones (`m6_24V62QQM_20230223` + 0.700→0.667, `m6_24V62QQM_20240820` 0.646→0.612, still accepted). + Concentricity also improved (median mid 12.0 / std 21.9 mm vs 21.1 / 28.0; + 3/9 accepted). It is the slowest of the six (~216 s/case, the deformable + optimization + KonfAI wrapper; the old rigid config was ~78 s). The + `m6_2BBN5XA4` ~30° tilt case is still its weakest (0.258) — the single + CenterOfGravity start remains the limiting factor. Earlier Docker-era + variants: MIND+MRSegmentator median 0.417 (~72 s), single-model MIND + median 0.406 (~37 s), MRSegmentator-only Jacobian configs 0.631 + (500/2000) and 0.591 (2000/5000) at ~5×/~35× the cost (archived: + `results/bench_registration/impact_marseg_*.json`, `impact_2000_5000.json`, + `impact_500_2000_explicit_smoothing.json`, `impact_mind_500_ref.json`). - **NiftyReg (reg_aladin) is mid-pack**: single-start symmetric block matching is strong on moderate cases (`m6_23XIMO4M` 0.632, best in class; `m6_26TYQI3C` 0.609) but falls into the same tilt local minima @@ -212,23 +246,28 @@ Per-case MI (nats, unclipped score; **bold** = best for the case): It needs **no centroid init** (global deep method). - **`tv-in-CT` is ~0.93 for every pipeline**: the warped GTV lands inside the CT head everywhere, so coverage is not the discriminator — MI is. -- **Speed** (median wall, 4 cases in parallel, 24 threads): elastix 6.0 s, - SimpleITK 10.6 s, **NiftyReg 20.8 s** (v2.1.1, GPU block matching), - SynthMorph 24.5 s (GPU forward pass), ANTs 36.6 s. (IMPACT is by far the - slowest — 2534 s/case median at the 2000/5000 benchmark budget, ~276 s/case - at the authors' 500/2000; GPU feature extraction with backprop per - iteration. No-mask runs are faster per iteration than the masked-era - ~833 s/case @500/2000 because only the ~14% in-buffer samples are scored.) +- **Speed** (median wall, 4 cases in parallel, 24 threads, T1c-fixed re-run): + SimpleITK 9.6 s, elastix 9.8 s, **ANTs 15.5 s** (down from 36.6 s — the + smaller fixed image also shaves ANTs' per-level work), **NiftyReg 16.2 s** + (v2.1.1, GPU block matching), SynthMorph 27.9 s (GPU forward pass), + **IMPACT 216.4 s** (`Elastix_IMPACT_Static` preset — Static feature mode, no + backprop, but the deformable B-spline optimization at 1100 total iterations + + the KonfAI wrapper; the retired rigid Docker config was 78.0 s). The + earlier MRSegmentator Jacobian configs were far slower still: ~354 s/case at + 500/2000, 2534 s/case at 2000/5000 (GPU feature extraction with backprop + per iteration). - **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; the in-process Elastix ASGD is deterministic given its fixed seed; SynthMorph is deterministic; NiftyReg is deterministic on the CPU (conda 1.5.69) but shows small run-to-run variation on the GPU (v2.1.1, - ~±0.01–0.05 MI on the weak cases). **IMPACT is the most stochastic**: - its Jacobian mode draws random patches every iteration, so the final pose - changes run to run. Treat its per-case numbers as one sample from a - distribution. + ~±0.01–0.05 MI on the weak cases). **IMPACT in the Jacobian preset + (`Elastix_IMPACT_Jacobian`) is the most stochastic**: it draws random + patches through the network every iteration, so the final pose changes run + to run — treat such numbers as one sample from a distribution. The default + Static preset re-samples coordinates every iteration as well, but the + feature maps are deterministic given the pose. ## CT clipping A/B (`--clip-ct 0 80`) @@ -293,15 +332,16 @@ python scripts/bench_registration.py --pipelines simpleitk elastix antspyx nifty --n-cases 10 --workers 4 --threads 24 --simpleitk-mode benchmark # force the conda 1.5.69 CPU build: --niftyreg-backend conda -# 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 -# NOTE: impact now runs on the M6 pairs (the parameter map relaxes the -# sample-validity gate to RequiredRatioOfValidSamples 0.1) but is by far the -# slowest: ~5 min/case at 500/2000, ~42 min/case (2534 s median) at 2000/5000 -# (--impact-iters / --impact-samples) — see the IMPACT section +# impact (pip; GPU; ~120-220 s/case; one-time elastix shim on hosts too old +# for the stock binary — see the IMPACT section). --impact-preset selects +# another published preset (e.g. Elastix_IMPACT_Jacobian, Generic_Rigid); +# --impact-iters / --impact-samples are no-ops for the default preset's +# per-resolution matrix python scripts/bench_registration.py --pipelines impact --n-cases 10 \ + --workers 3 --threads 24 + +# synthmorph (Docker, needs the pulled image; see the SynthMorph section) +python scripts/bench_registration.py --pipelines synthmorph --n-cases 10 \ --workers 4 --threads 24 # A/B: does clipping the CT to a brain window help? (--clip-ct clips the CT @@ -323,9 +363,10 @@ Other flags: `--sample random --seed 0` (reproducible random 10), `--threads T` (also the OpenMP thread count for `reg_aladin`), `--simpleitk-mode {benchmark,production}`, `--synthmorph-model {rigid,affine}`, `--synthmorph-gpu `, -`--synthmorph-extent `, `--impact-model `, -`--impact-gpu `, `--impact-iters N` (default 2000), -`--impact-samples N` (default 5000), +`--synthmorph-extent `, `--impact-preset ` (default +`Elastix_IMPACT_Static`), `--impact-gpu `, `--impact-iters N` +(preset max_iterations override; no-op for matrix presets), +`--impact-samples N` (preset spatial_samples override), `--out-dir `, `--qa-images` / `--qa-out ` (visual QA, below). ## Visual QA images (checkerboard + color overlay) @@ -391,10 +432,11 @@ No extra model assets: the rigid weights ship inside the image visible (auto-spread over the host GPUs by pid; `--synthmorph-gpu -1` forces CPU). SynthMorph is a global deep method, so it recovers the gross inter-session offset on its own and needs **no** initial transform. -2. **Read-back** — the `-T` LTA (a `LINEAR_RAS_TO_RAS` 4×4 in shared world - space; with fixed=CT/moving=T1c the `-T` matrix is the T1c→CT contract - transform) is parsed into a SimpleITK `AffineTransform` and run through the - usual `fix_orientation` gate. +2. **Read-back** — the model estimates first→second, so the pair is passed as + `(fixed=T1c, moving=CT)` and the `-T` LTA (a `LINEAR_RAS_TO_RAS` 4×4 in + shared world space) is the T1c→CT contract transform directly; it is + parsed into a SimpleITK `AffineTransform` and run through the usual + `fix_orientation` gate. 3. **GPU sharing** — the container is run as the host uid (clean scratch removal, TF still reaches the GPU) with `TF_FORCE_GPU_ALLOW_GROWTH` set so concurrent benchmark workers can share a GPU. @@ -434,112 +476,114 @@ conda install -c conda-forge niftyreg ### How the adapter drives the tool -1. **Rigid 6-DOF, full volumes** — `reg_aladin -ref CT -flo T1c -rigOnly` - (symmetric variant on by default, which matters: the non-symmetric mode - drifts grossly on these repositioned pairs). `-ln 3` pyramid, - `-maxit 15` TLS iterations per level (the tool's default of 5 is too - coarse), `-interp 1` (GPU) / `-interp 3` (CPU) internal warping, default - initialisation (image centres). GPU: `--gpus all`, `-platf 1`, - `-gpuid `; CPU/conda: `-platf 0` (v2.1.1) or no `-platf` - (1.5.69). +1. **Rigid 6-DOF, full volumes** — `reg_aladin -ref T1c -flo CT -rigOnly` + (the T1c is the reference/fixed image, the CT the floating/moving one; + the CT is clipped to [0,80] HU first). Symmetric variant on by default, + which matters: the non-symmetric mode drifts grossly on these + repositioned pairs. `-ln 3` pyramid, `-maxit 15` TLS iterations per level + (the tool's default of 5 is too coarse), `-interp 1` (GPU) / `-interp 3` + (CPU) internal warping, default initialisation (image centres). GPU: + `--gpus all`, `-platf 1`, `-gpuid `; CPU/conda: `-platf 0` + (v2.1.1) or no `-platf` (1.5.69). 2. **Read-back** — the output affine (`-aff`) is a 4×4 row-major matrix, - documented as `Affine * Reference = Floating` (CT→T1c); the adapter parses - it into a SimpleITK `AffineTransform` and runs it through the usual - `fix_orientation` gate (skull-ring concentricity). On the M6 cases the gate - keeps the file's affine as-is (`auto_inverted: False`), i.e. it is already - the contract transform (T1c→CT). + documented as `Affine * Reference = Floating`; the adapter parses it into + a SimpleITK `AffineTransform` and runs it through the usual + `fix_orientation` gate (skull-ring concentricity), which resolves the + direction empirically rather than trusting the docs. Under + `-ref T1c -flo CT` the file's affine comes out the other way, so the gate + inverts it on the M6 cases (8/9 `auto_inverted: True`) to the contract + transform (T1c→CT). -## IMPACT pipeline (Docker) — one-time setup +## IMPACT pipeline (impact-reg-konfai) — one-time setup -The `impact` pipeline runs the [`vboussot/elastix_impact`](https://hub.docker.com/r/vboussot/elastix_impact) -image: **Elastix 5.2.0** with the [IMPACT](https://arxiv.org/abs/2503.24121) -*semantic* similarity metric compiled in as a native Elastix metric, plus a -TorchScript segmentation feature extractor (**MRSegmentator** `MRSeg/MRSeg.pt` -by default; TotalSegmentator `TS/M730.pt` and other models from -`VBoussot/impact-torchscript-models` are selectable). It is GPU-based and -Docker-based (alongside `synthmorph`). It follows the shared rules (full -volumes, no mask, fixed = CT, moving = T1c). The metric only evaluates samples -that map inside the moving buffer, and the CT head-neck FOV (~300×300×251 mm) -is much larger than the T1c head crop (~220×220×150 mm), so ~14% of the samples -are in-buffer at the auto-init. The image build's sample-validity gate -defaults to `RequiredRatioOfValidSamples 0.25` and used to abort the M6 pairs -on that ("too many samples map outside moving image buffer"); the parameter -map now sets it to `0.1` (see the IMPACT specifics below). +The `impact` pipeline runs the +[impact-reg-konfai](https://pypi.org/project/impact-reg-konfai/) package: +[IMPACT](https://arxiv.org/abs/2503.24121) registration presets from the +[`VBoussot/ImpactReg`](https://huggingface.co/VBoussot/ImpactReg) Hugging Face +repo, driven through the [KonfAI](https://github.com/fideus-labs/KonfAI) +runtime — no Docker. The default preset is **`Elastix_IMPACT_Static`**: +Elastix 5.2 with the IMPACT *semantic* similarity metric — a per-level +combination of the **MIND** handcrafted descriptor (`MIND/R1D2_3D.pt`, L1, +subset 32) and **TotalSegmentator** `TS/M730.pt` (deepest layer, Dice, subset +64), layer weights 0.2/0.8 → 0.7/0.3 from coarse to fine, **Static** mode +(whole-image features, no backprop), a 4-scale B-spline (FinalGridSpacing +14 mm), 400/300/200/200 iterations per level, 2000 spatial samples, + Mattes +MI (w 0.4) + bending energy — the published successor of the `MR_CT_TS` +reference parameter map this repo used to hand-build (and replace the +`vboussot/elastix_impact` Docker image). Other presets are selectable with +`--impact-preset` (`Elastix_IMPACT_Jacobian`, `Generic_Rigid`, +`Generic_Rigid_BSpline`, `FireANTs_Anatomix/IMPACT/SyN`, +`ConvexAdam_Coarse/Composite`; the live database is listed by +`impact-reg-konfai register --help`). It is GPU-based (like `synthmorph`), +but pip-installed. Prerequisites (check with `RP.availability()`; a missing piece raises `PipelineUnavailable` with the exact command): ```bash -# 1. Docker (with NVIDIA container toolkit so --gpus works) -docker --version - -# 2. Pull the image (~23.5 GB, one-time) -docker pull vboussot/elastix_impact - -# 3. Fetch the TorchScript model (~123 MB, one-time) into the repo's model -# cache (data/impact_models/). MRSeg.pt (MRSegmentator) is the default: -# it is trained on both MRI and CT, the right prior for this -# cross-modality CT->T1c task (all models share the same 2^l+3 FOV -# architecture, so the parameter map is unchanged between them). -mkdir -p /mnt/t24/Public/xfr/longitudinal/impact_models/MRSeg -curl -L -o /mnt/t24/Public/xfr/longitudinal/impact_models/MRSeg/MRSeg.pt \ - https://huggingface.co/VBoussot/impact-torchscript-models/resolve/main/MRSeg/MRSeg.pt +# in the longitudinal-313 env (Python >= 3.11): +pip install impact-reg-konfai # CLI + KonfAI runtime (pulls h5py etc.) +bash scripts/setup_impact_konfai.sh # one-time, this host only: the elastix shim ``` -Other models (TotalSegmentator `TS/M730.pt`, `TS/M258/M291/M850`, MIND, SAM2.1, -Anatomix, …) are selectable with `--impact-model `; -see the -[impact-torchscript-models](https://huggingface.co/VBoussot/impact-torchscript-models) -repo for the full list and FOV/layer tables. The model choice does not affect -the sample-validity gate (below): the "too many samples outside moving buffer" -check is geometric (CT head-neck FOV vs the tight T1c crop) and fires at the -first sample evaluation, before feature extraction. +The CLI then auto-downloads the one-time assets on first `register`: the +preset app (HF `VBoussot/ImpactReg`), the TorchScript feature models (HF +`VBoussot/impact-torchscript-models`: MIND ~15 KB + TS/M730 ~200 MB) and the +**elastix-impact 1.0.0** binary (GitHub `vboussot/ImpactElastix`, ~40 MB) +into `~/.cache` (HF cache + `~/.cache/konfai`). -### How the adapter drives the image +### Why the elastix shim (this host only) -The image's `CMD.sh` is a generic launcher, so the adapter calls the in-image -binary directly (`/lib/elastix-install/bin/elastix`): +The stock binary needs **glibc ≥ 2.35** and links **LibTorch 2.8.0**; this +host (Ubuntu 20.04, glibc 2.31, env torch 2.14) satisfies neither, so +`scripts/setup_impact_konfai.sh` builds a self-contained shim at +`~/.cache/konfai/elastix-impact-shim`: the glibc 2.35 loader + libs extracted +from the `niftyreg:2.1.1-ubuntu` image (docker is needed once for that copy +only), the self-contained LibTorch 2.8.0 cu128 from download.pytorch.org +(~3.8 GB zip), and a wrapper `bin/elastix` that runs the stock binary through +the shim loader with the LibTorch 2.8.0 library path (stripping the inherited +`LD_LIBRARY_PATH`, so the env's torch 2.14 cannot leak in). The pipeline +selects the shim via `KONFAI_ELASTIX_DIR` (`_impact_elastix_root`); a host +new enough to run the stock binary directly needs no shim at all. -1. **Full volumes, fixed=CT, moving=T1c** — both volumes are used as-is (no - head box, no mask, no crop / pad, per the shared rules above). The metric - only evaluates spatial samples that map inside the moving buffer, and the - CT FOV (~300×300×251 mm, head+neck) is much larger than the tight T1c head - crop (~220×220×150 mm), so at the CenterOfGravity auto-init only ~14% of - the samples are in-buffer. The image build's default validity gate - (`RequiredRatioOfValidSamples 0.25`) aborted on that ("Too many samples map - outside moving image buffer: 704 / 5000"); the parameter map therefore sets - `(RequiredRatioOfValidSamples 0.1)` so the M6 pairs run on their in-buffer - samples while a genuinely diverging transform (<10% in-buffer) still - aborts. -2. **Auto init (no `-t0`)** — the parameter map sets - `AutomaticTransformInitialization "true"` with method `CenterOfGravity`, so - elastix aligns the intensity centroids; `AutomaticScalesEstimation` sets the - optimizer scales. -3. **Parameter map** — rigid `EulerTransform` (6-DOF), `Impact` metric in - **Jacobian** mode (differentiable, per the recommended config), 4 semantic - scales `ImpactVoxelSize` 6/3/1.5/1 mm (scaled to the T1c base spacing) with - **no image resampling** (all-1 pyramid schedules, as the IMPACT docs require), - ASGD, `Random` sampler, `(RequiredRatioOfValidSamples 0.1)`. A copy of the - exact parameter map is dumped to - `data/m6reg/impact/paramdump/ParameterMap__.txt` before every - docker run (also recorded as `param_dump` in the result JSON). GPU - auto-spread across the host GPUs by pid (`--impact-gpu -1` forces CPU). -4. **Intensity standardization happens inside the metric** (no pipeline-side - normalization — the model card's "standardize intensities to zero mean and - unit variance" is already applied): at init the metric computes - whole-volume `[min, max, mean, sigma]` for each image - (`itk::StatisticsImageFilter`) and passes it to the model, whose - `normalize` branch rescales every extracted patch with the image's own - `(mean, sigma)`. It is computed over the whole buffer (CT air included) and - identically for fixed and moving, so both modalities' features sit on the - same normalized scale. Pre-standardizing the MHA files in the pipeline - would be an exact no-op (z-scoring is affine-invariant: the metric - recomputes the stats of the image it receives) and would only corrupt the -HU scale used by MI scoring and the QA brain-window images. - 5. **Read-back** — the result `TransformParameters.0.itk.txt` is an - `Euler3DTransform` in the world physical frame (the - `AutomaticTransformInitialization` is folded in by elastix); the adapter - runs it through the usual `fix_orientation` gate. +### How the adapter drives the CLI + +1. **Full volumes, fixed=T1c (MR), moving=CT** — the ImpactReg **MR_CT** + convention, which coincides with the shared rule (the T1c is the fixed + image): the preset's output is a displacement field on the **fixed** grid + mapping fixed→moving, so the field is the contract transform T1c→CT + **directly** (displacement fields cannot be inverted with `GetInverse`, + which a fixed=CT direction would require). It also keeps + the metric's sample-validity gate healthy: samples drawn from the tight + T1c head crop (~220×220×150 mm) map mostly inside the large CT head-neck + FOV (~300×300×251 mm), unlike fixed=CT where only ~14% were in-buffer at + the auto-init. +2. **M6 robustness overrides** — the preset's map is tuned for well-aligned + pairs, so the adapter forwards raw elastix lines through the package's + `parameter_overrides` knob: `AutomaticTransformInitialization=true` + + `AutomaticTransformInitializationMethod=CenterOfGravity` (the preset's + GeometricalCenter start leaves the intensity centroids ~50-100 mm apart on + the head+neck FOV mismatch) and `RequiredRatioOfValidSamples 0.1` (build + default 0.25; keeps the divergence gate while letting the FOV-mismatched + runs proceed). +3. **`--fields-only`** — the adapter derives the warp itself (the package's + derived "Moved" image would be an extra full-size resample of the same + content). +4. **GPU auto-spread** across the host GPUs by pid (`--impact-gpu -1` forces + CPU; the engine then turns mixed precision off). +5. **Read-back** — the transform comes from `/P000/Transform.h5` (an + ITK-form HDF5 `DisplacementFieldTransform` on the T1c grid; the SimpleITK + build here has no ITK HDF5 transform reader, so it is parsed with h5py) + and direction-checked by `_orient_impact`: the skull-ring gap of the field + vs a sign-flipped (first-order reverse) proxy — the `fix_orientation` + analogue for a field with no `GetInverse`. +6. **Intensity normalization happens inside the metric / model** (no + pipeline-side normalization): the MIND descriptor normalizes each 5-voxel + neighbourhood to [0, 1] locally (modality-agnostic by construction); the + segmentation models (TS/MRSeg) are standardized to zero mean / unit + variance per the model card. Pre-normalizing the volumes would be an + exact no-op for the z-scored models and would only corrupt the HU scale + used by MI scoring and the QA brain-window images. ## Notes / caveats @@ -556,15 +600,20 @@ HU scale used by MI scoring and the QA brain-window images. subsampling), so its per-case numbers are stable run to run — unlike ANTs / IMPACT. Its low MI (median 0.238; 0.06–0.10 on the repositioned cases) is a genuine cross-modality / tilt weakness, not variance. -- `fix_orientation` (elastix/ants/niftyreg/synthmorph/impact) chooses between a +- `fix_orientation` (elastix/ants/niftyreg/synthmorph) 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. With the fixed=CT/moving=T1c swap, `elastix`'s combination - transform is now the inverse of the contract, so it reports - `auto_inverted: True` (the gate corrects it); antspyx/synthmorph/impact - return the contract direction directly (`auto_inverted: False`). `niftyreg`'s - `-aff` file is documented as CT→T1c but the gate keeps it as-is on these - cases (`auto_inverted: False`) — the empirical gate, not the docs, decides. + convention, so the fixed/moving swap is absorbed by the gate rather than + re-derived per toolkit. Under the shared fixed=T1c / moving=CT rule the raw + transforms come out in the contract direction for elastix (combination + transform = fixed→moving), antspyx (ANTs forward transform) and synthmorph + (the `-T` LTA, model estimates first→second) — completed runs report + `auto_inverted: False` for those; `niftyreg`'s `-aff` file comes out the + other way and the gate inverts it (8/9 `auto_inverted: True`). `impact` + returns a displacement field (no `GetInverse` in ITK), so it uses + `_orient_impact` — the same skull-ring decision against a sign-flipped + (first-order reverse) proxy field; its MR_CT field is in the contract + direction (`auto_inverted: False`). - **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 @@ -577,8 +626,8 @@ HU scale used by MI scoring and the QA brain-window images. path). - **GPU sharing.** The container runs as the host uid with `TF_FORCE_GPU_ALLOW_GROWTH` so concurrent benchmark workers share a GPU. - - **~25 s/case** (GPU: one pretrained forward pass), between SimpleITK - (~10 s) / ANTs (~37 s) and IMPACT (~834 s). +- **~25 s/case** (GPU: one pretrained forward pass), between SimpleITK + (~10 s) / ANTs (~37 s) and IMPACT (~216 s, deformable preset). - **NiftyReg specifics** (conda-forge `niftyreg` 1.5.69, `reg_aladin`, symmetric block matching + TLS, CPU/OpenMP): - **Different algorithmic family.** No histogram MI objective: 4-voxel @@ -618,61 +667,60 @@ HU scale used by MI scoring and the QA brain-window images. - `WriteResultImage` must be the string `"false"` (Python `False`); the default (absent) makes the filter raise "Could not read result image" after a *successful* registration. - - With fixed=CT/moving=T1c the combination transform is CT→T1c (the inverse - of the contract), so `fix_orientation` reports `auto_inverted: True` and - inverts it. The result is recovered by point-mapping - (`GetCombinationTransform().TransformPoint`) into an affine — independent - of elastix's Euler-angle parameterization — and wrapped in a SimpleITK - `AffineTransform(3)`. +- With fixed=T1c/moving=CT the combination transform is T1c→CT — the + contract direction — so `fix_orientation` keeps it as-is + (`auto_inverted: False`). The result is recovered by point-mapping + (`GetCombinationTransform().TransformPoint`) into an affine — independent + of elastix's Euler-angle parameterization — and wrapped in a SimpleITK + `AffineTransform(3)`. - Benchmark budget: `MaximumNumberOfIterations` 2000 per resolution (ASGD 2000/2000/2000 over 3 resolutions) with the sampler/metric at their defaults (`NumberOfSpatialSamples` 5000, `NumberOfHistogramBins` 32). -- **IMPACT specifics** (Docker `vboussot/elastix_impact`, Elastix 5.2.0, - GPU, Jacobian mode): -- **Slow, by design.** Every ASGD iteration runs the segmentation - extractor (MRSegmentator by default) on the sampled patches with - backpropagation. The benchmark -budget (MaximumNumberOfIterations 2000, NumberOfSpatialSamples 5000 = - the elastix default) is ~10× the authors' 500-iter / 2000-sample - Jacobian budget: measured 2534 s/case (median) on the no-mask M6 - pairs, ~276 s/case at 500/2000. `--impact-iters` / `--impact-samples` - trade accuracy for speed. -- **Runs on M6 with a relaxed sample-validity gate.** Like every other - pipeline it runs fixed=CT/moving=T1c on the full volumes (no mask, no - crop / pad). The CT FOV (~300×300×251 mm) is much larger than the T1c - crop (~220×220×150 mm), so at the CenterOfGravity auto-init only ~14% of - the samples map inside the T1c moving buffer. The metric calls - `CheckNumberOfSamples()` after every evaluation and the image build - defaults to `RequiredRatioOfValidSamples 0.25`, so it used to abort - ("Too many samples map outside moving image buffer: 704 / 5000"). The - parameter map now sets `(RequiredRatioOfValidSamples 0.1)`: the gate - still aborts a genuinely diverging transform (<10% in-buffer) while the - FOV-mismatched M6 pairs run on their ~14% in-buffer samples — exactly - the native "only in-buffer samples are evaluated" behaviour. In the - 2000/5000 re-run the gate still fired once, mid-registration: on - `m6_26TYQI3C` the transform drifted until only 468/5000 samples (9.4%) - stayed in-buffer at iteration 256 and the run aborted. The transform - contract (T1c→CT) holds; `fix_orientation` decides the direction as - usual (all 8 completed runs report `auto_inverted: True`). - - **Auto init.** `AutomaticTransformInitialization` (CenterOfGravity) aligns - the intensity centroids; no manual `-t0` translation is used. -- **History: single-init convergence.** The semantic landscape has more +- **IMPACT specifics** (`impact-reg-konfai` 1.8.6, KonfAI runtime, + elastix-impact 5.2, GPU; preset `Elastix_IMPACT_Static`): + - **Runs on M6 with the T1c-fixed direction + relaxed sample-validity + gate.** Full volumes, no mask, no crop / pad, fixed=T1c/moving=CT (the + shared rule; the ImpactReg MR_CT convention coincides with it): the + preset's output field is on the fixed grid, so that direction is the + contract transform (T1c→CT) directly, and the metric's samples (drawn + from the tight T1c crop) map mostly inside the large CT FOV. The package's generated map still defaults to + `RequiredRatioOfValidSamples 0.25`, so the adapter overrides it to 0.1 + (a genuinely diverging transform, <10% in-buffer, still aborts). In the + old Docker era (fixed=CT) the gate aborted the no-mask M6 pairs until + the same 0.1 relaxation. The transform contract holds; `_orient_impact` + decides the direction as usual (completed runs report + `auto_inverted: False` — the MR_CT field is already the contract). + - **Auto init.** `AutomaticTransformInitialization` (CenterOfGravity) is + applied as an override (the preset's GeometricalCenter start is too + crude for the head+neck FOV mismatch); no manual translation is used. + - **Deformable, not rigid.** The preset's native transform is a + RecursiveBSpline (FinalGridSpacing 14 mm, 4 scales) regularized by the + bending-energy metric; the old pipeline forced a rigid `EulerTransform` + in place of it. The benchmark contract (warp the CT concentric with the + MR head) is transform-agnostic, so this is a native-config comparison + point like every other row. + - **History: single-init convergence.** The semantic landscape has more local minima than Mattes-MI; earlier (masked, fixed=CT) runs no longer nailed the hardest repositioned case (`m6_22IRRKU6` mid 22.7 / std 27.0; the pre-mask head-box flow got it to mid 8.3 / std 13.2). The no-mask - MRSegmentator runs show the same pattern: strong MI (0.635 on - `m6_22IRRKU6` at 2000/5000) but the pose still misses the concentricity - acceptance rule there (mid 26.1 / std 28.6); it *is* accepted on - `m6_2BBN5XA4` (mid 15.1 / std 10.7). - - **Stochastic.** Jacobian mode resamples patches each iteration, so results - vary run to run (see Stochasticity). - - **Elastix 5.x differences vs the in-process `elastix` pipeline:** the image - bundles 5.2.0, so the in-file multi-`Transform` chain and `UseScalingFactors` - are gone. The result is read back as an `Euler3DTransform` (`.itk.txt`). - The image's own `CMD.sh` is a generic launcher, so the adapter invokes - `/lib/elastix-install/bin/elastix` directly. - - **Models are not in the image.** The TorchScript model(s) must be present - under the mounted `/Data/Models` (the repo caches them in - `data/impact_models/`, bind-mounted read-only). - - **CPU fallback:** `--impact-gpu -1` runs without `--gpus` and with - `ImpactGPU -1` / mixed precision off (much slower). \ No newline at end of file + MRSegmentator runs showed the same pattern (0.635 MI on `m6_22IRRKU6` + at 2000/5000 but mid 26.1 / std 28.6 there; accepted on `m6_2BBN5XA4`, + mid 15.1 / std 10.7). The single CenterOfGravity start on the ~30° + tilt pairs remains IMPACT's limiting factor — a multi-start (or + `--tta` flips through the preset app) would be the next step. + - **Stochastic.** The Jacobian preset (`Elastix_IMPACT_Jacobian`) + resamples patches through the network each iteration, so results vary + run to run; the default Static preset re-samples coordinates each + iteration but the feature maps are deterministic given the pose (see + Stochasticity). + - **The elastix shim (this host).** The stock binary needs glibc ≥ 2.35 + + LibTorch 2.8.0; the host (glibc 2.31, torch 2.14) runs it through + `~/.cache/konfai/elastix-impact-shim` (built by + `scripts/setup_impact_konfai.sh`, selected via `KONFAI_ELASTIX_DIR`). + Same glibc story as the niftyreg v2.x binaries. + - **Models / presets are not in the repo.** They auto-download from HF + (`VBoussot/ImpactReg`, `VBoussot/impact-torchscript-models`) into + `~/.cache/huggingface` on first use; the old `data/impact_models/` + cache is obsolete. + - **CPU fallback:** `--impact-gpu -1` (no `--gpu` flag; the engine runs + the feature models on CPU with mixed precision off — much slower). \ No newline at end of file diff --git a/scripts/bench_registration.py b/scripts/bench_registration.py index c84e6e8..3a6ab8c 100644 --- a/scripts/bench_registration.py +++ b/scripts/bench_registration.py @@ -31,18 +31,20 @@ 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``); niftyreg runs natively (``conda install -c conda-forge niftyreg`` provides ``reg_aladin``); -``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. +``impact`` runs the ``impact-reg-konfai`` package (pip; GPU; one-time elastix +shim setup on hosts too old for the stock binary — see +REGISTRATION_BENCHMARK.md); ``synthmorph`` runs in a Docker image +(``freesurfer/synthmorph``; needs Docker + the pulled image). 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 + niftyreg) + # main env (SimpleITK + elastix + antspyx + niftyreg + impact) source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 - python scripts/bench_registration.py --pipelines simpleitk elastix antspyx niftyreg + python scripts/bench_registration.py --pipelines simpleitk elastix antspyx niftyreg impact - # synthmorph / impact (GPU, Docker) — one-time setup in REGISTRATION_BENCHMARK.md - python scripts/bench_registration.py --pipelines synthmorph impact + # synthmorph (GPU, Docker) — one-time setup in REGISTRATION_BENCHMARK.md + python scripts/bench_registration.py --pipelines synthmorph # any pipeline that is not importable in the current env is recorded as # unavailable, so the script is env-portable; --merge combines the @@ -246,7 +248,7 @@ def _call(pipeline, ct, t1c, tv, args): 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, + return RP.register_impact(ct, t1c, preset=args.impact_preset, gpu=args.impact_gpu, iters=args.impact_iters, samples=args.impact_samples) if pipeline == "niftyreg": return RP.register_niftyreg(ct, t1c, threads=args.threads, @@ -494,18 +496,24 @@ def parse_args(): 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") + ap.add_argument("--impact-preset", default=RP.IMPACT_DEFAULT_PRESET, + help="ImpactReg preset from the published database " + f"(default: {RP.IMPACT_DEFAULT_PRESET}; e.g. " + "Elastix_IMPACT_Jacobian, Generic_Rigid, " + "Generic_Rigid_BSpline, FireANTs_SyN, FireANTs_IMPACT, " + "ConvexAdam_Coarse/Composite; list live with: " + "impact-reg-konfai register --help)") ap.add_argument("--impact-gpu", type=int, default=None, help="GPU device index for IMPACT (default: auto-spread over " "the host GPUs by pid; -1 forces CPU)") - ap.add_argument("--impact-iters", type=int, default=2000, - help="IMPACT MaximumNumberOfIterations (default 2000, the benchmark " - "budget; ~10x the authors' 500-iter Jacobian budget, much slower)") - ap.add_argument("--impact-samples", type=int, default=5000, - help="IMPACT NumberOfSpatialSamples per iteration (default 5000, the " - "elastix default)") + ap.add_argument("--impact-iters", type=int, default=None, + help="preset max_iterations override, one value for all " + "resolutions (only effective for presets without a " + "per-resolution model matrix; default: the preset's " + "own budget, 400/300/200/200 for Elastix_IMPACT_Static)") + ap.add_argument("--impact-samples", type=int, default=None, + help="preset spatial_samples override (NumberOfSpatialSamples; " + "default: the preset's own value, 2000)") ap.add_argument("--niftyreg-backend", choices=["auto", "conda", "docker"], default="auto", help="niftyreg reg_aladin build: docker = official v2.1.1 in the " diff --git a/scripts/setup_impact_konfai.sh b/scripts/setup_impact_konfai.sh new file mode 100644 index 0000000..31e2352 --- /dev/null +++ b/scripts/setup_impact_konfai.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +# One-time host setup for the IMPACT registration pipeline (impact-reg-konfai). +# +# The stock elastix-impact 1.0.0 binary (vboussot/ImpactElastix) needs +# glibc >= 2.35 and links LibTorch 2.8.0. This host (Ubuntu 20.04, glibc 2.31, +# env torch 2.14) satisfies neither, so this script builds a self-contained +# shim at ~/.cache/konfai/elastix-impact-shim: +# * glibc 2.35 loader + libs, extracted from the niftyreg:2.1.1-ubuntu image +# (ubuntu 22.04 base) -- docker is needed ONLY for this one-time copy; +# * LibTorch 2.8.0 cu128, self-contained (its own CUDA runtime), from +# download.pytorch.org; +# * a wrapper bin/elastix that runs the stock binary through the glibc 2.35 +# loader with the LibTorch 2.8.0 library path (ignoring the inherited +# LD_LIBRARY_PATH, so the env's torch 2.14 cannot leak in). +# The pipeline selects the shim via KONFAI_ELASTIX_DIR +# (src/registration_pipelines.py::_impact_elastix_root). +# +# Run from the repo root, inside the conda env: +# source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 +# pip install impact-reg-konfai # the CLI itself +# bash scripts/setup_impact_konfai.sh # this script +# +# First `register` afterwards auto-downloads from HF/GitHub: the preset app +# (VBoussot/ImpactReg), the feature models (VBoussot/impact-torchscript-models) +# and, on hosts that CAN run the stock binary directly, the binary itself. +set -euo pipefail + +KONFAI_CACHE="${KONFAI_CACHE:-$HOME/.cache/konfai}" +SHIM="$KONFAI_CACHE/elastix-impact-shim" +STOCK="$KONFAI_CACHE/elastix-impact" +GLIBC_SRC_DIR="${GLIBC_SRC_DIR:-glibc-2.35}" +GLIBC_IMG="${GLIBC_IMG:-niftyreg:2.1.1-ubuntu}" +LT_VER="2.8.0" +LT_DIR="libtorch-${LT_VER}-cu128" +LT_URL="https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-${LT_VER}%2Bcu128.zip" + +echo "[1/5] stock elastix-impact binary (package auto-install)" +if [ ! -x "$STOCK/bin/elastix" ]; then + python - <<'PY' +from pathlib import Path +from impact_reg_konfai.models.elastix_install import ( + ELASTIX_CACHE, get_elastix_bin, install_elastix_impact) +if not get_elastix_bin(ELASTIX_CACHE).exists(): + # flavor auto-detected (cu128 on this host); downloads + extracts the zip + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) +PY +fi +[ -x "$STOCK/bin/elastix" ] || { echo "stock binary missing: $STOCK/bin/elastix"; exit 1; } + +echo "[2/5] glibc 2.35 from $GLIBC_IMG (docker, one-time)" +if [ ! -x "$KONFAI_CACHE/$GLIBC_SRC_DIR/ld-linux-x86-64.so.2" ]; then + docker run --rm --entrypoint sh -v "$KONFAI_CACHE":/dest "$GLIBC_IMG" -c \ + "rm -rf /dest/$GLIBC_SRC_DIR && mkdir -p /dest/$GLIBC_SRC_DIR && \ + cp -a /lib/x86_64-linux-gnu/. /dest/$GLIBC_SRC_DIR/ && \ + cp -a /usr/lib/x86_64-linux-gnu/libstdc++.so.6* /dest/$GLIBC_SRC_DIR/ 2>/dev/null || true" +fi + +echo "[3/5] LibTorch ${LT_VER} cu128 (self-contained)" +if [ ! -f "$SHIM/lib/$LT_DIR/libtorch/lib/libc10.so" ]; then + mkdir -p "$SHIM/lib" + curl -L --fail -o "$KONFAI_CACHE/$LT_DIR.zip" "$LT_URL" + unzip -q "$KONFAI_CACHE/$LT_DIR.zip" -d "$SHIM/lib/$LT_DIR" + rm "$KONFAI_CACHE/$LT_DIR.zip" +fi +# the with-deps zip carries the CUDA libs under hashed names; expose the public +# soname elastix links directly +LT_LIB="$SHIM/lib/$LT_DIR/libtorch/lib" +CUDART="$(ls "$LT_LIB"/libcudart-*.so.12 | head -1)" +ln -sf "$(basename "$CUDART")" "$LT_LIB/libcudart.so.12" + +echo "[4/5] shim layout + wrapper" +mkdir -p "$SHIM/bin" +cp "$STOCK/bin/elastix" "$SHIM/bin/elastix-real" +[ -d "$STOCK/lib" ] && cp -a "$STOCK/lib/." "$SHIM/lib/" +[ -d "$KONFAI_CACHE/$GLIBC_SRC_DIR" ] && cp -a "$KONFAI_CACHE/$GLIBC_SRC_DIR" "$SHIM/lib/$GLIBC_SRC_DIR" +cat > "$SHIM/bin/elastix" <<'WRAP' +#!/bin/sh +# elastix-impact 1.0.0 shim: the host glibc (2.31) is too old for the binary +# (built against glibc 2.35) and the binary expects LibTorch 2.8.0, which the +# conda env's torch (2.14) does not provide. Run it through the bundled +# glibc 2.35 loader with the bundled self-contained LibTorch 2.8.0 (cu128). +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +LT="$ROOT/lib/libtorch-2.8.0-cu128/libtorch/lib" +GLIBC="$ROOT/lib/glibc-2.35" +exec env -u LD_LIBRARY_PATH -u LD_PRELOAD \ + "$GLIBC/ld-linux-x86-64.so.2" \ + --library-path "$LT:$ROOT/lib:$GLIBC" \ + "$ROOT/bin/elastix-real" "$@" +WRAP +chmod +x "$SHIM/bin/elastix" + +echo "[5/5] verify" +KONFAI_ELASTIX_DIR="$SHIM" python - <<'PY' +import os +from pathlib import Path +from impact_reg_konfai.models.elastix_install import get_elastix_bin, loader_env, try_elastix +import subprocess +root = Path(os.environ["KONFAI_ELASTIX_DIR"]) +try_elastix(root) +r = subprocess.run([str(get_elastix_bin(root).resolve()), "--extended-version"], + capture_output=True, text=True, env=loader_env(root)) +assert r.returncode == 0, r.stderr +print(r.stdout.splitlines()[0] if r.stdout else "elastix OK") +print("elastix-impact shim OK:", root) +PY + +echo +echo "Done. The impact pipeline picks the shim up automatically" +echo "(src/registration_pipelines.py::_impact_elastix_root), e.g.:" +echo " python scripts/bench_registration.py --pipelines impact --n-cases 2" \ No newline at end of file diff --git a/src/registration_pipelines.py b/src/registration_pipelines.py index a3a8d1e..edb38ef 100644 --- a/src/registration_pipelines.py +++ b/src/registration_pipelines.py @@ -28,14 +28,17 @@ Shared rules for every pipeline: - **No mask, no crop / pad**: registration always runs on the full input volumes exactly as given (no ``FixedMask`` / ``MovingMask``, no head-box resampling, no padding). -- **CT is always the fixed image** (moving = T1c). When both inputs are MR - (no CT present) the first argument is kept as fixed (see - :func:`fixed_moving`). +- **The T1c (MR) is always the fixed image** (moving = CT). When both inputs + are the same class (both CT, or a clipped CT that no longer reads as CT) + the T1c-slot (second) argument is kept as fixed; see :func:`fixed_moving`. + Every pipeline -- including ``impact``, whose ImpactReg MR_CT convention + coincides with this rule -- registers fixed=T1c / moving=CT. - **CT brain-window clip for the pipelines that benefit from it**: per the CT-clipping A/B (see REGISTRATION_BENCHMARK.md), ``niftyreg`` and - ``synthmorph`` clip the CT to [0, 80] HU before registration - (``clip_ct`` parameter, default on, MR inputs never clipped); the - Mattes-MI pipelines (antspyx / elastix / simpleitk) register the CT as-is. + ``synthmorph`` clip the CT (the moving image) to [0, 80] HU before + registration (``clip_ct`` parameter, default on, MR inputs never clipped); + the Mattes-MI pipelines (antspyx / elastix / simpleitk) register the CT + as-is. Each pipeline uses its toolkit's **native automatic transform initialization** (no manual ``centroid_offset``): elastix / impact use elastix's @@ -50,12 +53,18 @@ default center-of-mass alignment, simpleitk starts from identity, niftyreg - ``register_elastix`` : ITK-Elastix Python API (pip ``itk`` + ``itk-elastix``, elastix 5.x: ``EulerTransform`` rigid + Mattes MI, ASGD). ``AutomaticTransformInitialization`` (CenterOfGravity) start. -- ``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). full volumes, fixed=CT/moving=T1c (no mask -- shared - rules above); elastix ``AutomaticTransformInitialization`` (CenterOfGravity) - provides the start. Needs Docker + the image + a model asset. +- ``register_impact`` : IMPACT (semantic-similarity metric) via the + ``impact-reg-konfai`` package (pip; the KonfAI runtime drives the + elastix-impact 5.2 binary -- no Docker). Default preset + ``Elastix_IMPACT_Static``: the published successor of the ImpactReg + ``MR_CT_TS`` reference map (elastix B-spline + per-level MIND + + TotalSegmentator M730, Static mode). Registered fixed=T1c/moving=CT (the + ImpactReg MR_CT convention) so the preset's output field (on the fixed + grid) is directly the T1c->CT contract transform; CenterOfGravity auto + init + the 0.1 sample-validity gate as raw-elastix overrides. GPU. + Full volumes, no mask (shared rules above). First use auto-downloads the + preset app, models and binary into ``~/.cache``; on hosts too old for the + stock binary, ``scripts/setup_impact_konfai.sh`` builds the shim. - ``register_niftyreg`` : NiftyReg (KCL-BMEIS suite) ``reg_aladin`` -- symmetric block-matching + trimmed least squares (Modat et al., J. Med. Img. 2014), NCC-based block similarity (cross-modality tolerant). Rigid @@ -84,8 +93,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 / synthmorph / impact, not to the trusted -production SimpleITK path. +applied to elastix / antspyx / synthmorph, not to the trusted production +SimpleITK path. ``impact`` returns a displacement field (which has no +``GetInverse``), so it uses :func:`_orient_impact` -- the same skull-ring +decision against a sign-flipped (first-order reverse) proxy field. """ import os import re @@ -145,14 +156,17 @@ def _modality_name(img): def fixed_moving(img_a, img_b): - """Shared fixed-image rule: the CT is always the fixed image. When both - inputs are MR (or both CT) the given order is kept, i.e. ``img_a`` is the - fixed image. Returns ``(fixed, moving)``. + """Shared fixed-image rule: the T1c (MR) is always the fixed image + (moving = CT). Returns ``(fixed, moving)``. - Note: a CT that has been intensity-clipped to a non-negative window (e.g. - [0, 80] HU) is no longer detectable as CT, so the given order is kept -- - callers should pass the CT first (as all pipelines do).""" - if _is_ct(img_b) and not _is_ct(img_a): + Callers pass the CT first (``img_a``) and the T1c second (``img_b``). + A CT that has been intensity-clipped to a non-negative window (e.g. + [0, 80] HU) is no longer detectable as CT, so when neither input looks + like CT the second (T1c-slot) image is the fixed one; when both look + like CT the given order is kept (``img_a`` fixed).""" + if _is_ct(img_a) and _is_ct(img_b): + return img_a, img_b + if _is_ct(img_a) or not _is_ct(img_b): return img_b, img_a return img_a, img_b @@ -165,17 +179,21 @@ def _clip_intensity(img, lo, hi): return out -def _maybe_clip_ct(fixed, clip_ct): - """Apply the pipeline's CT brain-window clip to the fixed image. +def _maybe_clip_ct_pair(fixed, moving, clip_ct): + """Apply the pipeline's CT brain-window clip to the CT of the pair + (whichever side it is on). The A/B benchmark (``--clip-ct 0 80``) showed clipping the CT to a brain window helps the block-matching / deep pipelines (niftyreg, synthmorph) - and hurts the Mattes-MI ones (antspyx, elastix, simpleitk), so those two - pipelines clip the CT by default. Returns ``(fixed, clipped: bool)``; - MR fixed images (both-MR pairs) are never clipped.""" + and hurts the Mattes-MI ones (antspyx / elastix / simpleitk), so those two + pipelines clip the CT by default. Under the T1c-fixed rule the CT is the + moving image; MR / undetectable images are never clipped. Returns + ``(fixed, moving, clipped: bool)``.""" + if clip_ct and _is_ct(moving): + return fixed, _clip_intensity(moving, *clip_ct), True if clip_ct and _is_ct(fixed): - return _clip_intensity(fixed, *clip_ct), True - return fixed, False + return _clip_intensity(fixed, *clip_ct), moving, True + return fixed, moving, False def fix_orientation(ct, t1c, cand): @@ -258,13 +276,12 @@ def register_simpleitk(ct, t1c, threads=12, mode="benchmark", iters_full=2000): except AttributeError: pass ct = sitk.Cast(ct, sitk.sitkFloat32); t1c = sitk.Cast(t1c, sitk.sitkFloat32) - # full volumes, no mask / crop / pad; CT is the fixed image (or, if both - # inputs are MR, the first argument is) + # full volumes, no mask / crop / pad; the T1c is the fixed image fixed, moving = fixed_moving(ct, t1c) shape = np.array(ct.GetSize()) # Identity start (no centroid init); the multi-resolution Mattes-MI pyramid # recovers the gross inter-session offset. The Euler center is the fixed - # (CT) geometric center. + # (T1c) geometric center. init = sitk.Euler3DTransform() init.SetCenter(tuple(float(x) for x in fixed.TransformIndexToPhysicalPoint( ((fixed.GetSize()[0] - 1) // 2, (fixed.GetSize()[1] - 1) // 2, @@ -284,7 +301,10 @@ def register_simpleitk(ct, t1c, threads=12, mode="benchmark", iters_full=2000): r.SetInterpolator(sitk.sitkLinear); r.SetInitialTransform(init, inPlace=False) tr0 = r.Execute(fixed, moving) dt = time.time() - t0 - tr = tr0.GetInverse() # contract: T1c -> CT + # Execute(fixed, moving) returns the transform that resamples moving onto + # the fixed grid (fixed -> moving). With fixed=T1c / moving=CT (the + # shared rule) that is T1c -> CT -- already the contract transform. + tr = tr0 ha = sitk.GetArrayFromImage(t1c) gap = m6._skull_ring_gap(_warp_ct(ct, t1c, tr), ha) meta = { @@ -352,7 +372,7 @@ def _itk_transform_to_sitk_rigid(itk_tf): def register_elastix(ct, t1c, workdir=None, threads=8): - """ITK-Elastix rigid CT->T1c (fixed=CT, moving=T1c), Mattes MI, full images. + """ITK-Elastix rigid CT->T1c (fixed=T1c, moving=CT), Mattes MI, full images. Runs in-process through the ITK-Elastix python API (pip ``itk`` + ``itk-elastix``, elastix 5.x). Rigid 6-DOF via ``EulerTransform`` (the @@ -372,7 +392,7 @@ def register_elastix(ct, t1c, workdir=None, threads=8): after a successful registration. """ itk = _elastix_api() - # full volumes, no mask / crop / pad; CT is the fixed image + # full volumes, no mask / crop / pad; the T1c is the fixed image fixed, moving = fixed_moving(ct, t1c) params = { "ImageDimension": 3, @@ -414,9 +434,9 @@ def register_elastix(ct, t1c, workdir=None, threads=8): tail = "\n".join(buf.getvalue().splitlines()[-10:]) raise RuntimeError(f"elastix registration failed: {e}\n{tail}") from e dt = time.time() - t0 - # The combination transform maps fixed(CT) -> moving(T1c) (elastix - # convention), i.e. the inverse of the contract (T1c->CT); - # fix_orientation resolves the direction against the skull ring. + # The combination transform maps fixed -> moving (elastix convention); + # with fixed=T1c / moving=CT that is T1c -> CT, the contract direction. + # fix_orientation verifies it against the skull ring. cand = _itk_transform_to_sitk_rigid(reg.GetCombinationTransform()) T, inverted = fix_orientation(ct, t1c, cand) try: @@ -442,7 +462,7 @@ def register_elastix(ct, t1c, workdir=None, threads=8): def register_antspyx(ct, t1c, transform_type="Rigid", aff_iterations=None, aff_shrink_factors=None, aff_smoothing_sigmas=None, aff_metric="mattes"): - """ANTs rigid CT->T1c via antspyx (fixed=CT, moving=T1c), full volumes. + """ANTs rigid CT->T1c via antspyx (fixed=T1c, moving=CT), full volumes. antspyx 0.6.3: ``ants.registration(fixed, moving, type_of_transform='Rigid', aff_metric='mattes', aff_iterations=..., aff_shrink_factors=..., @@ -465,7 +485,7 @@ def register_antspyx(ct, t1c, transform_type="Rigid", aff_iterations=None, 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) - # full volumes, no mask / crop / pad; CT is the fixed image + # full volumes, no mask / crop / pad; the T1c is the fixed image fx, mv = fixed_moving(ct, t1c) import time t0 = time.time() @@ -478,11 +498,10 @@ def register_antspyx(ct, t1c, transform_type="Rigid", aff_iterations=None, aff_shrink_factors=aff_shrink_factors, aff_smoothing_sigmas=aff_smoothing_sigmas) dt = time.time() - t0 - # r['fwdtransforms'][0] is the ANTs forward transform: fixed(CT) -> moving - # (T1c). The contract is T1c -> CT, i.e. its inverse; fix_orientation - # verifies the direction against the skull ring. - T_ants = sitk.ReadTransform(r["fwdtransforms"][0]) - cand = T_ants.GetInverse() + # r['fwdtransforms'][0] is the ANTs forward transform: fixed -> moving. + # With fixed=T1c / moving=CT that is T1c -> CT, the contract direction; + # fix_orientation verifies it against the skull ring. + cand = sitk.ReadTransform(r["fwdtransforms"][0]) T, inverted = fix_orientation(ct, t1c, cand) meta = { "pipeline": "antspyx", "toolkit": f"ANTsPy {getattr(ants, '__version__', '?')}", @@ -544,7 +563,7 @@ def _lta_to_sitk(path): def register_synthmorph(ct, t1c, model="rigid", gpu=None, threads=None, extent=256, keep_scratch=False, clip_ct=(0.0, 80.0)): - """SynthMorph CT->T1c (fixed=CT, moving=T1c) in the Docker image + """SynthMorph CT->T1c (fixed=T1c, moving=CT) in the Docker image ``freesurfer/synthmorph`` (TensorFlow ``VxmAffineFeatureDetector``, pretrained ``synthmorph.rigid.1.h5`` weights). @@ -586,11 +605,12 @@ def register_synthmorph(ct, t1c, model="rigid", gpu=None, threads=None, # per-case scratch on a local (non-NFS) tmp fs tmp = tempfile.mkdtemp(prefix="synthmorph_reg_") try: - # full volumes, no mask / crop / pad; CT is the fixed image + # full volumes, no mask / crop / pad; the T1c is the fixed image fixed, moving = fixed_moving(ct, t1c) - fixed_name = _modality_name(fixed) # before clipping (a clipped CT + fixed_name = _modality_name(fixed) + moving_name = _modality_name(moving) # before clipping (a clipped CT # is no longer intensity-detectable) - fixed, ct_clipped = _maybe_clip_ct(fixed, clip_ct) + fixed, moving, ct_clipped = _maybe_clip_ct_pair(fixed, moving, clip_ct) sitk.WriteImage(sitk.Cast(fixed, sitk.sitkFloat32), os.path.join(tmp, "fixed.nii.gz"), True) sitk.WriteImage(sitk.Cast(moving, sitk.sitkFloat32), @@ -609,9 +629,11 @@ def register_synthmorph(ct, t1c, model="rigid", gpu=None, threads=None, cmd.append("-g") if threads: cmd += ["-j", str(threads)] + # the model estimates first -> second, so (fixed, moving) puts the + # -T LTA in the fixed->moving = T1c->CT contract direction cmd += ["-e", str(extent), "-T", "/data/f2m.lta", - "/data/moving.nii.gz", "/data/fixed.nii.gz"] + "/data/fixed.nii.gz", "/data/moving.nii.gz"] t0 = time.time() proc = subprocess.run(cmd, capture_output=True, text=True, timeout=7200) dt = time.time() - t0 @@ -636,7 +658,7 @@ def register_synthmorph(ct, t1c, model="rigid", gpu=None, threads=None, f"device={'GPU' + str(gpu) if gpu >= 0 else 'CPU'}, " f"extent {extent})"), "convention": "T1c->CT (verified)", "auto_inverted": inverted, - "fixed": fixed_name, "moving": _modality_name(moving), + "fixed": fixed_name, "moving": moving_name, "seconds_reg": round(dt, 1), "docker_image": SYNTHMORPH_IMAGE, "model": model, "gpu": gpu, "extent": extent, @@ -648,31 +670,91 @@ def register_synthmorph(ct, t1c, model="rigid", gpu=None, threads=None, # --------------------------------------------------------------------------- # -# 5. IMPACT (Docker image vboussot/elastix_impact) +# 5. IMPACT (impact-reg-konfai -- KonfAI runtime, no Docker) # --------------------------------------------------------------------------- # -IMPACT_IMAGE = "vboussot/elastix_impact" -IMPACT_ELASTIX_BIN = "/lib/elastix-install/bin/elastix" -IMPACT_MODEL_CACHE = "data/impact_models" -IMPACT_DEFAULT_MODEL = "MRSeg/MRSeg.pt" -IMPACT_HF_MODEL_URL = ( - "https://huggingface.co/VBoussot/impact-torchscript-models/resolve/main/{rel}" +IMPACT_CLI = "impact-reg-konfai" +# default preset from the published ImpactReg database (VBoussot/ImpactReg, +# resolved by the CLI): elastix B-spline driven by the IMPACT semantic metric +# -- per-level MIND (L1) + TotalSegmentator M730 (deepest layer, Dice), Static +# mode (whole-image features, no backprop), 4 scales, 400/300/200/200 +# iterations, 2000 spatial samples, + Mattes MI (w 0.4) + bending energy. +# The published successor of the old MR_CT_TS reference parameter map. +IMPACT_DEFAULT_PRESET = "Elastix_IMPACT_Static" +# M6 robustness overrides -- raw elastix parameter lines the package applies +# on top of the preset's generated map (the ``parameter_overrides`` knob): +# * CenterOfGravity auto-init (instead of the preset's GeometricalCenter): +# aligns the intensity centroids -- the T1c head-crop and the CT head+neck +# FOV have their buffer centres ~50-100 mm apart, and the GeometricalCenter +# start leaves the metric with almost nothing in-buffer; +# * RequiredRatioOfValidSamples 0.1 (build default 0.25): the FOV-mismatched +# pairs keep only a minority of the random spatial samples in-buffer at the +# init; 0.1 keeps the divergence gate while letting the runs proceed (a +# genuinely diverging transform, <10% in-buffer, still aborts). +IMPACT_PARAM_OVERRIDES = ( + "AutomaticTransformInitialization=true", + "AutomaticTransformInitializationMethod=CenterOfGravity", + "RequiredRatioOfValidSamples=0.1", ) +# The stock elastix-impact 1.0.0 binary (vboussot/ImpactElastix) needs +# glibc >= 2.35 and links LibTorch 2.8.0; a host that cannot run it directly +# (this one: glibc 2.31 + env torch 2.14) builds a self-contained shim with +# scripts/setup_impact_konfai.sh, selected via KONFAI_ELASTIX_DIR. +IMPACT_ELASTIX_SHIM = os.path.join( + os.path.expanduser("~"), ".cache", "konfai", "elastix-impact-shim") +IMPACT_ELASTIX_STOCK = os.path.join( + os.path.expanduser("~"), ".cache", "konfai", "elastix-impact") +# orientation-gate threshold: a candidate DVF is kept as-is when its +# skull-ring gap score (std_mm + mid_mm) is below this. A misdirected field +# lands the head on the far side of the box (gap ~100-300 mm), so the margin +# is wide; see _orient_impact. +IMPACT_GAP_ACCEPT = 60.0 -def _impact_docker_ok(): - """Return (ok, note): is the Docker CLI + the IMPACT image available?""" - if shutil.which("docker") is None: - return False, "docker CLI not found on PATH" +def _impact_cli_ok(): + """Return (ok, note): is the impact-reg-konfai CLI installed?""" + if shutil.which(IMPACT_CLI) is None: + return False, f"{IMPACT_CLI} not on PATH (pip install impact-reg-konfai)" try: - r = subprocess.run(["docker", "image", "inspect", IMPACT_IMAGE], - capture_output=True, text=True, timeout=30) - if r.returncode != 0: - return False, (f"docker image {IMPACT_IMAGE} not pulled " - f"(run: docker pull {IMPACT_IMAGE})") - except Exception as e: # noqa - return False, f"docker image check failed ({type(e).__name__}: {e})" - return True, IMPACT_IMAGE + import importlib.metadata + ver = importlib.metadata.version("impact-reg-konfai") + except Exception: # noqa + ver = "?" + return True, f"{IMPACT_CLI} {ver}" + + +_ELASTIX_ROOT_CACHE = None + + +def _impact_elastix_root(): + """The elastix-impact install root that runs on this host, or None. + + Tried in order: a caller-set ``KONFAI_ELASTIX_DIR``, the setup-script shim + (glibc 2.35 loader + LibTorch 2.8.0, for hosts too old for the stock + binary), and the package's own auto-install cache (populated by the first + run when the host can run the stock binary directly). Each candidate is + actually executed (``elastix -h``) so a stale/broken install is skipped.""" + global _ELASTIX_ROOT_CACHE + if _ELASTIX_ROOT_CACHE is not None: + return _ELASTIX_ROOT_CACHE + from pathlib import Path + from impact_reg_konfai.models import elastix_install as _ei + + root = os.environ.get("KONFAI_ELASTIX_DIR", "") + candidates = [c for c in (root, IMPACT_ELASTIX_SHIM, IMPACT_ELASTIX_STOCK) + if c] + found = None + for c in candidates: + if not os.path.exists(os.path.join(c, "bin", "elastix")): + continue + try: + _ei.try_elastix(Path(c)) + found = c + break + except Exception: # noqa + continue + _ELASTIX_ROOT_CACHE = found + return found _NGPUS_CACHE = None @@ -691,200 +773,156 @@ def _n_gpus(): return _NGPUS_CACHE -def _impact_model_host(model): - """Resolve a model to a file under the IMPACT model cache, returning - ``(host_path, rel)`` where ``rel`` is the path relative to the cache (which - the pipeline bind-mounts at ``/Data/Models`` inside the container). - - ``model`` is either a path relative to the HF model layout (e.g. - ``TS/M730.pt``) or an absolute host path. An absolute path outside the - cache is copied into it (once) so the single ``/Data/Models`` mount covers - it. Raises :class:`PipelineUnavailable` if the model is missing, with the - exact download command for the HF-relative case. - """ - rel = model or IMPACT_DEFAULT_MODEL - cache = d(IMPACT_MODEL_CACHE) - if os.path.isabs(rel): - src = rel - dst = os.path.join(cache, os.path.basename(rel)) - if not os.path.exists(dst): - if not os.path.exists(src): - raise PipelineUnavailable(f"impact: model file not found: {src}") - os.makedirs(os.path.dirname(dst), exist_ok=True) - shutil.copy2(src, dst) - rel = os.path.basename(rel) - host = os.path.join(cache, rel) - if not os.path.exists(host): - url = IMPACT_HF_MODEL_URL.format(rel=rel) - raise PipelineUnavailable( - f"impact: model not found at {host}. Download it once with:\n" - f" mkdir -p {os.path.dirname(host)} && curl -L -o {host} {url}\n" - f"(or `sh Data/Models/download_models.sh` from the ImpactLoss repo).") - return host, rel +def _impact_dvf_from_h5(path): + """An ITK-form HDF5 transform (the ``:itktransform`` form the KonfAI + presets write) read with h5py -- the SimpleITK build in this env has no + ITK HDF5 transform reader. Returns ``(DisplacementFieldTransform, + dvf_image)``; the field is in mm, channel-last ``(z, y, x, 3)``.""" + import h5py + with h5py.File(str(path), "r") as f: + grp = f["TransformGroup/0"] + ttype = grp["TransformType"][()][0] + if isinstance(ttype, bytes): + ttype = ttype.decode() + fixed = grp["TransformFixedParameters"][()] + params = grp["TransformParameters"][()] + if not str(ttype).startswith("DisplacementFieldTransform"): + raise RuntimeError(f"impact: unsupported h5 transform type {ttype!r}") + size = fixed[:3][::-1].astype(int) # stored (x, y, z) -> array (z, y, x) + arr = params.reshape(int(size[0]), int(size[1]), int(size[2]), 3) + img = sitk.GetImageFromArray(arr) + img.SetOrigin(tuple(float(v) for v in fixed[3:6])) + img.SetSpacing(tuple(float(v) for v in fixed[6:9])) + img.SetDirection(tuple(float(v) for v in fixed[9:18])) + # this SimpleITK build invalidates the image wrapper handed to the DVF + # constructor (it comes back as an empty uint8 image), so build the + # transform from a copy and keep ``img`` for the caller + return sitk.DisplacementFieldTransform(sitk.Image(img)), img -def _impact_param_map(model_cont_path, gpu, iters, samples, voxel_sizes, - patch=11, subset=64, rescale=None, smoothing=None): - """Build the Elastix-5.2 parameter file for a RIGID (Euler) IMPACT - registration with a multi-scale semantic metric. By default the image - pyramid is NOT resampled (all-1 schedules) and the scale comes from the - per-level ImpactVoxelSize (as the IMPACT docs recommend); pass explicit - ``rescale``/``smoothing`` space-separated schedules to use a real - coarse-to-fine shrinking pyramid instead.""" - n = len(voxel_sizes) - ones = " ".join(["1"] * (n * 3)) - rescale = rescale or ones - smoothing = smoothing or ones - mixed = "true" if gpu >= 0 else "false" - L = ["// IMPACT semantic-similarity metric (vboussot/ImpactLoss)"] - L.append('(Metric "Impact")') - for i, vs in enumerate(voxel_sizes): - L += [ - f'(ImpactModelsPath{i} "{model_cont_path}")', - f"(ImpactDimension{i} 3)", - f"(ImpactNumberOfChannels{i} 1)", - f"(ImpactPatchSize{i} {patch} {patch} {patch})", - f"(ImpactVoxelSize{i} {vs:g} {vs:g} {vs:g})", - f'(ImpactLayersMask{i} "01")', - f"(ImpactPCA{i} 0)", - f"(ImpactSubsetFeatures{i} {subset})", - f'(ImpactDistance{i} "L2")', - f"(ImpactLayersWeight{i} 1)", - ] - L += [ - '(ImpactMode "Jacobian")', - f"(ImpactGPU {gpu})", - f'(ImpactUseMixedPrecision "{mixed}")', - "(ImpactFeaturesMapUpdateInterval -1)", - '(ImpactWriteFeatureMaps "false")', - # safety gate: abort if < ratio of the spatial samples map inside the - # moving buffer. The build default is 0.25, but the no-mask M6 pairs - # (CT head-neck FOV vs the tight T1c crop) only have ~14% in-buffer at - # the CenterOfGravity auto-init, so elastix aborts at the first metric - # evaluation. 0.1 keeps the gate (a diverging transform still aborts) - # while letting the FOV-mismatched runs proceed -- the metric natively - # evaluates only the in-buffer samples. - "(RequiredRatioOfValidSamples 0.1)", - "", - "// rigid 6-DOF to be optimized (auto center init, no -t0)", - '(Transform "EulerTransform")', - '(TransformBase "ElastixTransformBase")', - '(TransformParameterReadoutFormat "AllParametersOneLine")', - '(FixedCenteredTransform "true")', - '(HowToCombineTransforms "Compose")', - '(AutomaticTransformInitialization "true")', - # align intensity centroids (heads), not buffer corners -- the CT FOV is - # much larger than the T1c crop, so GeometricalCenter leaves the heads - # ~100 mm apart and the metric aborts ("too many samples outside buffer") - '(AutomaticTransformInitializationMethod "CenterOfGravity")', - '(AutomaticScalesEstimation "true")', - "", - "(ImageDimension 3)", - '(FixedInternalImagePixelType "float")', - '(MovingInternalImagePixelType "float")', - '(UseDirectionCosines "true")', - ] - L += [ - '(Registration "MultiMetricMultiResolutionRegistration")', - f"(NumberOfResolutions {n})", - f"(MaximumNumberOfIterations {iters})", - f"(NumberOfSpatialSamples {samples})", - "", - '(FixedImagePyramid "FixedGenericImagePyramid")', - '(MovingImagePyramid "MovingGenericImagePyramid")', - f"(FixedImagePyramidRescaleSchedule {rescale})", - f"(MovingImagePyramidRescaleSchedule {rescale})", - f"(FixedImagePyramidSmoothingSchedule {smoothing})", - f"(MovingImagePyramidSmoothingSchedule {smoothing})", - "", - '(Optimizer "AdaptiveStochasticGradientDescent")', - '(UseAdaptiveStepSizes "true")', - '(UseMultiThreadingForMetrics "true")', - "(MaximumNumberOfSamplingAttempts 8)", - "(NumberOfSamplesForExactGradient 4096)", - "(SigmoidInitialTime 0.0)", - '(NoiseCompensation "true")', - '(ASGDParameterEstimationMethod "DisplacementDistribution")', - "", - '(ImageSampler "Random")', - '(NewSamplesEveryIteration "true")', - '(Interpolator "BSplineInterpolator")', - "(BSplineInterpolationOrder 3)", - "", - '(WriteResultImage "false")', - "(DefaultPixelValue -1024)", - '(ResultImageFormat "mha")', - '(ITKTransformOutputFileNameExtension "itk.txt")', - '(WriteTransformParametersEachIteration "false")', - '(WriteTransformParametersEachResolution "false")', - ] - return "\n".join(L) + "\n" +def _impact_read_transform(path): + """The preset's single-case transform output as + ``(sitk.Transform, dvf_image_or_None)``. + + KonfAI presets write the registration in whatever form the preset + declares: an ITK-form HDF5 DisplacementFieldTransform (``Transform.h5``) + by default, or an ITK transform file (``.itk.txt`` / ``.tfm``) or a vector + image (``.mha``) for other presets.""" + from pathlib import Path + p = Path(path) + if p.name.endswith(".itk.txt") or p.suffix in (".h5", ".tfm"): + if p.suffix == ".h5": + return _impact_dvf_from_h5(p) + return sitk.ReadTransform(str(p)), None + img = sitk.ReadImage(str(p)) + return (sitk.DisplacementFieldTransform( + sitk.Cast(img, sitk.sitkVectorFloat64)), img) -def register_impact(ct, t1c, model=None, gpu=None, iters=2000, samples=5000, - threads=12, keep_scratch=False, scales=(6.0, 3.0, 1.5, 1.0), - downsample=False): - """IMPACT rigid CT->T1c (fixed=CT, moving=T1c) in the Docker image - ``vboussot/elastix_impact`` (Elastix 5.2.0 + IMPACT semantic metric). +def _orient_impact(ct, t1c, cand, dvf_img): + """``(T, inverted)`` in the contract direction, for an impact candidate. + + ``cand`` is nominally T1c->CT: the preset ran with fixed=T1c / moving=CT + and its output field is on the fixed (T1c) grid. Displacement fields + cannot be inverted with ``GetInverse`` (ITK), so the direction gate does + not compare ``cand`` with ``cand^-1`` (see :func:`fix_orientation`); it + compares the skull-ring gap of ``cand`` against a reverse proxy -- the + sign-flipped field (first-order inverse, exact enough for the decision: + a misdirected DVF lands the head on the far side of the box, gap + ~100+ mm instead of tens of mm) -- and falls back to it only when the + nominal direction is clearly wrong.""" + ha = sitk.GetArrayFromImage(t1c).astype("float32") + s0 = _gap_score(_skull_gap(_warp_ct(ct, t1c, cand), ha)) + if s0 <= IMPACT_GAP_ACCEPT: + return cand, False + if dvf_img is not None: + img = sitk.GetImageFromArray(-sitk.GetArrayViewFromImage(dvf_img)) + img.CopyInformation(dvf_img) + rev = sitk.DisplacementFieldTransform(img) + else: + try: + rev = cand.GetInverse() + except Exception: # noqa: parametric transform with no inverse + return cand, False + s1 = _gap_score(_skull_gap(_warp_ct(ct, t1c, rev), ha)) + if s1 < s0: + return rev, True + return cand, False + + +def register_impact(ct, t1c, preset=None, gpu=None, iters=None, samples=None, + keep_scratch=False): + """IMPACT (semantic-similarity metric) CT->T1c via the + ``impact-reg-konfai`` package (pip; KonfAI runtime driving the + elastix-impact 5.2 binary; no Docker). IMPACT (arXiv:2503.24121) replaces the intensity metric with a *semantic* - one: a TorchScript segmentation model (``MRSeg/MRSeg.pt`` by default, - MRSegmentator -- robust multi-modality segmentation on MRI *and* CT, the - right prior for this cross-modality CT->T1c task; TotalSegmentator - ``TS/M730.pt`` and the other models in the HF repo - ``VBoussot/impact-torchscript-models`` work too via ``model=``) is run on - randomly sampled patches of the fixed and moving images and the - registration minimizes the distance between the extracted feature maps - (L2, per the recommended config). Intensity standardization (zero mean / - unit variance per the model card) is applied *inside the metric*: at init - it computes whole-volume [min, max, mean, sigma] per image and the model's - normalize branch rescales every patch with that image's own (mean, sigma), - so no pipeline-side pre-normalization is needed (and would be an exact - no-op -- z-scoring is affine-invariant). This is a genuinely different - metric from the Mattes-MI pipelines above -- it is the cross-modality tool - the image was built for. + one. The default preset ``Elastix_IMPACT_Static`` is the published + successor of the old ImpactReg ``MR_CT_TS`` reference parameter map: an + elastix **B-spline** (4 scales, FinalGridSpacing 14 mm, 400/300/200/200 + iterations, 2000 spatial samples, bending-energy regularizer) driven by + the IMPACT metric -- a per-level **MIND** (L1) + **TotalSegmentator M730** + (deepest layer, Dice) combination in **Static** mode (whole-image + features, no backprop), plus a Mattes-MI secondary metric (w 0.4). Other + presets (``Elastix_IMPACT_Jacobian``, ``Generic_Rigid``, + ``Generic_Rigid_BSpline``, ``FireANTs_*``, ``ConvexAdam_*``) are + selectable with ``preset=``; the live database is listed by + ``impact-reg-konfai register --help``. - The pipeline registers on the full CT / T1c volumes (no head box) with no - manual centroid init: elastix's ``AutomaticTransformInitialization`` - (CenterOfGravity) aligns the intensity centroids. It uses a rigid 6-DOF - ``EulerTransform`` (the 5.x name for the classic 3D rigid), ASGD, a 4-level - scale pyramid (no image resampling). +Fixed/moving: the pair is registered in the ImpactReg **MR_CT** + convention (fixed=T1c, moving=CT), which coincides with the shared + fixed-image rule (the T1c is the fixed image). The preset's output is a + displacement field on the **fixed** grid mapping fixed->moving, so the + field comes out directly in the contract direction T1c->CT (displacement + fields cannot be inverted with ``GetInverse``, which the fixed=CT + direction would have required). That also keeps the metric's + sample-validity gate healthy: samples drawn from the tight T1c head crop + (fixed) map mostly inside the large CT head+neck FOV (moving), unlike + fixed=CT where only ~14% of the samples are in-buffer at the auto-init. + The transform is the preset's native deformable B-spline (a DVF) -- the + contract (warp the CT concentric with the MR head) is unchanged. - Like every other pipeline this runs with **fixed=CT, moving=T1c** on the - full volumes: no mask, no crop / pad (shared rules above). The IMPACT - metric draws random spatial samples from the fixed (CT) buffer and - evaluates them only where they map inside the moving (T1c) buffer; when - too few samples are in-buffer, elastix aborts ("too many samples map - outside moving image buffer"). On the M6 head-neck CTs (FOV - ~300x300x251 mm) vs the tight T1c head crop (~220x220x150 mm) only ~14% - of the samples are in-buffer at the CenterOfGravity auto-init, so the - M6 runs abort -- IMPACT needs a fixed volume whose FOV is comparable to - the moving crop (e.g. a head-only CT). The returned transform satisfies - the unified contract (T1c->CT). - Budget: the default ``iters=2000`` (MaximumNumberOfIterations) / - ``samples=5000`` (NumberOfSpatialSamples, the elastix default) runs the - full benchmark budget for best results -- roughly 10x the authors' - 500/2000 Jacobian budget, i.e. hours per case on full volumes. Lower - ``iters``/``samples`` to trade accuracy for speed. - The result is read back from ``TransformParameters.0.itk.txt`` - (an ``Euler3DTransform`` in the world physical frame) and checked by - :func:`fix_orientation`. + M6 robustness overrides (raw elastix lines the package applies on top of + the preset's generated map): CenterOfGravity auto init (the preset's + GeometricalCenter start leaves the intensity centroids ~50-100 mm apart + on the head+neck FOV mismatch) and ``RequiredRatioOfValidSamples 0.1`` + (the build default 0.25 aborts the FOV-mismatched pairs at the init; 0.1 + keeps the divergence gate). - Preconditions (else :class:`PipelineUnavailable`): the ``docker`` CLI, the - ``vboussot/elastix_impact`` image, and the model asset under the cache. A - GPU is used when one is visible (auto-spread over the host GPUs by pid); - pass ``gpu=-1`` to force CPU. + The package auto-downloads its one-time assets from HF/GitHub on first + use: the preset app (``VBoussot/ImpactReg``), the TorchScript feature + models (``VBoussot/impact-torchscript-models``) and the elastix-impact + binary (``vboussot/ImpactElastix``) into ``~/.cache``. The binary needs + glibc >= 2.35 and LibTorch 2.8.0; on a host that cannot run it directly + (this one: glibc 2.31 + torch 2.14) ``scripts/setup_impact_konfai.sh`` + builds a self-contained shim, selected via ``KONFAI_ELASTIX_DIR``. + + Budget: ``iters`` / ``samples`` forward to the preset's + ``max_iterations`` / ``spatial_samples`` knobs (only effective for + presets without a per-resolution model matrix; the default preset's + matrix wins). A GPU is used when one is visible (auto-spread over the + host GPUs by pid); pass ``gpu=-1`` to force CPU (the engine then turns + mixed precision off). The result is read back from + ``/P000/.h5`` (an ITK-form DisplacementFieldTransform on the + T1c grid) and direction-checked by :func:`_orient_impact` (skull-ring + gap vs the reverse-proxy field). """ import time - ok, note = _impact_docker_ok() + ok, note = _impact_cli_ok() if not ok: raise PipelineUnavailable(f"impact: {note}") - model_host, model_rel = _impact_model_host(model) + root = _impact_elastix_root() + if root is None: + raise PipelineUnavailable( + "impact: no runnable elastix-impact binary on this host (the stock " + "build needs glibc >= 2.35 + LibTorch 2.8.0). Run " + "scripts/setup_impact_konfai.sh once, or point KONFAI_ELASTIX_DIR " + "at a working install.") + preset = preset or IMPACT_DEFAULT_PRESET - # full volumes, no mask / crop / pad; CT is the fixed image (or, if both - # inputs are MR, the first argument). No manual init -- elastix's - # AutomaticTransformInitialization aligns the image centers. + # fixed=T1c (MR), moving=CT: see the docstring (contract direction comes + # out of the preset directly; MR_CT convention = the shared rule) fixed, moving = fixed_moving(ct, t1c) # GPU: auto-spread across the host GPUs by pid, or forced value @@ -892,105 +930,78 @@ def register_impact(ct, t1c, model=None, gpu=None, iters=2000, samples=5000, 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 - # the finest level matches the moving image's base spacing - base = float(np.median(moving.GetSpacing())) - n = len(scales) - voxel_sizes = [base * s for s in scales] - # image pyramid: IMPACT's docs recommend NO image resampling (all-1 schedules) - # with the scale carried by ImpactVoxelSize; `downsample` instead uses a real - # shrinking pyramid (2x per level) with ImpactVoxelSize matching the resampled - # spacing (classical coarse-to-fine). - if downsample: - rescale = [] - smoothing = [] - for i in range(n): - f = 2.0 ** (n - 1 - i) - rescale += [f, f, f] - s = n - 1 - i - smoothing += [s, s, s] - schedule = " ".join(f"{x:g}" for x in rescale) - smoothe = " ".join(f"{x:g}" for x in smoothing) - voxel_sizes = [base * (2.0 ** (n - 1 - i)) for i in range(n)] - else: - schedule = " ".join(["1"] * (n * 3)) - smoothe = schedule - - # per-case scratch on a local (non-NFS) tmp fs: Data/ + Out/ tmp = tempfile.mkdtemp(prefix="impact_reg_") - datadir = os.path.join(tmp, "Data") - outdir = os.path.join(tmp, "Out") - os.makedirs(datadir, exist_ok=True) - os.makedirs(outdir, exist_ok=True) + tail = "" try: - sitk.WriteImage(sitk.Cast(fixed, sitk.sitkFloat32), - os.path.join(datadir, "Fixed_image.mha"), True) - sitk.WriteImage(sitk.Cast(moving, sitk.sitkFloat32), - os.path.join(datadir, "Moving_image.mha"), True) - param_text = _impact_param_map(os.path.join("/Data/Models", model_rel), - gpu, iters, samples, voxel_sizes, - rescale=schedule, smoothing=smoothe) - with open(os.path.join(datadir, "ParameterMap.txt"), "w") as f: - f.write(param_text) + fixed_path = os.path.join(tmp, "Fixed.mha") + moving_path = os.path.join(tmp, "Moving.mha") + outdir = os.path.join(tmp, "Out") + sitk.WriteImage(sitk.Cast(fixed, sitk.sitkFloat32), fixed_path, True) + sitk.WriteImage(sitk.Cast(moving, sitk.sitkFloat32), moving_path, True) - # debug aid: keep a persistent copy of the exact parameter map elastix - # is about to run, written BEFORE the docker call so an abort leaves it - # behind for inspection (per-pid name: parallel workers do not clobber - # each other's dumps) - dump_dir = d("data/m6reg/impact/paramdump") - os.makedirs(dump_dir, exist_ok=True) - family = os.path.basename(os.path.dirname(model_rel)) or "model" - param_dump = os.path.join(dump_dir, f"ParameterMap_{family}_{os.getpid()}.txt") - with open(param_dump, "w") as f: - f.write(param_text) - - cmd = ["docker", "run", "--rm"] + cmd = [IMPACT_CLI, "register", preset, + "-f", fixed_path, "-m", moving_path, + "-o", outdir, + "--fields-only", # we derive the warp ourselves + "--tmp-dir", os.path.join(tmp, "Work")] if gpu >= 0: - cmd += ["--gpus", "all"] - cmd += [ - "-v", f"{datadir}:/Data", - "-v", f"{outdir}:/Out", - "-v", f"{d(IMPACT_MODEL_CACHE)}:/Data/Models", - IMPACT_IMAGE, IMPACT_ELASTIX_BIN, - "-f", "/Data/Fixed_image.mha", "-m", "/Data/Moving_image.mha", - "-p", "/Data/ParameterMap.txt", - "-out", "/Out/", "-threads", str(threads), - ] + cmd += ["--gpu", str(gpu)] + # preset tuning (YAML values, forwarded by the CLI to the preset app): + # the M6 robustness overrides, always; the budget knobs, when set + # (no-op for presets whose per-resolution matrix wins) + ov = "[" + ",".join(f'"{o}"' for o in IMPACT_PARAM_OVERRIDES) + "]" + cmd += ["--set", f"parameter_overrides={ov}"] + if iters is not None: + cmd += ["--set", f"max_iterations={int(iters)}"] + if samples is not None: + cmd += ["--set", f"spatial_samples={int(samples)}"] + + env = os.environ.copy() + env["KONFAI_ELASTIX_DIR"] = root + t0 = time.time() - proc = subprocess.run(cmd, capture_output=True, text=True, timeout=7200) + proc = subprocess.run(cmd, capture_output=True, text=True, + timeout=7200, env=env) dt = time.time() - t0 + tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:]) if proc.returncode != 0: - tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:]) - raise RuntimeError(f"impact: docker elastix failed (rc={proc.returncode}):\n{tail}") + raise RuntimeError(f"impact: {IMPACT_CLI} failed " + f"(rc={proc.returncode}):\n{tail}") - itk_tf = os.path.join(outdir, "TransformParameters.0.itk.txt") - if not os.path.exists(itk_tf): - tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:]) - raise RuntimeError(f"impact: no TransformParameters.0.itk.txt produced:\n{tail}") - - # the optimized Euler (world frame) is the full combination transform - # (AutomaticTransformInitialization is folded in by elastix); elastix - # maps moving(T1c) -> fixed(CT), which is the contract direction; - # fix_orientation verifies it against the skull ring. - T_main = sitk.ReadTransform(itk_tf) - T, inverted = fix_orientation(ct, t1c, T_main) + case_dirs = sorted(p for p in os.listdir(outdir) + if os.path.isdir(os.path.join(outdir, p))) + if not case_dirs: + raise RuntimeError(f"impact: no case output under {outdir}:\n{tail}") + case_dir = os.path.join(outdir, case_dirs[0]) + entries = [p for p in os.listdir(case_dir) + if os.path.isfile(os.path.join(case_dir, p)) + and not p.startswith("Moved")] + if len(entries) != 1: + raise RuntimeError( + f"impact: expected exactly one transform output in {case_dir}, " + f"got {entries}:\n{tail}") + T_main, dvf_img = _impact_read_transform(os.path.join(case_dir, entries[0])) + T, inverted = _orient_impact(ct, t1c, T_main, dvf_img) finally: if not keep_scratch: shutil.rmtree(tmp, ignore_errors=True) - family = os.path.basename(os.path.dirname(model_rel)) \ - or os.path.splitext(os.path.basename(model_rel))[0] meta = { "pipeline": "impact", - "toolkit": f"Docker {IMPACT_IMAGE} (Elastix 5.2.0 + IMPACT)", - "metric": f"IMPACT semantic ({family} {os.path.basename(model_rel)}, L2, Jacobian)", - "transform": (f"Rigid (EulerTransform, ASGD, {n} scales, auto center init; " + "toolkit": (f"impact-reg-konfai (KonfAI; elastix-impact 5.2, " + f"preset {preset})"), + "metric": ("IMPACT semantic (MIND R1D2_3D L1 + TotalSegmentator M730 " + "deepest-layer Dice, per-level weights, Static) + " + "Mattes MI (w 0.4) + bending energy"), + "transform": (f"Deformable (preset-native B-spline DVF on the T1c grid; " + f"CenterOfGravity auto init; " f"device={'GPU' + str(gpu) if gpu >= 0 else 'CPU'})"), "convention": "T1c->CT (verified)", "auto_inverted": inverted, "fixed": _modality_name(fixed), "moving": _modality_name(moving), "seconds_reg": round(dt, 1), - "docker_image": IMPACT_IMAGE, "model": model_rel, "gpu": gpu, - "iters": iters, "samples": samples, "param_dump": param_dump, + "preset": preset, "gpu": gpu, "iters": iters, "samples": samples, + "parameter_overrides": list(IMPACT_PARAM_OVERRIDES), + "elastix_root": os.path.basename(root), } if keep_scratch: meta["scratch"] = tmp @@ -1064,7 +1075,7 @@ def _parse_niftyreg_affine(path): def register_niftyreg(ct, t1c, threads=24, levels=3, maxit=15, keep_scratch=False, backend="auto", gpu=None, clip_ct=(0.0, 80.0)): - """NiftyReg rigid CT->T1c (reference=CT, floating=T1c) via ``reg_aladin``. + """NiftyReg rigid CT->T1c (reference=T1c, floating=CT) via ``reg_aladin``. reg_aladin (Modat et al., J. Med. Img. 1(2) 024003, 2014) is a global block-matching registration: at each level of a ``levels``-level Gaussian @@ -1133,11 +1144,12 @@ def register_niftyreg(ct, t1c, threads=24, levels=3, maxit=15, keep_scratch=Fals gpu, use_gpu, platf = -1, False, None interp = "3" toolkit = _niftyreg_version(binpath) - # full volumes, no mask / crop / pad; reference (fixed) = CT, floating = T1c + # full volumes, no mask / crop / pad; reference (fixed) = T1c, floating = CT fixed, moving = fixed_moving(ct, t1c) - fixed_name = _modality_name(fixed) # before clipping (a clipped CT is no - # longer intensity-detectable) - fixed, ct_clipped = _maybe_clip_ct(fixed, clip_ct) + fixed_name = _modality_name(fixed) + moving_name = _modality_name(moving) # before clipping (a clipped CT is + # no longer intensity-detectable) + fixed, moving, ct_clipped = _maybe_clip_ct_pair(fixed, moving, clip_ct) tmp = tempfile.mkdtemp(prefix="niftyreg_reg_") dt = 0.0 try: @@ -1178,9 +1190,10 @@ def register_niftyreg(ct, t1c, threads=24, levels=3, maxit=15, keep_scratch=Fals tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:]) raise RuntimeError(f"niftyreg: no affine output produced:\n{tail}") T = _parse_niftyreg_affine(aff) - # aff.txt is documented as CT -> T1c (Affine*Reference=Floating); - # the contract is T1c -> CT. fix_orientation resolves the direction - # empirically (on these cases it keeps the file's affine as-is). + # aff.txt is documented as Affine * Reference = Floating (ref -> flo); + # with reference=T1c / floating=CT the direction the file actually + # holds is resolved empirically by fix_orientation against the skull + # ring (see _parse_niftyreg_affine). T, inverted = fix_orientation(ct, t1c, T) finally: if not keep_scratch: @@ -1194,7 +1207,7 @@ def register_niftyreg(ct, t1c, threads=24, levels=3, maxit=15, keep_scratch=Fals "transform": (f"Rigid (reg_aladin -rigOnly, {levels}-level pyramid, " f"{maxit} TLS iters/level, cubic interp)"), "convention": "T1c->CT (verified)", "auto_inverted": inverted, - "fixed": fixed_name, "moving": _modality_name(moving), + "fixed": fixed_name, "moving": moving_name, "seconds_reg": round(dt, 1), "levels": levels, "maxit": maxit, "interp": interp, "gpu": gpu, "clip_ct": [float(clip_ct[0]), float(clip_ct[1])] if ct_clipped else None, @@ -1259,15 +1272,21 @@ def availability(): else: out[name] = (True, f"Docker {SYNTHMORPH_IMAGE}") elif name == "impact": - ok, note = _impact_docker_ok() + ok, note = _impact_cli_ok() if not ok: out[name] = (False, note) continue - try: - host, rel = _impact_model_host(None) - out[name] = (True, f"Docker {IMPACT_IMAGE} + model {rel}") - except Exception as e: # noqa - out[name] = (False, str(e)) + root = _impact_elastix_root() + if root is None: + out[name] = (False, + note + "; no runnable elastix-impact binary " + "(needs glibc >= 2.35 + LibTorch 2.8.0) -- run " + "scripts/setup_impact_konfai.sh once") + else: + out[name] = (True, + note + "; elastix-impact under " + f"{os.path.basename(root)}; preset " + f"{IMPACT_DEFAULT_PRESET}") except Exception as e: # noqa out[name] = (False, repr(e)) return out \ No newline at end of file