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
|
||||
|
||||
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:
|
||||
|
||||
```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.
|
||||
|
|
@ -20,6 +20,6 @@ Longitudinal (repeated-measures) analysis of medical imaging data. Repo is at an
|
|||
|
||||
## 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).
|
||||
- No lint/test tooling is configured yet; run scripts directly with `python <script>` from the repo root.
|
||||
|
|
@ -33,9 +33,9 @@ the study.
|
|||
|
||||
## Prerequisites
|
||||
|
||||
- The `longitudinal` conda env and the GPU(s) it sees:
|
||||
- The `longitudinal-313` conda env and the GPU(s) it sees:
|
||||
```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/proc/*` + `data/manifests/split_{train,val,test}.jsonl`, i.e.
|
||||
|
|
@ -51,7 +51,7 @@ the study.
|
|||
|
||||
```bash
|
||||
cd /mnt/b4/xfr/git26/longitudinal
|
||||
conda activate longitudinal
|
||||
conda activate longitudinal-313
|
||||
|
||||
# small/fast check (what the validation run used): ~minutes per pipeline
|
||||
python scripts/benchmark_pipelines.py \
|
||||
|
|
|
|||
|
|
@ -13,11 +13,11 @@ native volumes straight into nnU-Net, which performs its own preprocessing
|
|||
|
||||
## Environment
|
||||
|
||||
Conda env `longitudinal` (Python 3.14, torch 2.14 +cu126). Activate before any
|
||||
command:
|
||||
Conda env `longitudinal-313` (Python 3.13, torch 2.14 +cu126). Activate before
|
||||
any command:
|
||||
|
||||
```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
|
||||
|
|
|
|||
|
|
@ -147,7 +147,7 @@ Per-case (mid / std, `OK` = accepted):
|
|||
```bash
|
||||
# main env — SimpleITK + Elastix (itk-elastix is pip-installed here;
|
||||
# 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 \
|
||||
--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:
|
||||
|
||||
```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
|
||||
/home/xfr/.conda/envs/ants312/bin/python scripts/bench_registration.py \
|
||||
--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):
|
||||
simpleitk | elastix | antspyx | deepreg
|
||||
simpleitk | elastix | antspyx | deepreg | impact
|
||||
|
||||
Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI
|
||||
(cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto
|
||||
|
|
@ -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).
|
||||
- reg_wall_s : wall-clock registration time (excludes I/O + scoring).
|
||||
|
||||
Because two of the four toolkits are not pip-installable into the main
|
||||
`longitudinal` env (antspyx has no cp314 wheel; PyPI deepreg is a broken
|
||||
stub), the benchmark is env-portable: run it under whichever conda env has a
|
||||
given toolkit installed, and it records only the pipelines that are
|
||||
importable there. elastix runs in the main env through the ITK-Elastix
|
||||
python API (pip ``itk`` + ``itk-elastix``). The `--merge` mode then combines
|
||||
Because some toolkits are not pip-installable into the main
|
||||
`longitudinal-313` env (antspyx is kept in an isolated env, see
|
||||
REGISTRATION_BENCHMARK.md; PyPI deepreg is a broken stub), the benchmark is
|
||||
env-portable: run it under whichever conda env has a given toolkit installed,
|
||||
and it records only the pipelines that are importable there. elastix runs in the main env through the ITK-Elastix
|
||||
python API (pip ``itk`` + ``itk-elastix``). ``impact`` runs in the
|
||||
``vboussot/elastix_impact`` Docker image (needs Docker + the image + a model,
|
||||
see REGISTRATION_BENCHMARK.md). The `--merge` mode then combines
|
||||
the per-pipeline JSON files (accumulated across envs) into one summary + plot.
|
||||
|
||||
Usage (repo root):
|
||||
# main env (SimpleITK + elastix; deepreg reports as unavailable)
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
|
||||
python scripts/bench_registration.py --pipelines simpleitk elastix deepreg
|
||||
# main env (SimpleITK + elastix; impact needs Docker + the image; deepreg unavailable)
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
||||
python scripts/bench_registration.py --pipelines simpleitk elastix impact
|
||||
|
||||
# isolated env that has antspyx (see REGISTRATION_BENCHMARK.md)
|
||||
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":
|
||||
return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir,
|
||||
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)
|
||||
|
||||
|
||||
|
|
@ -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("--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("--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",
|
||||
help="also write a checkreg-style PNG (checkerboard blend + color "
|
||||
"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
|
||||
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 --pipelines A C
|
||||
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``)
|
||||
--------------------------------------------------------
|
||||
|
|
@ -31,8 +31,13 @@ Pipelines
|
|||
- ``register_antspyx`` : ANTs / ANTsPy ``ants.registration`` (Rigid).
|
||||
- ``register_deepreg`` : Oxford DeepReg (CycleGAN deformation-field
|
||||
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
|
||||
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
|
||||
|
|
@ -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
|
||||
concentricity gap (a pure correctness gate -- a wrongly-oriented rigid transform
|
||||
lands the head on the far side of the box and yields a pathological gap). It is
|
||||
applied to elastix / antspyx / deepreg, not to the trusted production
|
||||
applied to elastix / antspyx / deepreg / impact, not to the trusted production
|
||||
SimpleITK path.
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import importlib.util
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import SimpleITK as sitk
|
||||
|
|
@ -580,11 +587,337 @@ def _displacement_to_sitk(disp, ref):
|
|||
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 = {
|
||||
"simpleitk": register_simpleitk,
|
||||
"elastix": register_elastix,
|
||||
"antspyx": register_antspyx,
|
||||
"deepreg": register_deepreg,
|
||||
"impact": register_impact,
|
||||
}
|
||||
ALL_PIPELINES = list(REGISTERS)
|
||||
|
||||
|
|
@ -616,6 +949,16 @@ def availability():
|
|||
out[name] = (True, "deepreg importable")
|
||||
except Exception as e: # noqa
|
||||
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
|
||||
out[name] = (False, repr(e))
|
||||
return out
|
||||
Loading…
Reference in a new issue