fix(docs): update conda environment name to longitudinal-313 and Python version to 3.13
feat(benchmark): add IMPACT pipeline for CT to T1c registration
This commit is contained in:
parent
a9a4d9bba0
commit
7ac9979c62
7 changed files with 386 additions and 26 deletions
|
|
@ -2,12 +2,12 @@
|
||||||
|
|
||||||
## Environment
|
## Environment
|
||||||
|
|
||||||
This project uses a conda environment named `longitudinal` (Python 3.14).
|
This project uses a conda environment named `longitudinal-313` (Python 3.13).
|
||||||
|
|
||||||
Activate with:
|
Activate with:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
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.
|
Key packages: torch 2.14 (+cu126), torchvision, monai 1.6 (pip, Pipeline C), numpy, scipy, pandas, scikit-learn, scikit-image, matplotlib, nibabel, SimpleITK.
|
||||||
|
|
@ -20,6 +20,6 @@ Longitudinal (repeated-measures) analysis of medical imaging data. Repo is at an
|
||||||
|
|
||||||
## Conventions
|
## Conventions
|
||||||
|
|
||||||
- All Python commands must run inside the `longitudinal` conda environment.
|
- All Python commands must run inside the `longitudinal-313` conda environment.
|
||||||
- GPU is available (CUDA 12.6); use `torch.device('cuda')` when appropriate. Multi-GPU stages of Pipelines A and C run via `torchrun --standalone --nproc_per_node N` (rank = GPU).
|
- GPU is available (CUDA 12.6); use `torch.device('cuda')` when appropriate. Multi-GPU stages of Pipelines A and C run via `torchrun --standalone --nproc_per_node N` (rank = GPU).
|
||||||
- No lint/test tooling is configured yet; run scripts directly with `python <script>` from the repo root.
|
- No lint/test tooling is configured yet; run scripts directly with `python <script>` from the repo root.
|
||||||
|
|
@ -33,9 +33,9 @@ the study.
|
||||||
|
|
||||||
## Prerequisites
|
## Prerequisites
|
||||||
|
|
||||||
- The `longitudinal` conda env and the GPU(s) it sees:
|
- The `longitudinal-313` conda env and the GPU(s) it sees:
|
||||||
```bash
|
```bash
|
||||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
||||||
```
|
```
|
||||||
- Data for Pipelines A and C (and split membership for all three):
|
- Data for Pipelines A and C (and split membership for all three):
|
||||||
`data/proc/*` + `data/manifests/split_{train,val,test}.jsonl`, i.e.
|
`data/proc/*` + `data/manifests/split_{train,val,test}.jsonl`, i.e.
|
||||||
|
|
@ -51,7 +51,7 @@ the study.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
cd /mnt/b4/xfr/git26/longitudinal
|
cd /mnt/b4/xfr/git26/longitudinal
|
||||||
conda activate longitudinal
|
conda activate longitudinal-313
|
||||||
|
|
||||||
# small/fast check (what the validation run used): ~minutes per pipeline
|
# small/fast check (what the validation run used): ~minutes per pipeline
|
||||||
python scripts/benchmark_pipelines.py \
|
python scripts/benchmark_pipelines.py \
|
||||||
|
|
|
||||||
|
|
@ -13,11 +13,11 @@ native volumes straight into nnU-Net, which performs its own preprocessing
|
||||||
|
|
||||||
## Environment
|
## Environment
|
||||||
|
|
||||||
Conda env `longitudinal` (Python 3.14, torch 2.14 +cu126). Activate before any
|
Conda env `longitudinal-313` (Python 3.13, torch 2.14 +cu126). Activate before
|
||||||
command:
|
any command:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
||||||
```
|
```
|
||||||
|
|
||||||
GPU (CUDA 12.6) is available; multi-GPU jobs use `torchrun` (in-house) or
|
GPU (CUDA 12.6) is available; multi-GPU jobs use `torchrun` (in-house) or
|
||||||
|
|
|
||||||
|
|
@ -147,7 +147,7 @@ Per-case (mid / std, `OK` = accepted):
|
||||||
```bash
|
```bash
|
||||||
# main env — SimpleITK + Elastix (itk-elastix is pip-installed here;
|
# main env — SimpleITK + Elastix (itk-elastix is pip-installed here;
|
||||||
# antspyx/deepreg report unavailable and are recorded as such)
|
# antspyx/deepreg report unavailable and are recorded as such)
|
||||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
||||||
python scripts/bench_registration.py --n-cases 10 \
|
python scripts/bench_registration.py --n-cases 10 \
|
||||||
--workers 4 --threads 24 --simpleitk-mode benchmark
|
--workers 4 --threads 24 --simpleitk-mode benchmark
|
||||||
|
|
||||||
|
|
@ -189,7 +189,7 @@ The title carries the pipeline, case key, Mattes MI, and the skull-ring gap
|
||||||
visuals agree. Generate for the available pipelines with:
|
visuals agree. Generate for the available pipelines with:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
/home/xfr/.conda/envs/longitudinal/bin/python scripts/bench_registration.py \
|
/home/xfr/.conda/envs/longitudinal-313/bin/python scripts/bench_registration.py \
|
||||||
--pipelines simpleitk elastix --n-cases 10 --qa-images
|
--pipelines simpleitk elastix --n-cases 10 --qa-images
|
||||||
/home/xfr/.conda/envs/ants312/bin/python scripts/bench_registration.py \
|
/home/xfr/.conda/envs/ants312/bin/python scripts/bench_registration.py \
|
||||||
--pipelines antspyx --n-cases 10 --qa-images
|
--pipelines antspyx --n-cases 10 --qa-images
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,7 @@
|
||||||
"""Benchmark the 4 rigid CT->T1c registration pipelines on N M6 test cases.
|
"""Benchmark the 5 rigid CT->T1c registration pipelines on N M6 test cases.
|
||||||
|
|
||||||
Pipelines (see src/registration_pipelines.py for the shared transform contract):
|
Pipelines (see src/registration_pipelines.py for the shared transform contract):
|
||||||
simpleitk | elastix | antspyx | deepreg
|
simpleitk | elastix | antspyx | deepreg | impact
|
||||||
|
|
||||||
Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI
|
Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI
|
||||||
(cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto
|
(cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto
|
||||||
|
|
@ -17,18 +17,20 @@ geometric QA used in production (scripts/02_build_m6_dataset.py):
|
||||||
- accepted : production rule (nvox>0 and tv_in_ct and concentric).
|
- accepted : production rule (nvox>0 and tv_in_ct and concentric).
|
||||||
- reg_wall_s : wall-clock registration time (excludes I/O + scoring).
|
- reg_wall_s : wall-clock registration time (excludes I/O + scoring).
|
||||||
|
|
||||||
Because two of the four toolkits are not pip-installable into the main
|
Because some toolkits are not pip-installable into the main
|
||||||
`longitudinal` env (antspyx has no cp314 wheel; PyPI deepreg is a broken
|
`longitudinal-313` env (antspyx is kept in an isolated env, see
|
||||||
stub), the benchmark is env-portable: run it under whichever conda env has a
|
REGISTRATION_BENCHMARK.md; PyPI deepreg is a broken stub), the benchmark is
|
||||||
given toolkit installed, and it records only the pipelines that are
|
env-portable: run it under whichever conda env has a given toolkit installed,
|
||||||
importable there. elastix runs in the main env through the ITK-Elastix
|
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``). The `--merge` mode then combines
|
python API (pip ``itk`` + ``itk-elastix``). ``impact`` runs in the
|
||||||
|
``vboussot/elastix_impact`` Docker image (needs Docker + the image + a model,
|
||||||
|
see REGISTRATION_BENCHMARK.md). The `--merge` mode then combines
|
||||||
the per-pipeline JSON files (accumulated across envs) into one summary + plot.
|
the per-pipeline JSON files (accumulated across envs) into one summary + plot.
|
||||||
|
|
||||||
Usage (repo root):
|
Usage (repo root):
|
||||||
# main env (SimpleITK + elastix; deepreg reports as unavailable)
|
# main env (SimpleITK + elastix; impact needs Docker + the image; deepreg unavailable)
|
||||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
||||||
python scripts/bench_registration.py --pipelines simpleitk elastix deepreg
|
python scripts/bench_registration.py --pipelines simpleitk elastix impact
|
||||||
|
|
||||||
# isolated env that has antspyx (see REGISTRATION_BENCHMARK.md)
|
# isolated env that has antspyx (see REGISTRATION_BENCHMARK.md)
|
||||||
source /opt/conda/etc/profile.d/conda.sh && conda activate <ants312>
|
source /opt/conda/etc/profile.d/conda.sh && conda activate <ants312>
|
||||||
|
|
@ -224,6 +226,9 @@ def _call(pipeline, ct, t1c, tv, args):
|
||||||
if pipeline == "deepreg":
|
if pipeline == "deepreg":
|
||||||
return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir,
|
return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir,
|
||||||
checkpoint=args.deepreg_ckpt)
|
checkpoint=args.deepreg_ckpt)
|
||||||
|
if pipeline == "impact":
|
||||||
|
return RP.register_impact(ct, t1c, model=args.impact_model, gpu=args.impact_gpu,
|
||||||
|
iters=args.impact_iters, samples=args.impact_samples)
|
||||||
raise KeyError(pipeline)
|
raise KeyError(pipeline)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -471,6 +476,18 @@ def parse_args():
|
||||||
ap.add_argument("--elastix-dir", default=None, help="elastix scratch dir (default data/m6reg/elastix)")
|
ap.add_argument("--elastix-dir", default=None, help="elastix scratch dir (default data/m6reg/elastix)")
|
||||||
ap.add_argument("--deepreg-dir", default=None, help="deepreg scratch dir")
|
ap.add_argument("--deepreg-dir", default=None, help="deepreg scratch dir")
|
||||||
ap.add_argument("--deepreg-ckpt", default=None, help="path to a DeepReg 3D pretrained checkpoint")
|
ap.add_argument("--deepreg-ckpt", default=None, help="path to a DeepReg 3D pretrained checkpoint")
|
||||||
|
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-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=400,
|
||||||
|
help="IMPACT per-resolution max iterations (default 400, a converged "
|
||||||
|
"rigid budget; the authors' full recommended is 500)")
|
||||||
|
ap.add_argument("--impact-samples", type=int, default=1500,
|
||||||
|
help="IMPACT random spatial samples per iteration (default 1500; "
|
||||||
|
"full recommended 2000)")
|
||||||
ap.add_argument("--qa-images", action="store_true",
|
ap.add_argument("--qa-images", action="store_true",
|
||||||
help="also write a checkreg-style PNG (checkerboard blend + color "
|
help="also write a checkreg-style PNG (checkerboard blend + color "
|
||||||
"overlay of the inverse-warped CT vs the original MR) per case")
|
"overlay of the inverse-warped CT vs the original MR) per case")
|
||||||
|
|
|
||||||
|
|
@ -26,7 +26,7 @@ Protocol
|
||||||
- VRAM training: max allocation over the training run (A/C in-process via
|
- VRAM training: max allocation over the training run (A/C in-process via
|
||||||
torch; B via nvidia-smi polling around the nnUNetv2_train subprocess).
|
torch; B via nvidia-smi polling around the nnUNetv2_train subprocess).
|
||||||
|
|
||||||
Usage (repo root, inside the `longitudinal` conda env):
|
Usage (repo root, inside the `longitudinal-313` conda env):
|
||||||
python scripts/benchmark_pipelines.py # train + eval A B C
|
python scripts/benchmark_pipelines.py # train + eval A B C
|
||||||
python scripts/benchmark_pipelines.py --pipelines A C
|
python scripts/benchmark_pipelines.py --pipelines A C
|
||||||
python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240
|
python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,4 @@
|
||||||
"""Four rigid CT -> T1c registration pipelines behind one unified contract.
|
"""Five rigid CT -> T1c registration pipelines behind one unified contract.
|
||||||
|
|
||||||
Problem (shared with ``scripts/02_build_m6_dataset.py``)
|
Problem (shared with ``scripts/02_build_m6_dataset.py``)
|
||||||
--------------------------------------------------------
|
--------------------------------------------------------
|
||||||
|
|
@ -31,8 +31,13 @@ Pipelines
|
||||||
- ``register_antspyx`` : ANTs / ANTsPy ``ants.registration`` (Rigid).
|
- ``register_antspyx`` : ANTs / ANTsPy ``ants.registration`` (Rigid).
|
||||||
- ``register_deepreg`` : Oxford DeepReg (CycleGAN deformation-field
|
- ``register_deepreg`` : Oxford DeepReg (CycleGAN deformation-field
|
||||||
registration; needs a pretrained checkpoint, see module docstring).
|
registration; needs a pretrained checkpoint, see module docstring).
|
||||||
|
- ``register_impact`` : IMPACT (semantic-similarity metric) inside the
|
||||||
|
``vboussot/elastix_impact`` Docker image (Elastix 5.2 + a TorchScript
|
||||||
|
TotalSegmentator feature extractor as a native Elastix metric). Rigid
|
||||||
|
6-DOF, GPU (Jacobian mode). Needs Docker + the image + a model asset.
|
||||||
|
|
||||||
A toolkit that is not installed (or, for DeepReg, missing its checkpoint)
|
A toolkit that is not installed (or, for DeepReg, missing its checkpoint, or,
|
||||||
|
for IMPACT, missing its Docker image / model)
|
||||||
raises :class:`PipelineUnavailable` with the exact install / precondition. The
|
raises :class:`PipelineUnavailable` with the exact install / precondition. The
|
||||||
benchmark catches that and reports the pipeline as *not available* instead of
|
benchmark catches that and reports the pipeline as *not available* instead of
|
||||||
failing the whole run, so the script is env-portable: run it under whichever
|
failing the whole run, so the script is env-portable: run it under whichever
|
||||||
|
|
@ -45,13 +50,15 @@ 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
|
``T`` and ``T^-1`` and keeping whichever gives the smaller skull-ring
|
||||||
concentricity gap (a pure correctness gate -- a wrongly-oriented rigid transform
|
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
|
lands the head on the far side of the box and yields a pathological gap). It is
|
||||||
applied to elastix / antspyx / deepreg, not to the trusted production
|
applied to elastix / antspyx / deepreg / impact, not to the trusted production
|
||||||
SimpleITK path.
|
SimpleITK path.
|
||||||
"""
|
"""
|
||||||
import os
|
import os
|
||||||
import json
|
import json
|
||||||
|
import shutil
|
||||||
import subprocess
|
import subprocess
|
||||||
import importlib.util
|
import importlib.util
|
||||||
|
import tempfile
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import SimpleITK as sitk
|
import SimpleITK as sitk
|
||||||
|
|
@ -580,11 +587,337 @@ def _displacement_to_sitk(disp, ref):
|
||||||
return t
|
return t
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# 5. IMPACT (Docker image vboussot/elastix_impact)
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
IMPACT_IMAGE = "vboussot/elastix_impact"
|
||||||
|
IMPACT_ELASTIX_BIN = "/lib/elastix-install/bin/elastix"
|
||||||
|
IMPACT_MODEL_CACHE = "data/impact_models"
|
||||||
|
IMPACT_DEFAULT_MODEL = "TS/M730.pt"
|
||||||
|
IMPACT_HF_MODEL_URL = (
|
||||||
|
"https://huggingface.co/VBoussot/impact-torchscript-models/resolve/main/{rel}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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"
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
_NGPUS_CACHE = None
|
||||||
|
|
||||||
|
|
||||||
|
def _impact_n_gpus():
|
||||||
|
global _NGPUS_CACHE
|
||||||
|
if _NGPUS_CACHE is None:
|
||||||
|
n = 0
|
||||||
|
try:
|
||||||
|
r = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True, timeout=30)
|
||||||
|
n = len([l for l in r.stdout.splitlines() if l.startswith("GPU ")])
|
||||||
|
except Exception: # noqa
|
||||||
|
n = 0
|
||||||
|
_NGPUS_CACHE = n
|
||||||
|
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_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")',
|
||||||
|
"",
|
||||||
|
"// rigid 6-DOF to be optimized (the shared centered start comes via -t0)",
|
||||||
|
'(Transform "EulerTransform")',
|
||||||
|
'(TransformBase "ElastixTransformBase")',
|
||||||
|
'(TransformParameterReadoutFormat "AllParametersOneLine")',
|
||||||
|
'(FixedCenteredTransform "true")',
|
||||||
|
'(HowToCombineTransforms "Compose")',
|
||||||
|
"",
|
||||||
|
"(ImageDimension 3)",
|
||||||
|
'(FixedInternalImagePixelType "float")',
|
||||||
|
'(MovingInternalImagePixelType "float")',
|
||||||
|
'(UseDirectionCosines "true")',
|
||||||
|
"",
|
||||||
|
'(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 register_impact(ct, t1c, model=None, gpu=None, iters=400, samples=1500,
|
||||||
|
threads=12, keep_scratch=False, scales=(6.0, 3.0, 1.5, 1.0),
|
||||||
|
downsample=False):
|
||||||
|
"""IMPACT rigid CT->T1c (fixed=T1c, moving=CT) in the Docker image
|
||||||
|
``vboussot/elastix_impact`` (Elastix 5.2.0 + IMPACT semantic metric).
|
||||||
|
|
||||||
|
IMPACT (arXiv:2503.24121) replaces the intensity metric with a *semantic*
|
||||||
|
one: a TorchScript TotalSegmentator model (``TS/M730.pt`` by default) 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). This is a genuinely different metric from
|
||||||
|
the Mattes-MI pipelines above -- it is the cross-modality tool the image was
|
||||||
|
built for.
|
||||||
|
|
||||||
|
To stay on the benchmark's equal footing the pipeline:
|
||||||
|
- registers on the same CT-space head box (both images resampled into it),
|
||||||
|
- starts from the *shared* centroid-aligned translation (passed to Elastix
|
||||||
|
as a fixed initial transform via ``-t0``; with ``HowToCombineTransforms
|
||||||
|
Compose`` the final transform = ``T_main(T_init(p))``),
|
||||||
|
- uses a rigid 6-DOF ``EulerTransform`` (the 5.x name for the classic
|
||||||
|
3D rigid), ASGD, a 4-level scale pyramid (no image resampling).
|
||||||
|
Budget: the default ``iters=400`` / ``samples=1500`` is a *converged rigid*
|
||||||
|
budget (measured: recovers the inter-session tilt the intensity-metric
|
||||||
|
pipelines miss). The authors' fully-recommended Jacobian budget is
|
||||||
|
``iters=500`` / ``samples=2000`` (deformable-oriented, ~25% slower); pass
|
||||||
|
those for maximum fidelity. Expect minutes per case (GPU), not seconds.
|
||||||
|
The result is read back from ``TransformParameters.0.itk.txt``
|
||||||
|
(an ``Euler3DTransform`` in the box/world physical frame) and composed with
|
||||||
|
the ``-t0`` translation, then checked by :func:`fix_orientation`.
|
||||||
|
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
import time
|
||||||
|
|
||||||
|
ok, note = _impact_docker_ok()
|
||||||
|
if not ok:
|
||||||
|
raise PipelineUnavailable(f"impact: {note}")
|
||||||
|
model_host, model_rel = _impact_model_host(model)
|
||||||
|
|
||||||
|
lo, hi = _ct_head_box_mm(ct)
|
||||||
|
if lo is None:
|
||||||
|
raise PipelineUnavailable("impact: no CT head box found (corrupt CT?)")
|
||||||
|
|
||||||
|
# resample both into the shared CT-space head box (mirrors the other pipelines)
|
||||||
|
tsp = np.array(t1c.GetSpacing())
|
||||||
|
shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3))
|
||||||
|
ref = sitk.GetImageFromArray(np.zeros(shape, np.float32))
|
||||||
|
ref.SetSpacing(t1c.GetSpacing())
|
||||||
|
ref.SetDirection(t1c.GetDirection())
|
||||||
|
ref.SetOrigin(tuple(float(x) for x in lo))
|
||||||
|
fixed = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # T1c = reference
|
||||||
|
moving = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # CT = moving
|
||||||
|
|
||||||
|
# shared centered start: elastix maps fixed(T1c)->moving(CT), i.e. the
|
||||||
|
# inverse of centroid_offset, hence the negation.
|
||||||
|
off = centroid_offset(ct, t1c, ct_box=moving) or (0.0, 0.0, 0.0)
|
||||||
|
init = (-float(off[0]), -float(off[1]), -float(off[2]))
|
||||||
|
|
||||||
|
# GPU: auto-spread across the host GPUs by pid, or forced value
|
||||||
|
if gpu is None:
|
||||||
|
ngpu = _impact_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 box's (T1c) base spacing
|
||||||
|
base = float(np.median(t1c.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)
|
||||||
|
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)
|
||||||
|
with open(os.path.join(datadir, "ParameterMap.txt"), "w") as f:
|
||||||
|
f.write(_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, "InitialTransform.txt"), "w") as f:
|
||||||
|
f.write(
|
||||||
|
'(Transform "TranslationTransform")\n'
|
||||||
|
'(TransformBase "ElastixTransformBase")\n'
|
||||||
|
'(TransformParameterReadoutFormat "AllParametersOneLine")\n'
|
||||||
|
"(NumberOfParameters 3)\n"
|
||||||
|
f"(TransformParameters {init[0]:g} {init[1]:g} {init[2]:g})\n")
|
||||||
|
|
||||||
|
cmd = ["docker", "run", "--rm"]
|
||||||
|
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",
|
||||||
|
"-t0", "/Data/InitialTransform.txt",
|
||||||
|
"-out", "/Out/", "-threads", str(threads),
|
||||||
|
]
|
||||||
|
t0 = time.time()
|
||||||
|
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=7200)
|
||||||
|
dt = time.time() - t0
|
||||||
|
if proc.returncode != 0:
|
||||||
|
tail = "\n".join((proc.stdout + proc.stderr).splitlines()[-25:])
|
||||||
|
raise RuntimeError(f"impact: docker elastix failed (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 (box/world frame) composed with the fixed -t0
|
||||||
|
# translation: T = T_euler(T_init(p))
|
||||||
|
T_main = sitk.ReadTransform(itk_tf)
|
||||||
|
T_init = sitk.TranslationTransform(3)
|
||||||
|
T_init.SetOffset(init)
|
||||||
|
cand = sitk.CompositeTransform([T_init, T_main])
|
||||||
|
T, inverted = fix_orientation(ct, t1c, cand)
|
||||||
|
finally:
|
||||||
|
if not keep_scratch:
|
||||||
|
shutil.rmtree(tmp, ignore_errors=True)
|
||||||
|
|
||||||
|
meta = {
|
||||||
|
"pipeline": "impact",
|
||||||
|
"toolkit": f"Docker {IMPACT_IMAGE} (Elastix 5.2.0 + IMPACT)",
|
||||||
|
"metric": f"IMPACT semantic (TotalSegmentator {os.path.basename(model_rel)}, L2, Jacobian)",
|
||||||
|
"transform": (f"Rigid (EulerTransform, ASGD, {n} scales, centroid init; "
|
||||||
|
f"device={'GPU' + str(gpu) if gpu >= 0 else 'CPU'})"),
|
||||||
|
"convention": "T1c->CT (verified)", "auto_inverted": inverted,
|
||||||
|
"fixed": "T1c", "moving": "CT",
|
||||||
|
"seconds_reg": round(dt, 1),
|
||||||
|
"docker_image": IMPACT_IMAGE, "model": model_rel, "gpu": gpu,
|
||||||
|
"iters": iters, "samples": samples,
|
||||||
|
}
|
||||||
|
if keep_scratch:
|
||||||
|
meta["scratch"] = tmp
|
||||||
|
return T, meta
|
||||||
|
|
||||||
|
|
||||||
REGISTERS = {
|
REGISTERS = {
|
||||||
"simpleitk": register_simpleitk,
|
"simpleitk": register_simpleitk,
|
||||||
"elastix": register_elastix,
|
"elastix": register_elastix,
|
||||||
"antspyx": register_antspyx,
|
"antspyx": register_antspyx,
|
||||||
"deepreg": register_deepreg,
|
"deepreg": register_deepreg,
|
||||||
|
"impact": register_impact,
|
||||||
}
|
}
|
||||||
ALL_PIPELINES = list(REGISTERS)
|
ALL_PIPELINES = list(REGISTERS)
|
||||||
|
|
||||||
|
|
@ -616,6 +949,16 @@ def availability():
|
||||||
out[name] = (True, "deepreg importable")
|
out[name] = (True, "deepreg importable")
|
||||||
except Exception as e: # noqa
|
except Exception as e: # noqa
|
||||||
out[name] = (False, "deepreg not importable")
|
out[name] = (False, "deepreg not importable")
|
||||||
|
elif name == "impact":
|
||||||
|
ok, note = _impact_docker_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))
|
||||||
except Exception as e: # noqa
|
except Exception as e: # noqa
|
||||||
out[name] = (False, repr(e))
|
out[name] = (False, repr(e))
|
||||||
return out
|
return out
|
||||||
Loading…
Reference in a new issue