Restructure IMPACT-Reg into per-preset KonfAI apps

Each preset is now a self-contained app (app.json + Prediction.yml + Model.py +
install.py + model.pt + parameter maps + evaluation configs). The registration
model takes fixed/moving plus optional fixed/moving masks (whole-image mask when
omitted); evaluation is declared per modality (image / segmentation / landmarks).
Removes the legacy monolithic app and the flat parameter-map layout.

CBCT_CT_HeadNeck/Model.py

@@ -0,0 +1,291 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+
+``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
+(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
+onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
+(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
+``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
+needs to register in physical space.
+
+Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+
+    Patch.patch_size: None        # one patch per case (registration is global)
+    batch_size: 1                 # one fixed/moving pair at a time
+
+NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
+    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._models = models
+        self._iterations = iterations
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Copy parameter maps into the work dir, patching the ImpactGPU device field when present."""
+        staged = []
+        for src in self._parameter_maps:
+            dst = work / src.name
+            lines = []
+            for line in src.read_text(encoding="utf-8").splitlines():
+                if line.strip().startswith("(ImpactGPU"):
+                    line = f"(ImpactGPU {device_index})"
+                lines.append(line)
+            dst.write_text("\n".join(lines) + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
+        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the parameter maps reference
+            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
+            # the "N% done/total" progress line) shows real progress during the long registration.
+            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            # ``iterations`` is the total iteration budget declared for the preset (summed over the
+            # chained parameter maps), so the bar spans the whole chain of registration stages.
+            progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            for line in proc.stdout:
+                captured.append(line)
+                if iteration_line.match(line):
+                    progress.update(1)
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
+    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(parameter_maps, models, iterations)
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
+        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
+        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)
+
+
+class ChannelSelect(torch.nn.Module):
+    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
+    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
+    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
+    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
+    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        engine: str = "elastix",
+        parameter_maps: list[str] = [],
+        models: list[str] = [],
+        iterations: int = 1000,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        self.add_module(
+            "Registration",
+            ElastixRegistration(engine, parameter_maps, models, iterations),
+            in_branch=[0, 1, 2, 3],
+            out_branch=["registration"],
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

CBCT_CT_HeadNeck/Prediction.yml

@@ -0,0 +1,108 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      engine: elastix
+      parameter_maps:
+      - ParameterMap_CBCT_HN.txt
+      models:
+      - VBoussot/impact-torchscript-models:TS/M732.pt
+      - VBoussot/impact-torchscript-models:TS/M731.pt
+      - VBoussot/impact-torchscript-models:TS/M730.pt
+      iterations: 1150
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-CBCT-CT-HeadNeck
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

CBCT_CT_HeadNeck/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "CBCT/CT Head&Neck",
+    "short_description": "Optimized preset for CBCT/CT registration on head & neck",
+    "description": "A five-level recursive B-spline deformable registration optimized for CBCT/CT head-and-neck alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 200, 200, and 150 iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm) using L1 distances on selected internal layers of the network. A composite objective (IMPACT + mutual information + bending energy penalty, with increased MI weight) ensures robust cross-modality alignment in complex head-and-neck anatomy while enforcing smooth, physically plausible deformations.",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

CBCT_CT_HeadNeck/install.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import platform
+import re
+import shutil
+import stat
+import subprocess  # nosec B404
+import zipfile
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+# -----------------------------------------------------------------------------
+# Elastix + IMPACT binary assets hosted on GitHub Releases.
+#
+# Key format: (OS, ARCH, FLAVOR)
+#   - OS     : platform.system() -> "Linux", "Windows", "Darwin"
+#   - ARCH   : normalized architecture -> "x86_64"
+#   - FLAVOR : "cpu" or "cu128"
+#
+# CPU assets are standalone (LibTorch CPU bundled).
+# CUDA assets do NOT bundle LibTorch (too large for GitHub limits).
+# -----------------------------------------------------------------------------
+ELX_ASSET_TEMPLATE = {
+    ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip",
+    ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip",
+    ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip",
+    ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip",
+    ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Official LibTorch downloads (PyTorch).
+#
+# IMPORTANT:
+# - Version MUST match the one used at build time (ABI compatibility).
+# - Using "shared-with-deps" ensures CUDA runtime libraries are included
+#   (except the NVIDIA driver, which must be installed system-wide).
+# -----------------------------------------------------------------------------
+
+LIBTORCH_URL = {
+    (
+        "Linux",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Darwin",
+        "arm64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Linux",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Minimum NVIDIA driver versions required for CUDA 12.8.
+#
+# The CUDA Toolkit itself is NOT required.
+# Only a sufficiently recent NVIDIA driver must be installed.
+# -----------------------------------------------------------------------------
+CUDA128_MIN_DRIVER_LINUX = (570, 26)
+CUDA128_MIN_DRIVER_WINDOWS = (570, 65)
+
+GITHUB_OWNER = "vboussot"
+GITHUB_REPO = "ImpactElastix"
+GITHUB_TAG = "1.0.0"
+
+
+DEFAULT_PREFIX = Path.cwd() / "elastix-impact"
+
+
+def run_cmd(cmd: list[str]) -> str:
+    return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace")  # nosec B603
+
+
+def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]:
+    """
+    Detect presence of an NVIDIA GPU and extract the driver version.
+    Returns:
+        (True, (major, minor)) if detected
+        (False, None) if nvidia-smi is not available
+    """
+    try:
+        nvidia_smi = shutil.which("nvidia-smi")
+        if nvidia_smi is None:
+            raise RuntimeError("nvidia-smi not found in PATH")
+
+        out = (
+            subprocess.check_output(
+                [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL
+            )  # nosec B603
+            .decode("utf-8")
+            .strip()
+        )
+    except Exception:
+        return (False, None)
+
+    # Exemple: "575.64"
+    m = re.match(r"([0-9]+)\.([0-9]+)", out)
+    if not m:
+        return (True, None)
+
+    return (True, (int(m.group(1)), int(m.group(2))))
+
+
+def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool:
+    """
+    Check whether the detected NVIDIA driver satisfies the minimum
+    requirement for CUDA 12.8 on the given operating system.
+    """
+    if drv is None:
+        return False
+    if os_name == "Linux":
+        return drv >= CUDA128_MIN_DRIVER_LINUX
+    if os_name == "Windows":
+        return drv >= CUDA128_MIN_DRIVER_WINDOWS
+    return False
+
+
+def normalize_arch(machine: str) -> str:
+    """
+    Normalize platform.machine() output across operating systems.
+
+    Examples:
+        AMD64   -> x86_64
+        aarch64 -> arm64
+    """
+    m = machine.lower()
+    if m in ("x86_64", "amd64"):
+        return "x86_64"
+    if m in ("aarch64", "arm64"):
+        return "arm64"
+    return machine
+
+
+def download_file(url: str, dst: Path) -> None:
+    """
+    Download a file from the given URL with a progress indicator.
+    """
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading: {url}", flush=True)
+    print(f"       to: {dst}", flush=True)
+
+    try:
+        with requests.get(url, stream=True, timeout=60) as r:
+            r.raise_for_status()
+            total = int(r.headers.get("content-length", 0))
+            with open(dst, "wb") as f:
+                with tqdm(
+                    total=total,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Downloading {dst.name}",
+                ) as pbar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+    except Exception as e:
+        raise e
+
+
+def extract_archive(archive: Path, dst_dir: Path) -> None:
+    """
+    Extract a ZIP archive to the destination directory.
+    """
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Extracting: {archive} -> {dst_dir}", flush=True)
+    with zipfile.ZipFile(archive, "r") as z:
+        z.extractall(dst_dir)
+    archive.unlink()
+
+
+def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool):
+    os_name = platform.system()
+    arch = normalize_arch(platform.machine())
+    has_nvidia, drv = detect_nvidia_driver()
+
+    if os_name not in ("Linux", "Windows", "Darwin"):
+        raise NameError(f"Unsupported OS: {os_name}")
+
+    if arch not in ("x86_64", "arm64"):
+        raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)")
+
+    flavor = "cpu"
+    if force_cuda:
+        if not has_nvidia or not driver_ok_for_cuda(os_name, drv):
+            raise NameError(
+                "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}"
+            )
+
+        flavor = "cu128"
+    elif not force_cpu:
+        if has_nvidia and driver_ok_for_cuda(os_name, drv):
+            flavor = "cu128"
+        else:
+            flavor = "cpu"
+
+    print(f"System: {os_name} {arch}", flush=True)
+    print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True)
+    print(f"Selected flavor: {flavor}", flush=True)
+
+    key = (os_name, arch, flavor)
+    if key not in ELX_ASSET_TEMPLATE:
+        raise NameError(f"No elastix asset configured for {key}")
+    if flavor != "cpu" and key not in LIBTORCH_URL:
+        raise NameError(f"No libtorch url configured for {key}")
+
+    install_path = install_path.resolve()
+    install_path.mkdir(parents=True, exist_ok=True)
+
+    elx_asset = ELX_ASSET_TEMPLATE[key]
+    elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}"
+    elx_archive = install_path / elx_asset
+    download_file(elx_url, elx_archive)
+    extract_archive(elx_archive, install_path)
+
+    # -------------------------------------------------------------------------
+    # ZIP archives may drop executable permissions.
+    # Ensure elastix and transformix are executable on Unix platforms.
+    # -------------------------------------------------------------------------
+    if os_name in ("Linux", "Darwin"):
+        for exe in ("elastix", "transformix"):
+            p = install_path / "bin" / exe
+            if p.exists():
+                p.chmod(p.stat().st_mode | stat.S_IEXEC)
+
+    # ---------------------------------------------------------------------
+    # - Download matching LibTorch 12.8 with or without CUDA 12.8
+    # - Extract it
+    # - Move runtime libraries to a location visible to the dynamic loader
+    #
+    # Linux / macOS : prefix/lib
+    # Windows       : prefix (next to elastix.exe)
+    # ---------------------------------------------------------------------
+    lt_url = LIBTORCH_URL[key]
+    lt_archive = install_path / Path(lt_url).name
+    download_file(lt_url, lt_archive)
+    extract_archive(lt_archive, install_path)
+    for p in (install_path / "libtorch" / "lib").iterdir():
+        if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name:
+            shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name)
+
+    shutil.rmtree(install_path / "libtorch")
+    if os_name == "Linux" and flavor == "cu128":
+        shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12")
+
+
+def get_elastix_bin(install_path: Path) -> Path:
+    return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix"))
+
+
+def try_elastix(install_path: Path) -> None:
+    try:
+        subprocess.run(
+            [str(get_elastix_bin(install_path)), "-h"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )  # nosec B603
+    except subprocess.CalledProcessError as e:
+        msg = "Elastix execution failed.\n\n"
+
+        msg += f"Command:\n{' '.join(e.cmd)}\n"
+        msg += f"Return code: {e.returncode}\n\n"
+
+        if e.stderr:
+            msg += "Error output:\n"
+            msg += e.stderr.strip()
+        raise NameError(msg)
+
+    except OSError as e:
+        msg = (
+            "Elastix could not be started.\n\n"
+            "This is usually caused by missing shared libraries "
+            "(e.g. LibTorch or CUDA runtime).\n\n"
+            f"System error:\n{str(e)}"
+        )
+
+        raise NameError(msg)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)"
+    )
+    group = ap.add_mutually_exclusive_group()
+    group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present")
+    group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)")
+    args = vars(ap.parse_args())
+    install_elastix_impact(**args)
+
+
+if __name__ == "__main__":
+    main()

ImpactReg/Uncertainty.yml → CBCT_CT_MRSeg/Evaluation_with_fid.yml

@@ -1,22 +1,20 @@
 Evaluator:
   metrics:
-    Uncertainty:
+    FixedFid:
       targets_criterions:
-        None:
+        MovingFid:
           criterions_loader:
-            Mean:
-              name: Uncertainty
+            TRE: {}
   Dataset:
     groups_src:
       Volume_0:
         groups_dest:
-          Uncertainty:
-            transforms: 
-              UncertaintyDVF:DVFToNorm: {}
-              Variance: {}
-              Save:
-                dataset: ./Uncertainties/ImpactReg/Output:mha
-                group: None
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
     subset: None
     dataset_filenames:
     - ./Dataset:mha

CBCT_CT_MRSeg/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

CBCT_CT_MRSeg/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

CBCT_CT_MRSeg/Model.py

@@ -0,0 +1,291 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+
+``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
+(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
+onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
+(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
+``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
+needs to register in physical space.
+
+Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+
+    Patch.patch_size: None        # one patch per case (registration is global)
+    batch_size: 1                 # one fixed/moving pair at a time
+
+NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
+    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._models = models
+        self._iterations = iterations
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Copy parameter maps into the work dir, patching the ImpactGPU device field when present."""
+        staged = []
+        for src in self._parameter_maps:
+            dst = work / src.name
+            lines = []
+            for line in src.read_text(encoding="utf-8").splitlines():
+                if line.strip().startswith("(ImpactGPU"):
+                    line = f"(ImpactGPU {device_index})"
+                lines.append(line)
+            dst.write_text("\n".join(lines) + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
+        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the parameter maps reference
+            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
+            # the "N% done/total" progress line) shows real progress during the long registration.
+            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            # ``iterations`` is the total iteration budget declared for the preset (summed over the
+            # chained parameter maps), so the bar spans the whole chain of registration stages.
+            progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            for line in proc.stdout:
+                captured.append(line)
+                if iteration_line.match(line):
+                    progress.update(1)
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
+    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(parameter_maps, models, iterations)
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
+        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
+        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)
+
+
+class ChannelSelect(torch.nn.Module):
+    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
+    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
+    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
+    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
+    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        engine: str = "elastix",
+        parameter_maps: list[str] = [],
+        models: list[str] = [],
+        iterations: int = 1000,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        self.add_module(
+            "Registration",
+            ElastixRegistration(engine, parameter_maps, models, iterations),
+            in_branch=[0, 1, 2, 3],
+            out_branch=["registration"],
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

CBCT_CT_MRSeg/Prediction.yml

@@ -0,0 +1,107 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      engine: elastix
+      parameter_maps:
+      - ParameterMap_CBCT_generic_MRSeg.txt
+      models:
+      - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt
+      - VBoussot/impact-torchscript-models:MRSeg/MRSeg.pt
+      iterations: 1100
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-CBCT-CT-MRSeg
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

CBCT_CT_MRSeg/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "CBCT/CT preset with MRSegmentator",
+    "short_description": "Generic CBCT/CT deformable registration using MRSegmentator features",
+    "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from the pretrained MRSegmentator model. The scheme follows a multi-resolution strategy with up to 300, 300, 250, and 200 ASGD iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), with a level-dependent combination of Dice-based segmentation overlap and L1 feature distances on selected internal layers of MRSegmentator. Early levels rely on pure Dice supervision, while finer stages progressively integrate feature-level alignment with increasing L1 contribution (0.3/0.7, 0.5/0.5) and a final purely feature-based stage. The optimization minimizes a composite objective (IMPACT + mutual information + bending energy penalty), enabling robust cross-modality alignment between CBCT and CT while enforcing smooth, physically plausible deformations.",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

CBCT_CT_MRSeg/install.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import platform
+import re
+import shutil
+import stat
+import subprocess  # nosec B404
+import zipfile
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+# -----------------------------------------------------------------------------
+# Elastix + IMPACT binary assets hosted on GitHub Releases.
+#
+# Key format: (OS, ARCH, FLAVOR)
+#   - OS     : platform.system() -> "Linux", "Windows", "Darwin"
+#   - ARCH   : normalized architecture -> "x86_64"
+#   - FLAVOR : "cpu" or "cu128"
+#
+# CPU assets are standalone (LibTorch CPU bundled).
+# CUDA assets do NOT bundle LibTorch (too large for GitHub limits).
+# -----------------------------------------------------------------------------
+ELX_ASSET_TEMPLATE = {
+    ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip",
+    ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip",
+    ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip",
+    ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip",
+    ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Official LibTorch downloads (PyTorch).
+#
+# IMPORTANT:
+# - Version MUST match the one used at build time (ABI compatibility).
+# - Using "shared-with-deps" ensures CUDA runtime libraries are included
+#   (except the NVIDIA driver, which must be installed system-wide).
+# -----------------------------------------------------------------------------
+
+LIBTORCH_URL = {
+    (
+        "Linux",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Darwin",
+        "arm64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Linux",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Minimum NVIDIA driver versions required for CUDA 12.8.
+#
+# The CUDA Toolkit itself is NOT required.
+# Only a sufficiently recent NVIDIA driver must be installed.
+# -----------------------------------------------------------------------------
+CUDA128_MIN_DRIVER_LINUX = (570, 26)
+CUDA128_MIN_DRIVER_WINDOWS = (570, 65)
+
+GITHUB_OWNER = "vboussot"
+GITHUB_REPO = "ImpactElastix"
+GITHUB_TAG = "1.0.0"
+
+
+DEFAULT_PREFIX = Path.cwd() / "elastix-impact"
+
+
+def run_cmd(cmd: list[str]) -> str:
+    return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace")  # nosec B603
+
+
+def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]:
+    """
+    Detect presence of an NVIDIA GPU and extract the driver version.
+    Returns:
+        (True, (major, minor)) if detected
+        (False, None) if nvidia-smi is not available
+    """
+    try:
+        nvidia_smi = shutil.which("nvidia-smi")
+        if nvidia_smi is None:
+            raise RuntimeError("nvidia-smi not found in PATH")
+
+        out = (
+            subprocess.check_output(
+                [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL
+            )  # nosec B603
+            .decode("utf-8")
+            .strip()
+        )
+    except Exception:
+        return (False, None)
+
+    # Exemple: "575.64"
+    m = re.match(r"([0-9]+)\.([0-9]+)", out)
+    if not m:
+        return (True, None)
+
+    return (True, (int(m.group(1)), int(m.group(2))))
+
+
+def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool:
+    """
+    Check whether the detected NVIDIA driver satisfies the minimum
+    requirement for CUDA 12.8 on the given operating system.
+    """
+    if drv is None:
+        return False
+    if os_name == "Linux":
+        return drv >= CUDA128_MIN_DRIVER_LINUX
+    if os_name == "Windows":
+        return drv >= CUDA128_MIN_DRIVER_WINDOWS
+    return False
+
+
+def normalize_arch(machine: str) -> str:
+    """
+    Normalize platform.machine() output across operating systems.
+
+    Examples:
+        AMD64   -> x86_64
+        aarch64 -> arm64
+    """
+    m = machine.lower()
+    if m in ("x86_64", "amd64"):
+        return "x86_64"
+    if m in ("aarch64", "arm64"):
+        return "arm64"
+    return machine
+
+
+def download_file(url: str, dst: Path) -> None:
+    """
+    Download a file from the given URL with a progress indicator.
+    """
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading: {url}", flush=True)
+    print(f"       to: {dst}", flush=True)
+
+    try:
+        with requests.get(url, stream=True, timeout=60) as r:
+            r.raise_for_status()
+            total = int(r.headers.get("content-length", 0))
+            with open(dst, "wb") as f:
+                with tqdm(
+                    total=total,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Downloading {dst.name}",
+                ) as pbar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+    except Exception as e:
+        raise e
+
+
+def extract_archive(archive: Path, dst_dir: Path) -> None:
+    """
+    Extract a ZIP archive to the destination directory.
+    """
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Extracting: {archive} -> {dst_dir}", flush=True)
+    with zipfile.ZipFile(archive, "r") as z:
+        z.extractall(dst_dir)
+    archive.unlink()
+
+
+def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool):
+    os_name = platform.system()
+    arch = normalize_arch(platform.machine())
+    has_nvidia, drv = detect_nvidia_driver()
+
+    if os_name not in ("Linux", "Windows", "Darwin"):
+        raise NameError(f"Unsupported OS: {os_name}")
+
+    if arch not in ("x86_64", "arm64"):
+        raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)")
+
+    flavor = "cpu"
+    if force_cuda:
+        if not has_nvidia or not driver_ok_for_cuda(os_name, drv):
+            raise NameError(
+                "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}"
+            )
+
+        flavor = "cu128"
+    elif not force_cpu:
+        if has_nvidia and driver_ok_for_cuda(os_name, drv):
+            flavor = "cu128"
+        else:
+            flavor = "cpu"
+
+    print(f"System: {os_name} {arch}", flush=True)
+    print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True)
+    print(f"Selected flavor: {flavor}", flush=True)
+
+    key = (os_name, arch, flavor)
+    if key not in ELX_ASSET_TEMPLATE:
+        raise NameError(f"No elastix asset configured for {key}")
+    if flavor != "cpu" and key not in LIBTORCH_URL:
+        raise NameError(f"No libtorch url configured for {key}")
+
+    install_path = install_path.resolve()
+    install_path.mkdir(parents=True, exist_ok=True)
+
+    elx_asset = ELX_ASSET_TEMPLATE[key]
+    elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}"
+    elx_archive = install_path / elx_asset
+    download_file(elx_url, elx_archive)
+    extract_archive(elx_archive, install_path)
+
+    # -------------------------------------------------------------------------
+    # ZIP archives may drop executable permissions.
+    # Ensure elastix and transformix are executable on Unix platforms.
+    # -------------------------------------------------------------------------
+    if os_name in ("Linux", "Darwin"):
+        for exe in ("elastix", "transformix"):
+            p = install_path / "bin" / exe
+            if p.exists():
+                p.chmod(p.stat().st_mode | stat.S_IEXEC)
+
+    # ---------------------------------------------------------------------
+    # - Download matching LibTorch 12.8 with or without CUDA 12.8
+    # - Extract it
+    # - Move runtime libraries to a location visible to the dynamic loader
+    #
+    # Linux / macOS : prefix/lib
+    # Windows       : prefix (next to elastix.exe)
+    # ---------------------------------------------------------------------
+    lt_url = LIBTORCH_URL[key]
+    lt_archive = install_path / Path(lt_url).name
+    download_file(lt_url, lt_archive)
+    extract_archive(lt_archive, install_path)
+    for p in (install_path / "libtorch" / "lib").iterdir():
+        if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name:
+            shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name)
+
+    shutil.rmtree(install_path / "libtorch")
+    if os_name == "Linux" and flavor == "cu128":
+        shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12")
+
+
+def get_elastix_bin(install_path: Path) -> Path:
+    return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix"))
+
+
+def try_elastix(install_path: Path) -> None:
+    try:
+        subprocess.run(
+            [str(get_elastix_bin(install_path)), "-h"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )  # nosec B603
+    except subprocess.CalledProcessError as e:
+        msg = "Elastix execution failed.\n\n"
+
+        msg += f"Command:\n{' '.join(e.cmd)}\n"
+        msg += f"Return code: {e.returncode}\n\n"
+
+        if e.stderr:
+            msg += "Error output:\n"
+            msg += e.stderr.strip()
+        raise NameError(msg)
+
+    except OSError as e:
+        msg = (
+            "Elastix could not be started.\n\n"
+            "This is usually caused by missing shared libraries "
+            "(e.g. LibTorch or CUDA runtime).\n\n"
+            f"System error:\n{str(e)}"
+        )
+
+        raise NameError(msg)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)"
+    )
+    group = ap.add_mutually_exclusive_group()
+    group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present")
+    group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)")
+    args = vars(ap.parse_args())
+    install_elastix_impact(**args)
+
+
+if __name__ == "__main__":
+    main()

CBCT_CT_TS/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

CBCT_CT_TS/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

CBCT_CT_TS/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

CBCT_CT_TS/Model.py

@@ -0,0 +1,291 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+
+``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
+(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
+onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
+(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
+``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
+needs to register in physical space.
+
+Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+
+    Patch.patch_size: None        # one patch per case (registration is global)
+    batch_size: 1                 # one fixed/moving pair at a time
+
+NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
+    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._models = models
+        self._iterations = iterations
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Copy parameter maps into the work dir, patching the ImpactGPU device field when present."""
+        staged = []
+        for src in self._parameter_maps:
+            dst = work / src.name
+            lines = []
+            for line in src.read_text(encoding="utf-8").splitlines():
+                if line.strip().startswith("(ImpactGPU"):
+                    line = f"(ImpactGPU {device_index})"
+                lines.append(line)
+            dst.write_text("\n".join(lines) + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
+        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the parameter maps reference
+            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
+            # the "N% done/total" progress line) shows real progress during the long registration.
+            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            # ``iterations`` is the total iteration budget declared for the preset (summed over the
+            # chained parameter maps), so the bar spans the whole chain of registration stages.
+            progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            for line in proc.stdout:
+                captured.append(line)
+                if iteration_line.match(line):
+                    progress.update(1)
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
+    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(parameter_maps, models, iterations)
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
+        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
+        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)
+
+
+class ChannelSelect(torch.nn.Module):
+    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
+    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
+    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
+    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
+    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        engine: str = "elastix",
+        parameter_maps: list[str] = [],
+        models: list[str] = [],
+        iterations: int = 1000,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        self.add_module(
+            "Registration",
+            ElastixRegistration(engine, parameter_maps, models, iterations),
+            in_branch=[0, 1, 2, 3],
+            out_branch=["registration"],
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

CBCT_CT_TS/Prediction.yml

@@ -0,0 +1,107 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      engine: elastix
+      parameter_maps:
+      - ParameterMap_CBCT_generic_TS.txt
+      models:
+      - VBoussot/impact-torchscript-models:TS/M852.pt
+      - VBoussot/impact-torchscript-models:TS/M850.pt
+      iterations: 1050
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-CBCT-CT-TS
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

CBCT_CT_TS/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "CBCT/CT preset with TotalSegmentator",
+    "short_description": "Generic CBCT/CT deformable registration using TotalSegmentator features",
+    "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 250, and 200 iterations using 2000 random spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), starting with Dice-based overlap on segmentation outputs and progressively integrating feature-level alignment via L1 distances on selected internal layers (0.3/0.7 then 0.5/0.5 L1/Dice), ending with a final purely feature-based stage. A composite objective (IMPACT + mutual information + bending energy penalty) ensures robust cross-modality alignment while enforcing smooth, physically plausible deformations.",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

CBCT_CT_TS/install.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import platform
+import re
+import shutil
+import stat
+import subprocess  # nosec B404
+import zipfile
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+# -----------------------------------------------------------------------------
+# Elastix + IMPACT binary assets hosted on GitHub Releases.
+#
+# Key format: (OS, ARCH, FLAVOR)
+#   - OS     : platform.system() -> "Linux", "Windows", "Darwin"
+#   - ARCH   : normalized architecture -> "x86_64"
+#   - FLAVOR : "cpu" or "cu128"
+#
+# CPU assets are standalone (LibTorch CPU bundled).
+# CUDA assets do NOT bundle LibTorch (too large for GitHub limits).
+# -----------------------------------------------------------------------------
+ELX_ASSET_TEMPLATE = {
+    ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip",
+    ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip",
+    ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip",
+    ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip",
+    ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Official LibTorch downloads (PyTorch).
+#
+# IMPORTANT:
+# - Version MUST match the one used at build time (ABI compatibility).
+# - Using "shared-with-deps" ensures CUDA runtime libraries are included
+#   (except the NVIDIA driver, which must be installed system-wide).
+# -----------------------------------------------------------------------------
+
+LIBTORCH_URL = {
+    (
+        "Linux",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Darwin",
+        "arm64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Linux",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Minimum NVIDIA driver versions required for CUDA 12.8.
+#
+# The CUDA Toolkit itself is NOT required.
+# Only a sufficiently recent NVIDIA driver must be installed.
+# -----------------------------------------------------------------------------
+CUDA128_MIN_DRIVER_LINUX = (570, 26)
+CUDA128_MIN_DRIVER_WINDOWS = (570, 65)
+
+GITHUB_OWNER = "vboussot"
+GITHUB_REPO = "ImpactElastix"
+GITHUB_TAG = "1.0.0"
+
+
+DEFAULT_PREFIX = Path.cwd() / "elastix-impact"
+
+
+def run_cmd(cmd: list[str]) -> str:
+    return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace")  # nosec B603
+
+
+def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]:
+    """
+    Detect presence of an NVIDIA GPU and extract the driver version.
+    Returns:
+        (True, (major, minor)) if detected
+        (False, None) if nvidia-smi is not available
+    """
+    try:
+        nvidia_smi = shutil.which("nvidia-smi")
+        if nvidia_smi is None:
+            raise RuntimeError("nvidia-smi not found in PATH")
+
+        out = (
+            subprocess.check_output(
+                [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL
+            )  # nosec B603
+            .decode("utf-8")
+            .strip()
+        )
+    except Exception:
+        return (False, None)
+
+    # Exemple: "575.64"
+    m = re.match(r"([0-9]+)\.([0-9]+)", out)
+    if not m:
+        return (True, None)
+
+    return (True, (int(m.group(1)), int(m.group(2))))
+
+
+def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool:
+    """
+    Check whether the detected NVIDIA driver satisfies the minimum
+    requirement for CUDA 12.8 on the given operating system.
+    """
+    if drv is None:
+        return False
+    if os_name == "Linux":
+        return drv >= CUDA128_MIN_DRIVER_LINUX
+    if os_name == "Windows":
+        return drv >= CUDA128_MIN_DRIVER_WINDOWS
+    return False
+
+
+def normalize_arch(machine: str) -> str:
+    """
+    Normalize platform.machine() output across operating systems.
+
+    Examples:
+        AMD64   -> x86_64
+        aarch64 -> arm64
+    """
+    m = machine.lower()
+    if m in ("x86_64", "amd64"):
+        return "x86_64"
+    if m in ("aarch64", "arm64"):
+        return "arm64"
+    return machine
+
+
+def download_file(url: str, dst: Path) -> None:
+    """
+    Download a file from the given URL with a progress indicator.
+    """
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading: {url}", flush=True)
+    print(f"       to: {dst}", flush=True)
+
+    try:
+        with requests.get(url, stream=True, timeout=60) as r:
+            r.raise_for_status()
+            total = int(r.headers.get("content-length", 0))
+            with open(dst, "wb") as f:
+                with tqdm(
+                    total=total,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Downloading {dst.name}",
+                ) as pbar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+    except Exception as e:
+        raise e
+
+
+def extract_archive(archive: Path, dst_dir: Path) -> None:
+    """
+    Extract a ZIP archive to the destination directory.
+    """
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Extracting: {archive} -> {dst_dir}", flush=True)
+    with zipfile.ZipFile(archive, "r") as z:
+        z.extractall(dst_dir)
+    archive.unlink()
+
+
+def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool):
+    os_name = platform.system()
+    arch = normalize_arch(platform.machine())
+    has_nvidia, drv = detect_nvidia_driver()
+
+    if os_name not in ("Linux", "Windows", "Darwin"):
+        raise NameError(f"Unsupported OS: {os_name}")
+
+    if arch not in ("x86_64", "arm64"):
+        raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)")
+
+    flavor = "cpu"
+    if force_cuda:
+        if not has_nvidia or not driver_ok_for_cuda(os_name, drv):
+            raise NameError(
+                "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}"
+            )
+
+        flavor = "cu128"
+    elif not force_cpu:
+        if has_nvidia and driver_ok_for_cuda(os_name, drv):
+            flavor = "cu128"
+        else:
+            flavor = "cpu"
+
+    print(f"System: {os_name} {arch}", flush=True)
+    print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True)
+    print(f"Selected flavor: {flavor}", flush=True)
+
+    key = (os_name, arch, flavor)
+    if key not in ELX_ASSET_TEMPLATE:
+        raise NameError(f"No elastix asset configured for {key}")
+    if flavor != "cpu" and key not in LIBTORCH_URL:
+        raise NameError(f"No libtorch url configured for {key}")
+
+    install_path = install_path.resolve()
+    install_path.mkdir(parents=True, exist_ok=True)
+
+    elx_asset = ELX_ASSET_TEMPLATE[key]
+    elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}"
+    elx_archive = install_path / elx_asset
+    download_file(elx_url, elx_archive)
+    extract_archive(elx_archive, install_path)
+
+    # -------------------------------------------------------------------------
+    # ZIP archives may drop executable permissions.
+    # Ensure elastix and transformix are executable on Unix platforms.
+    # -------------------------------------------------------------------------
+    if os_name in ("Linux", "Darwin"):
+        for exe in ("elastix", "transformix"):
+            p = install_path / "bin" / exe
+            if p.exists():
+                p.chmod(p.stat().st_mode | stat.S_IEXEC)
+
+    # ---------------------------------------------------------------------
+    # - Download matching LibTorch 12.8 with or without CUDA 12.8
+    # - Extract it
+    # - Move runtime libraries to a location visible to the dynamic loader
+    #
+    # Linux / macOS : prefix/lib
+    # Windows       : prefix (next to elastix.exe)
+    # ---------------------------------------------------------------------
+    lt_url = LIBTORCH_URL[key]
+    lt_archive = install_path / Path(lt_url).name
+    download_file(lt_url, lt_archive)
+    extract_archive(lt_archive, install_path)
+    for p in (install_path / "libtorch" / "lib").iterdir():
+        if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name:
+            shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name)
+
+    shutil.rmtree(install_path / "libtorch")
+    if os_name == "Linux" and flavor == "cu128":
+        shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12")
+
+
+def get_elastix_bin(install_path: Path) -> Path:
+    return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix"))
+
+
+def try_elastix(install_path: Path) -> None:
+    try:
+        subprocess.run(
+            [str(get_elastix_bin(install_path)), "-h"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )  # nosec B603
+    except subprocess.CalledProcessError as e:
+        msg = "Elastix execution failed.\n\n"
+
+        msg += f"Command:\n{' '.join(e.cmd)}\n"
+        msg += f"Return code: {e.returncode}\n\n"
+
+        if e.stderr:
+            msg += "Error output:\n"
+            msg += e.stderr.strip()
+        raise NameError(msg)
+
+    except OSError as e:
+        msg = (
+            "Elastix could not be started.\n\n"
+            "This is usually caused by missing shared libraries "
+            "(e.g. LibTorch or CUDA runtime).\n\n"
+            f"System error:\n{str(e)}"
+        )
+
+        raise NameError(msg)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)"
+    )
+    group = ap.add_mutually_exclusive_group()
+    group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present")
+    group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)")
+    args = vars(ap.parse_args())
+    install_elastix_impact(**args)
+
+
+if __name__ == "__main__":
+    main()

Generic_Rigid/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

Generic_Rigid/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

Generic_Rigid/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

Generic_Rigid/Model.py

@@ -0,0 +1,291 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+
+``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
+(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
+onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
+(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
+``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
+needs to register in physical space.
+
+Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+
+    Patch.patch_size: None        # one patch per case (registration is global)
+    batch_size: 1                 # one fixed/moving pair at a time
+
+NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
+    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._models = models
+        self._iterations = iterations
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Copy parameter maps into the work dir, patching the ImpactGPU device field when present."""
+        staged = []
+        for src in self._parameter_maps:
+            dst = work / src.name
+            lines = []
+            for line in src.read_text(encoding="utf-8").splitlines():
+                if line.strip().startswith("(ImpactGPU"):
+                    line = f"(ImpactGPU {device_index})"
+                lines.append(line)
+            dst.write_text("\n".join(lines) + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
+        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the parameter maps reference
+            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
+            # the "N% done/total" progress line) shows real progress during the long registration.
+            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            # ``iterations`` is the total iteration budget declared for the preset (summed over the
+            # chained parameter maps), so the bar spans the whole chain of registration stages.
+            progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            for line in proc.stdout:
+                captured.append(line)
+                if iteration_line.match(line):
+                    progress.update(1)
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
+    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(parameter_maps, models, iterations)
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
+        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
+        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)
+
+
+class ChannelSelect(torch.nn.Module):
+    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
+    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
+    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
+    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
+    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        engine: str = "elastix",
+        parameter_maps: list[str] = [],
+        models: list[str] = [],
+        iterations: int = 1000,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        self.add_module(
+            "Registration",
+            ElastixRegistration(engine, parameter_maps, models, iterations),
+            in_branch=[0, 1, 2, 3],
+            out_branch=["registration"],
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

Generic_Rigid/Prediction.yml

@@ -0,0 +1,105 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      engine: elastix
+      parameter_maps:
+      - Parameters_Rigid.txt
+      models: []
+      iterations: 250
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-Generic-Rigid
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

Generic_Rigid/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "Generic Rigid",
+    "short_description": "Rigid registration using mutual information and a multi-resolution pyramid.",
+    "description": "This preset performs rigid alignment using an Euler transform optimized with Adaptive Stochastic Gradient Descent. It uses a 4-level multi-resolution strategy and Mattes mutual information as similarity metric. Initial alignment based on image centers are enabled to ensure robust convergence for multimodal images.",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

Generic_Rigid/install.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import platform
+import re
+import shutil
+import stat
+import subprocess  # nosec B404
+import zipfile
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+# -----------------------------------------------------------------------------
+# Elastix + IMPACT binary assets hosted on GitHub Releases.
+#
+# Key format: (OS, ARCH, FLAVOR)
+#   - OS     : platform.system() -> "Linux", "Windows", "Darwin"
+#   - ARCH   : normalized architecture -> "x86_64"
+#   - FLAVOR : "cpu" or "cu128"
+#
+# CPU assets are standalone (LibTorch CPU bundled).
+# CUDA assets do NOT bundle LibTorch (too large for GitHub limits).
+# -----------------------------------------------------------------------------
+ELX_ASSET_TEMPLATE = {
+    ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip",
+    ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip",
+    ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip",
+    ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip",
+    ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Official LibTorch downloads (PyTorch).
+#
+# IMPORTANT:
+# - Version MUST match the one used at build time (ABI compatibility).
+# - Using "shared-with-deps" ensures CUDA runtime libraries are included
+#   (except the NVIDIA driver, which must be installed system-wide).
+# -----------------------------------------------------------------------------
+
+LIBTORCH_URL = {
+    (
+        "Linux",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Darwin",
+        "arm64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Linux",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Minimum NVIDIA driver versions required for CUDA 12.8.
+#
+# The CUDA Toolkit itself is NOT required.
+# Only a sufficiently recent NVIDIA driver must be installed.
+# -----------------------------------------------------------------------------
+CUDA128_MIN_DRIVER_LINUX = (570, 26)
+CUDA128_MIN_DRIVER_WINDOWS = (570, 65)
+
+GITHUB_OWNER = "vboussot"
+GITHUB_REPO = "ImpactElastix"
+GITHUB_TAG = "1.0.0"
+
+
+DEFAULT_PREFIX = Path.cwd() / "elastix-impact"
+
+
+def run_cmd(cmd: list[str]) -> str:
+    return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace")  # nosec B603
+
+
+def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]:
+    """
+    Detect presence of an NVIDIA GPU and extract the driver version.
+    Returns:
+        (True, (major, minor)) if detected
+        (False, None) if nvidia-smi is not available
+    """
+    try:
+        nvidia_smi = shutil.which("nvidia-smi")
+        if nvidia_smi is None:
+            raise RuntimeError("nvidia-smi not found in PATH")
+
+        out = (
+            subprocess.check_output(
+                [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL
+            )  # nosec B603
+            .decode("utf-8")
+            .strip()
+        )
+    except Exception:
+        return (False, None)
+
+    # Exemple: "575.64"
+    m = re.match(r"([0-9]+)\.([0-9]+)", out)
+    if not m:
+        return (True, None)
+
+    return (True, (int(m.group(1)), int(m.group(2))))
+
+
+def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool:
+    """
+    Check whether the detected NVIDIA driver satisfies the minimum
+    requirement for CUDA 12.8 on the given operating system.
+    """
+    if drv is None:
+        return False
+    if os_name == "Linux":
+        return drv >= CUDA128_MIN_DRIVER_LINUX
+    if os_name == "Windows":
+        return drv >= CUDA128_MIN_DRIVER_WINDOWS
+    return False
+
+
+def normalize_arch(machine: str) -> str:
+    """
+    Normalize platform.machine() output across operating systems.
+
+    Examples:
+        AMD64   -> x86_64
+        aarch64 -> arm64
+    """
+    m = machine.lower()
+    if m in ("x86_64", "amd64"):
+        return "x86_64"
+    if m in ("aarch64", "arm64"):
+        return "arm64"
+    return machine
+
+
+def download_file(url: str, dst: Path) -> None:
+    """
+    Download a file from the given URL with a progress indicator.
+    """
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading: {url}", flush=True)
+    print(f"       to: {dst}", flush=True)
+
+    try:
+        with requests.get(url, stream=True, timeout=60) as r:
+            r.raise_for_status()
+            total = int(r.headers.get("content-length", 0))
+            with open(dst, "wb") as f:
+                with tqdm(
+                    total=total,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Downloading {dst.name}",
+                ) as pbar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+    except Exception as e:
+        raise e
+
+
+def extract_archive(archive: Path, dst_dir: Path) -> None:
+    """
+    Extract a ZIP archive to the destination directory.
+    """
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Extracting: {archive} -> {dst_dir}", flush=True)
+    with zipfile.ZipFile(archive, "r") as z:
+        z.extractall(dst_dir)
+    archive.unlink()
+
+
+def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool):
+    os_name = platform.system()
+    arch = normalize_arch(platform.machine())
+    has_nvidia, drv = detect_nvidia_driver()
+
+    if os_name not in ("Linux", "Windows", "Darwin"):
+        raise NameError(f"Unsupported OS: {os_name}")
+
+    if arch not in ("x86_64", "arm64"):
+        raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)")
+
+    flavor = "cpu"
+    if force_cuda:
+        if not has_nvidia or not driver_ok_for_cuda(os_name, drv):
+            raise NameError(
+                "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}"
+            )
+
+        flavor = "cu128"
+    elif not force_cpu:
+        if has_nvidia and driver_ok_for_cuda(os_name, drv):
+            flavor = "cu128"
+        else:
+            flavor = "cpu"
+
+    print(f"System: {os_name} {arch}", flush=True)
+    print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True)
+    print(f"Selected flavor: {flavor}", flush=True)
+
+    key = (os_name, arch, flavor)
+    if key not in ELX_ASSET_TEMPLATE:
+        raise NameError(f"No elastix asset configured for {key}")
+    if flavor != "cpu" and key not in LIBTORCH_URL:
+        raise NameError(f"No libtorch url configured for {key}")
+
+    install_path = install_path.resolve()
+    install_path.mkdir(parents=True, exist_ok=True)
+
+    elx_asset = ELX_ASSET_TEMPLATE[key]
+    elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}"
+    elx_archive = install_path / elx_asset
+    download_file(elx_url, elx_archive)
+    extract_archive(elx_archive, install_path)
+
+    # -------------------------------------------------------------------------
+    # ZIP archives may drop executable permissions.
+    # Ensure elastix and transformix are executable on Unix platforms.
+    # -------------------------------------------------------------------------
+    if os_name in ("Linux", "Darwin"):
+        for exe in ("elastix", "transformix"):
+            p = install_path / "bin" / exe
+            if p.exists():
+                p.chmod(p.stat().st_mode | stat.S_IEXEC)
+
+    # ---------------------------------------------------------------------
+    # - Download matching LibTorch 12.8 with or without CUDA 12.8
+    # - Extract it
+    # - Move runtime libraries to a location visible to the dynamic loader
+    #
+    # Linux / macOS : prefix/lib
+    # Windows       : prefix (next to elastix.exe)
+    # ---------------------------------------------------------------------
+    lt_url = LIBTORCH_URL[key]
+    lt_archive = install_path / Path(lt_url).name
+    download_file(lt_url, lt_archive)
+    extract_archive(lt_archive, install_path)
+    for p in (install_path / "libtorch" / "lib").iterdir():
+        if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name:
+            shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name)
+
+    shutil.rmtree(install_path / "libtorch")
+    if os_name == "Linux" and flavor == "cu128":
+        shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12")
+
+
+def get_elastix_bin(install_path: Path) -> Path:
+    return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix"))
+
+
+def try_elastix(install_path: Path) -> None:
+    try:
+        subprocess.run(
+            [str(get_elastix_bin(install_path)), "-h"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )  # nosec B603
+    except subprocess.CalledProcessError as e:
+        msg = "Elastix execution failed.\n\n"
+
+        msg += f"Command:\n{' '.join(e.cmd)}\n"
+        msg += f"Return code: {e.returncode}\n\n"
+
+        if e.stderr:
+            msg += "Error output:\n"
+            msg += e.stderr.strip()
+        raise NameError(msg)
+
+    except OSError as e:
+        msg = (
+            "Elastix could not be started.\n\n"
+            "This is usually caused by missing shared libraries "
+            "(e.g. LibTorch or CUDA runtime).\n\n"
+            f"System error:\n{str(e)}"
+        )
+
+        raise NameError(msg)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)"
+    )
+    group = ap.add_mutually_exclusive_group()
+    group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present")
+    group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)")
+    args = vars(ap.parse_args())
+    install_elastix_impact(**args)
+
+
+if __name__ == "__main__":
+    main()

Generic_Rigid_BSpline/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

Generic_Rigid_BSpline/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

Generic_Rigid_BSpline/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

Generic_Rigid_BSpline/Model.py

@@ -0,0 +1,291 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+
+``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
+(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
+onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
+(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
+``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
+needs to register in physical space.
+
+Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+
+    Patch.patch_size: None        # one patch per case (registration is global)
+    batch_size: 1                 # one fixed/moving pair at a time
+
+NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
+    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._models = models
+        self._iterations = iterations
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Copy parameter maps into the work dir, patching the ImpactGPU device field when present."""
+        staged = []
+        for src in self._parameter_maps:
+            dst = work / src.name
+            lines = []
+            for line in src.read_text(encoding="utf-8").splitlines():
+                if line.strip().startswith("(ImpactGPU"):
+                    line = f"(ImpactGPU {device_index})"
+                lines.append(line)
+            dst.write_text("\n".join(lines) + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
+        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the parameter maps reference
+            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
+            # the "N% done/total" progress line) shows real progress during the long registration.
+            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            # ``iterations`` is the total iteration budget declared for the preset (summed over the
+            # chained parameter maps), so the bar spans the whole chain of registration stages.
+            progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            for line in proc.stdout:
+                captured.append(line)
+                if iteration_line.match(line):
+                    progress.update(1)
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
+    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(parameter_maps, models, iterations)
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
+        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
+        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)
+
+
+class ChannelSelect(torch.nn.Module):
+    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
+    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
+    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
+    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
+    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        engine: str = "elastix",
+        parameter_maps: list[str] = [],
+        models: list[str] = [],
+        iterations: int = 1000,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        self.add_module(
+            "Registration",
+            ElastixRegistration(engine, parameter_maps, models, iterations),
+            in_branch=[0, 1, 2, 3],
+            out_branch=["registration"],
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

Generic_Rigid_BSpline/Parameters_Rigid.txt

@@ -0,0 +1,147 @@
+// Example parameter file for rotation registration
+// C-style comments: //
+
+// The internal pixel type, used for internal computations
+// Leave to float in general. 
+// NB: this is not the type of the input images! The pixel 
+// type of the input images is automatically read from the 
+// images themselves.
+// This setting can be changed to "short" to save some memory
+// in case of very large 3D images.
+(FixedInternalImagePixelType "float")
+(MovingInternalImagePixelType "float")
+
+// **************** Main Components **************************
+
+// The following components should usually be left as they are:
+(Registration "MultiResolutionRegistration")
+(Interpolator "BSplineInterpolator")
+(ResampleInterpolator "FinalBSplineInterpolator")
+(Resampler "DefaultResampler")
+
+// These may be changed to Fixed/MovingSmoothingImagePyramid.
+// See the manual.
+(FixedImagePyramid "FixedRecursiveImagePyramid")
+(MovingImagePyramid "MovingRecursiveImagePyramid")
+
+// The following components are most important:
+// The optimizer AdaptiveStochasticGradientDescent (ASGD) works
+// quite ok in general. The Transform and Metric are important
+// and need to be chosen careful for each application. See manual.
+(Optimizer "AdaptiveStochasticGradientDescent")
+(Transform "EulerTransform")
+(Metric "AdvancedMattesMutualInformation")
+
+// ***************** Transformation **************************
+
+// Scales the rotations compared to the translations, to make
+// sure they are in the same range. In general, it's best to  
+// use automatic scales estimation:
+(AutomaticScalesEstimation "true")
+
+// Automatically guess an initial translation by aligning the
+// geometric centers of the fixed and moving.
+(AutomaticTransformInitialization "true")
+
+// Whether transforms are combined by composition or by addition.
+// In generally, Compose is the best option in most cases.
+// It does not influence the results very much.
+(HowToCombineTransforms "Compose")
+
+//Save composite ITK transform
+(ITKTransformOutputFileNameExtension "itk.txt")
+(WriteITKCompositeTransform "true")
+
+// ******************* Similarity measure *********************
+
+// Number of grey level bins in each resolution level,
+// for the mutual information. 16 or 32 usually works fine.
+// You could also employ a hierarchical strategy:
+//(NumberOfHistogramBins 16 32 64)
+(NumberOfHistogramBins 32)
+
+// If you use a mask, this option is important. 
+// If the mask serves as region of interest, set it to false.
+// If the mask indicates which pixels are valid, then set it to true.
+// If you do not use a mask, the option doesn't matter.
+(ErodeMask "false")
+
+// ******************** Multiresolution **********************
+
+// The number of resolutions. 1 Is only enough if the expected
+// deformations are small. 3 or 4 mostly works fine. For large
+// images and large deformations, 5 or 6 may even be useful.
+(NumberOfResolutions 4)
+
+// The downsampling/blurring factors for the image pyramids.
+// By default, the images are downsampled by a factor of 2
+// compared to the next resolution.
+// So, in 2D, with 4 resolutions, the following schedule is used:
+//(ImagePyramidSchedule 8 8  4 4  2 2  1 1 )
+// And in 3D:
+//(ImagePyramidSchedule 8 8 8  4 4 4  2 2 2  1 1 1 )
+// You can specify any schedule, for example:
+//(ImagePyramidSchedule 4 4  4 3  2 1  1 1 )
+// Make sure that the number of elements equals the number
+// of resolutions times the image dimension.
+
+// ******************* Optimizer ****************************
+
+// Maximum number of iterations in each resolution level:
+// 200-500 works usually fine for rigid registration.
+// For more robustness, you may increase this to 1000-2000.
+(MaximumNumberOfIterations 250)
+
+// The step size of the optimizer, in mm. By default the voxel size is used.
+// which usually works well. In case of unusual high-resolution images
+// (eg histology) it is necessary to increase this value a bit, to the size
+// of the "smallest visible structure" in the image:
+//(MaximumStepLength 1.0)
+
+// **************** Image sampling **********************
+
+// Number of spatial samples used to compute the mutual
+// information (and its derivative) in each iteration.
+// With an AdaptiveStochasticGradientDescent optimizer,
+// in combination with the two options below, around 2000
+// samples may already suffice.
+(NumberOfSpatialSamples 2048)
+
+// Refresh these spatial samples in every iteration, and select
+// them randomly. See the manual for information on other sampling
+// strategies.
+(NewSamplesEveryIteration "true")
+(ImageSampler "Random")
+
+// ************* Interpolation and Resampling ****************
+
+// Order of B-Spline interpolation used during registration/optimisation.
+// It may improve accuracy if you set this to 3. Never use 0.
+// An order of 1 gives linear interpolation. This is in most 
+// applications a good choice.
+(BSplineInterpolationOrder 1)
+
+// Order of B-Spline interpolation used for applying the final
+// deformation.
+// 3 gives good accuracy; recommended in most cases.
+// 1 gives worse accuracy (linear interpolation)
+// 0 gives worst accuracy, but is appropriate for binary images
+// (masks, segmentations); equivalent to nearest neighbor interpolation.
+(FinalBSplineInterpolationOrder 3)
+
+//Default pixel value for pixels that come from outside the picture:
+(DefaultPixelValue 0)
+
+// Choose whether to generate the deformed moving image.
+// You can save some time by setting this to false, if you are
+// only interested in the final (nonrigidly) deformed moving image
+// for example.
+(WriteResultImage "false")
+
+// The pixel type and format of the resulting deformed moving image
+(ResultImagePixelType "short")
+(ResultImageFormat "mha")
+(WriteTransformParametersEachIteration "false")
+(WriteTransformParametersEachResolution "false")
+(WriteIterationInfo "false")
+

Generic_Rigid_BSpline/Prediction.yml

@@ -0,0 +1,106 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      engine: elastix
+      parameter_maps:
+      - Parameters_Rigid.txt
+      - Parameters_BSpline.txt
+      models: []
+      iterations: 3000
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-Generic-Rigid-BSpline
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

Generic_Rigid_BSpline/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "Generic Rigid + BSpline",
+    "short_description": "Two-stage registration: rigid alignment followed by BSpline refinement.",
+    "description": "A combined registration strategy: first a rigid Euler transform corrects global misalignment, then a BSpline model captures localized anatomical deformations. Both stages use a multi-resolution pyramid, mutual information, and stochastic optimization for robust performance across a wide range of multimodal imaging scenarios.",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

Generic_Rigid_BSpline/install.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import platform
+import re
+import shutil
+import stat
+import subprocess  # nosec B404
+import zipfile
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+# -----------------------------------------------------------------------------
+# Elastix + IMPACT binary assets hosted on GitHub Releases.
+#
+# Key format: (OS, ARCH, FLAVOR)
+#   - OS     : platform.system() -> "Linux", "Windows", "Darwin"
+#   - ARCH   : normalized architecture -> "x86_64"
+#   - FLAVOR : "cpu" or "cu128"
+#
+# CPU assets are standalone (LibTorch CPU bundled).
+# CUDA assets do NOT bundle LibTorch (too large for GitHub limits).
+# -----------------------------------------------------------------------------
+ELX_ASSET_TEMPLATE = {
+    ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip",
+    ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip",
+    ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip",
+    ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip",
+    ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Official LibTorch downloads (PyTorch).
+#
+# IMPORTANT:
+# - Version MUST match the one used at build time (ABI compatibility).
+# - Using "shared-with-deps" ensures CUDA runtime libraries are included
+#   (except the NVIDIA driver, which must be installed system-wide).
+# -----------------------------------------------------------------------------
+
+LIBTORCH_URL = {
+    (
+        "Linux",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Darwin",
+        "arm64",
+        "cpu",
+    ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip",
+    (
+        "Linux",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip",
+    (
+        "Windows",
+        "x86_64",
+        "cu128",
+    ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip",
+}
+
+# -----------------------------------------------------------------------------
+# Minimum NVIDIA driver versions required for CUDA 12.8.
+#
+# The CUDA Toolkit itself is NOT required.
+# Only a sufficiently recent NVIDIA driver must be installed.
+# -----------------------------------------------------------------------------
+CUDA128_MIN_DRIVER_LINUX = (570, 26)
+CUDA128_MIN_DRIVER_WINDOWS = (570, 65)
+
+GITHUB_OWNER = "vboussot"
+GITHUB_REPO = "ImpactElastix"
+GITHUB_TAG = "1.0.0"
+
+
+DEFAULT_PREFIX = Path.cwd() / "elastix-impact"
+
+
+def run_cmd(cmd: list[str]) -> str:
+    return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace")  # nosec B603
+
+
+def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]:
+    """
+    Detect presence of an NVIDIA GPU and extract the driver version.
+    Returns:
+        (True, (major, minor)) if detected
+        (False, None) if nvidia-smi is not available
+    """
+    try:
+        nvidia_smi = shutil.which("nvidia-smi")
+        if nvidia_smi is None:
+            raise RuntimeError("nvidia-smi not found in PATH")
+
+        out = (
+            subprocess.check_output(
+                [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL
+            )  # nosec B603
+            .decode("utf-8")
+            .strip()
+        )
+    except Exception:
+        return (False, None)
+
+    # Exemple: "575.64"
+    m = re.match(r"([0-9]+)\.([0-9]+)", out)
+    if not m:
+        return (True, None)
+
+    return (True, (int(m.group(1)), int(m.group(2))))
+
+
+def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool:
+    """
+    Check whether the detected NVIDIA driver satisfies the minimum
+    requirement for CUDA 12.8 on the given operating system.
+    """
+    if drv is None:
+        return False
+    if os_name == "Linux":
+        return drv >= CUDA128_MIN_DRIVER_LINUX
+    if os_name == "Windows":
+        return drv >= CUDA128_MIN_DRIVER_WINDOWS
+    return False
+
+
+def normalize_arch(machine: str) -> str:
+    """
+    Normalize platform.machine() output across operating systems.
+
+    Examples:
+        AMD64   -> x86_64
+        aarch64 -> arm64
+    """
+    m = machine.lower()
+    if m in ("x86_64", "amd64"):
+        return "x86_64"
+    if m in ("aarch64", "arm64"):
+        return "arm64"
+    return machine
+
+
+def download_file(url: str, dst: Path) -> None:
+    """
+    Download a file from the given URL with a progress indicator.
+    """
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading: {url}", flush=True)
+    print(f"       to: {dst}", flush=True)
+
+    try:
+        with requests.get(url, stream=True, timeout=60) as r:
+            r.raise_for_status()
+            total = int(r.headers.get("content-length", 0))
+            with open(dst, "wb") as f:
+                with tqdm(
+                    total=total,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Downloading {dst.name}",
+                ) as pbar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+    except Exception as e:
+        raise e
+
+
+def extract_archive(archive: Path, dst_dir: Path) -> None:
+    """
+    Extract a ZIP archive to the destination directory.
+    """
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Extracting: {archive} -> {dst_dir}", flush=True)
+    with zipfile.ZipFile(archive, "r") as z:
+        z.extractall(dst_dir)
+    archive.unlink()
+
+
+def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool):
+    os_name = platform.system()
+    arch = normalize_arch(platform.machine())
+    has_nvidia, drv = detect_nvidia_driver()
+
+    if os_name not in ("Linux", "Windows", "Darwin"):
+        raise NameError(f"Unsupported OS: {os_name}")
+
+    if arch not in ("x86_64", "arm64"):
+        raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)")
+
+    flavor = "cpu"
+    if force_cuda:
+        if not has_nvidia or not driver_ok_for_cuda(os_name, drv):
+            raise NameError(
+                "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}"
+            )
+
+        flavor = "cu128"
+    elif not force_cpu:
+        if has_nvidia and driver_ok_for_cuda(os_name, drv):
+            flavor = "cu128"
+        else:
+            flavor = "cpu"
+
+    print(f"System: {os_name} {arch}", flush=True)
+    print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True)
+    print(f"Selected flavor: {flavor}", flush=True)
+
+    key = (os_name, arch, flavor)
+    if key not in ELX_ASSET_TEMPLATE:
+        raise NameError(f"No elastix asset configured for {key}")
+    if flavor != "cpu" and key not in LIBTORCH_URL:
+        raise NameError(f"No libtorch url configured for {key}")
+
+    install_path = install_path.resolve()
+    install_path.mkdir(parents=True, exist_ok=True)
+
+    elx_asset = ELX_ASSET_TEMPLATE[key]
+    elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}"
+    elx_archive = install_path / elx_asset
+    download_file(elx_url, elx_archive)
+    extract_archive(elx_archive, install_path)
+
+    # -------------------------------------------------------------------------
+    # ZIP archives may drop executable permissions.
+    # Ensure elastix and transformix are executable on Unix platforms.
+    # -------------------------------------------------------------------------
+    if os_name in ("Linux", "Darwin"):
+        for exe in ("elastix", "transformix"):
+            p = install_path / "bin" / exe
+            if p.exists():
+                p.chmod(p.stat().st_mode | stat.S_IEXEC)
+
+    # ---------------------------------------------------------------------
+    # - Download matching LibTorch 12.8 with or without CUDA 12.8
+    # - Extract it
+    # - Move runtime libraries to a location visible to the dynamic loader
+    #
+    # Linux / macOS : prefix/lib
+    # Windows       : prefix (next to elastix.exe)
+    # ---------------------------------------------------------------------
+    lt_url = LIBTORCH_URL[key]
+    lt_archive = install_path / Path(lt_url).name
+    download_file(lt_url, lt_archive)
+    extract_archive(lt_archive, install_path)
+    for p in (install_path / "libtorch" / "lib").iterdir():
+        if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name:
+            shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name)
+
+    shutil.rmtree(install_path / "libtorch")
+    if os_name == "Linux" and flavor == "cu128":
+        shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12")
+
+
+def get_elastix_bin(install_path: Path) -> Path:
+    return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix"))
+
+
+def try_elastix(install_path: Path) -> None:
+    try:
+        subprocess.run(
+            [str(get_elastix_bin(install_path)), "-h"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )  # nosec B603
+    except subprocess.CalledProcessError as e:
+        msg = "Elastix execution failed.\n\n"
+
+        msg += f"Command:\n{' '.join(e.cmd)}\n"
+        msg += f"Return code: {e.returncode}\n\n"
+
+        if e.stderr:
+            msg += "Error output:\n"
+            msg += e.stderr.strip()
+        raise NameError(msg)
+
+    except OSError as e:
+        msg = (
+            "Elastix could not be started.\n\n"
+            "This is usually caused by missing shared libraries "
+            "(e.g. LibTorch or CUDA runtime).\n\n"
+            f"System error:\n{str(e)}"
+        )
+
+        raise NameError(msg)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)"
+    )
+    group = ap.add_mutually_exclusive_group()
+    group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present")
+    group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)")
+    args = vars(ap.parse_args())
+    install_elastix_impact(**args)
+
+
+if __name__ == "__main__":
+    main()

ImpactReg/UncertaintyDVF.py

@@ -1,15 +0,0 @@
-from konfai.data.transform import Transform
-from konfai.utils.dataset import Attribute
-import torch
-
-class DVFToNorm(Transform):
-
-    def __init__(self) -> None:
-        super().__init__()
-
-    def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        cache_attribute["Origin"] = cache_attribute.pop_np_array("Origin")[:tensors.dim()-2]
-        cache_attribute["Spacing"] = cache_attribute.pop_np_array("Spacing")[:tensors.dim()-2]
-        cache_attribute.pop("Direction")
-        cache_attribute["Direction"] = torch.eye(tensors.dim()-2).flatten()
-        return torch.linalg.norm(tensors, dim=-1)

ImpactReg/app.json

@@ -1,7 +0,0 @@
-{
-    "display_name": "ImpactReg",
-    "short_description": "",
-    "description": "",
-    "tta": 2,
-    "mc_dropout": false
-}

MR_CT_HeadNeck/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

MR_CT_HeadNeck/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg