Add itk-impact ConvexAdam presets (coarse, fine, composite)

ConvexAdam_Coarse/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

ConvexAdam_Coarse/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

ConvexAdam_Coarse/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

ConvexAdam_Coarse/Model.py

@@ -0,0 +1,471 @@
+# 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
+
+"""ConvexAdam (itk-impact) registration as a self-contained KonfAI model.
+
+Same idiomatic ``add_module`` graph and the same output contract as the elastix preset
+(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``),
+so the orchestrator / app.json / ensemble / uncertainty are unchanged. The engine here is
+the native, in-memory itk-impact ConvexAdam pipeline (``pip install itk-impact``) instead of
+the elastix binary:
+
+    (optional) moments + affine Mattes-MI          [ITKv4 linear pre-align]
+      -> ImpactCoarseRegistration                   [coupled-convex init, IMPACT features]
+      -> ImpactFineRegistration                     [Adam instance optimisation, IMPACT features]
+
+The IMPACT feature models (e.g. MIND) are TorchScript ``.pt`` files fetched from Hugging Face
+and wrapped as ``itk.ModelConfiguration`` — the same models the elastix presets use.
+
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
+"""
+
+import itk
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+DIM = 3
+_IMAGE_F = itk.Image[itk.F, DIM]
+
+
+def _coarse_registration_type():
+    """The coupled-convex initializer, tolerant to the two names the wrapping has shipped under."""
+    cls = getattr(itk, "ImpactCoarseRegistration", None) or getattr(itk, "ImpactConvexAdamInitializer", None)
+    if cls is None:
+        raise RuntimeError(
+            "itk-impact does not expose ImpactCoarseRegistration / ImpactConvexAdamInitializer; "
+            "install a build with the ConvexAdam registration filters."
+        )
+    return cls[_IMAGE_F, _IMAGE_F]
+
+
+def _fine_registration_type():
+    """The Adam instance-optimisation stage, tolerant to the two names the wrapping has shipped under."""
+    cls = getattr(itk, "ImpactFineRegistration", None) or getattr(itk, "ImpactTorchAdamRegistration", None)
+    if cls is None:
+        raise RuntimeError(
+            "itk-impact does not expose ImpactFineRegistration / ImpactTorchAdamRegistration; "
+            "install a build with the ConvexAdam registration filters."
+        )
+    return cls[_IMAGE_F, _IMAGE_F]
+
+
+def _sitk_to_itk(image: sitk.Image) -> "itk.Image":
+    """Copy a scalar SimpleITK image (with its geometry) into an ``itk.Image[F, 3]``."""
+    itk_image = itk.image_from_array(sitk.GetArrayFromImage(image).astype(np.float32))
+    itk_image.SetOrigin([float(v) for v in image.GetOrigin()])
+    itk_image.SetSpacing([float(v) for v in image.GetSpacing()])
+    itk_image.SetDirection(itk.matrix_from_array(np.asarray(image.GetDirection(), dtype=float).reshape(DIM, DIM)))
+    return itk_image
+
+
+def _itk_field_to_sitk_transform(field: "itk.Image", reference: sitk.Image) -> sitk.Transform:
+    """Wrap an itk displacement field (on the fixed grid) as a SimpleITK ``DisplacementFieldTransform``."""
+    array = itk.array_from_image(field).astype(np.float64)  # [Z, Y, X, 3]
+    sitk_field = sitk.GetImageFromArray(array, isVector=True)
+    sitk_field.CopyInformation(reference)
+    return sitk.DisplacementFieldTransform(sitk.Cast(sitk_field, sitk.sitkVectorFloat64))
+
+
+def _itk_affine_to_sitk(affine: "itk.AffineTransform") -> sitk.AffineTransform:
+    """Convert an ``itk.AffineTransform[D, 3]`` into a SimpleITK ``AffineTransform`` (same LPS convention)."""
+    sitk_affine = sitk.AffineTransform(DIM)
+    sitk_affine.SetMatrix([float(v) for v in itk.array_from_matrix(affine.GetMatrix()).flatten()])
+    sitk_affine.SetTranslation([float(v) for v in affine.GetTranslation()])
+    sitk_affine.SetCenter([float(v) for v in affine.GetCenter()])
+    return sitk_affine
+
+
+class ConvexAdamEngine:
+    """Register a fixed/moving pair with the itk-impact ConvexAdam pipeline; return (moved, dvf) on the fixed grid.
+
+    The IMPACT feature models are downloaded once (``repo:filename`` on Hugging Face) and reused across cases.
+    Masks are accepted for signature compatibility with the elastix engine but ignored: the ConvexAdam
+    filters optimise over the whole image (no mask API is exposed by the coarse/fine stages).
+    """
+
+    def __init__(
+        self,
+        models: list[str],
+        voxel_size: list[float],
+        num_channels: int,
+        overlap: int,
+        layers_mask: list[bool],
+        mixed_precision: bool,
+        grid_spacing: int,
+        displacement_half_width: int,
+        iterations: int,
+        learning_rate: float,
+        regularization_weight: float,
+        grid_shrink: int,
+        distance: list[str],
+        layers_weight: list[float],
+        subset_features: list[int],
+        pca: list[int],
+        stages: list[str],
+        linear: bool,
+        linear_iterations: int,
+        seed: int,
+    ) -> None:
+        self._stages = stages
+        self._model_paths = self._download_models(models)
+        # Built lazily and cached: constructing an itk.ModelConfiguration loads the TorchScript model
+        # from disk in C++, so build the list once and reuse it across both stages and every case.
+        self._configurations: "list[itk.ModelConfiguration] | None" = None
+        self._voxel_size = voxel_size
+        self._num_channels = num_channels
+        self._overlap = overlap
+        self._layers_mask = layers_mask
+        self._mixed_precision = mixed_precision
+        self._grid_spacing = grid_spacing
+        self._displacement_half_width = displacement_half_width
+        self._iterations = iterations
+        self._learning_rate = learning_rate
+        self._regularization_weight = regularization_weight
+        self._grid_shrink = grid_shrink
+        self._distance = distance
+        self._layers_weight = layers_weight
+        self._subset_features = subset_features
+        self._pca = pca
+        self._linear = linear
+        self._linear_iterations = linear_iterations
+        self._seed = seed
+
+    @staticmethod
+    def _download_models(models: list[str]) -> list[str]:
+        """Fetch the TorchScript feature models (``repo:filename``); return their local paths."""
+        paths = []
+        for ref in models:
+            repo, filename = ref.split(":", 1)
+            paths.append(str(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")))  # nosec B615
+        return paths
+
+    def _model_configurations(self) -> list["itk.ModelConfiguration"]:
+        """Build one ``ModelConfiguration`` per feature model once, then reuse it across stages and cases.
+
+        Constructing an ``itk.ModelConfiguration`` loads the TorchScript module from disk on the C++ side, so
+        it is built lazily and cached. The coarse/fine filters copy each configuration by value in
+        ``AddModelConfiguration`` and the copy shares the loaded module through the configuration's internal
+        ``shared_ptr`` — so a single build is reused everywhere without any reload.
+        """
+        if self._configurations is None:
+            self._configurations = [
+                itk.ModelConfiguration(
+                    path,
+                    DIM,
+                    self._num_channels,
+                    [0, 0, 0],
+                    [float(v) for v in self._voxel_size],
+                    self._overlap,
+                    list(self._layers_mask),
+                    self._mixed_precision,
+                )
+                for path in self._model_paths
+            ]
+        return self._configurations
+
+    def _linear_align(self, fixed: "itk.Image", moving: "itk.Image") -> "itk.AffineTransform":
+        """Moments-initialised rigid + affine (Mattes MI), mapping fixed -> moving physical points."""
+        rigid = itk.VersorRigid3DTransform[itk.D].New()
+        initializer = itk.CenteredTransformInitializer[itk.VersorRigid3DTransform[itk.D], _IMAGE_F, _IMAGE_F].New(
+            Transform=rigid, FixedImage=fixed, MovingImage=moving
+        )
+        initializer.MomentsOn()
+        initializer.InitializeTransform()
+
+        affine = itk.AffineTransform[itk.D, DIM].New()
+        affine.SetCenter(rigid.GetCenter())
+        affine.SetMatrix(rigid.GetMatrix())
+        affine.SetOffset(rigid.GetOffset())
+        levels = 3
+        metric_type = itk.MattesMutualInformationImageToImageMetricv4[_IMAGE_F, _IMAGE_F]
+        metric = metric_type.New()
+        metric.SetNumberOfHistogramBins(32)
+        optimizer = itk.RegularStepGradientDescentOptimizerv4[itk.D].New()
+        optimizer.SetNumberOfIterations(self._linear_iterations)
+        optimizer.SetLearningRate(1.0)
+        optimizer.SetMinimumStepLength(1e-5)
+        optimizer.SetRelaxationFactor(0.6)
+        scales = itk.RegistrationParameterScalesFromPhysicalShift[metric_type].New()
+        scales.SetMetric(metric)
+        optimizer.SetScalesEstimator(scales)
+        registration = itk.ImageRegistrationMethodv4[_IMAGE_F, _IMAGE_F].New(
+            FixedImage=fixed, MovingImage=moving, Metric=metric, Optimizer=optimizer, InitialTransform=affine
+        )
+        registration.SetNumberOfLevels(levels)
+        registration.SetShrinkFactorsPerLevel([2 ** (levels - 1 - i) for i in range(levels)])
+        registration.SetSmoothingSigmasPerLevel([float(levels - 1 - i) for i in range(levels)])
+        registration.InPlaceOn()
+        registration.Update()
+        return affine
+
+    def _coarse(self, fixed: "itk.Image", moving: "itk.Image", device: str) -> "itk.Image":
+        """ConvexAdam coarse coupled-convex initializer -> robust low-resolution field on the fixed grid."""
+        coarse = _coarse_registration_type().New()
+        coarse.SetFixedImage(fixed)
+        coarse.SetMovingImage(moving)
+        for configuration in self._model_configurations():
+            coarse.AddModelConfiguration(configuration)
+        coarse.SetGridSpacing(self._grid_spacing)
+        coarse.SetDisplacementHalfWidth(self._displacement_half_width)
+        coarse.SetDevice(device)
+        coarse.SetSeed(self._seed)
+        coarse.Update()
+        field = coarse.GetOutput()
+        field.DisconnectPipeline()
+        return field
+
+    def _fine(
+        self, fixed: "itk.Image", moving: "itk.Image", initial_field: "itk.Image | None", device: str
+    ) -> "itk.Image":
+        """Adam instance-optimisation refinement, warm-started from ``initial_field`` (zero if none)."""
+        fine = _fine_registration_type().New()
+        fine.SetFixedImage(fixed)
+        fine.SetMovingImage(moving)
+        fine.SetInitialDisplacementField(initial_field if initial_field is not None else self._zero_field(fixed))
+        for configuration in self._model_configurations():
+            fine.AddModelConfiguration(configuration)
+        fine.SetDistance(list(self._distance))
+        fine.SetLayersWeight([float(v) for v in self._layers_weight])
+        fine.SetSubsetFeatures([int(v) for v in self._subset_features])
+        fine.SetPCA([int(v) for v in self._pca])
+        fine.SetNumberOfIterations(self._iterations)
+        fine.SetLearningRate(self._learning_rate)
+        fine.SetRegularizationWeight(self._regularization_weight)
+        fine.SetGridShrinkFactor(self._grid_shrink)
+        fine.SetDevice(device)
+        fine.SetSeed(self._seed)
+
+        # Mirror KonfAI's informative bars: drive a tqdm over the Adam iterations from the metric trace so
+        # SlicerKonfAI (which parses the "N% done/total" progress line) shows real progress. The observer is
+        # best-effort — if the filter does not emit IterationEvent the bar simply fills on completion.
+        progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+
+        def _update(*_: object) -> None:
+            values = list(fine.GetMetricValuesPerIteration())
+            progress.n = min(len(values), self._iterations)
+            if values:
+                progress.set_description(f"Registration : iter {len(values)} | metric {float(values[-1]):.4f}")
+            progress.refresh()
+
+        try:
+            fine.AddObserver(itk.IterationEvent(), _update)
+        except Exception:  # nosec B110 - progress is best-effort; never fail a run over the bar
+            pass
+        fine.Update()
+        progress.close()
+        field = fine.GetDisplacementField()
+        field.DisconnectPipeline()
+        return field
+
+    @staticmethod
+    def _zero_field(reference: "itk.Image") -> "itk.Image":
+        """An all-zero displacement field on ``reference``'s grid (identity warm-start for a lone fine stage)."""
+        field = itk.Image[itk.Vector[itk.F, DIM], DIM].New()
+        field.CopyInformation(reference)
+        field.SetRegions(reference.GetLargestPossibleRegion())
+        field.Allocate()
+        zero = itk.Vector[itk.F, DIM]()
+        zero.Fill(0)  # itk::Vector default ctor does not zero-initialise
+        field.FillBuffer(zero)
+        return field
+
+    def _run_stages(self, fixed: "itk.Image", moving: "itk.Image", device: str) -> "itk.Image | None":
+        """Run the configured coarse/fine chain; each fine warm-starts from the running field.
+
+        ``coarse`` produces a field from scratch; ``fine`` refines the running field. So ``['coarse']`` is a
+        coarse-only app, ``['fine']`` a fine-only app (zero warm-start), and ``['coarse', 'fine']`` chains both
+        (the composite, as before). Returns None when no deformable stage runs (e.g. a linear-only chain).
+        """
+        field: "itk.Image | None" = None
+        for stage in self._stages:
+            if stage == "coarse":
+                field = self._coarse(fixed, moving, device)
+            elif stage == "fine":
+                field = self._fine(fixed, moving, field, device)
+            else:
+                raise ValueError(f"Unknown registration stage '{stage}' (expected 'coarse' or 'fine').")
+        return field
+
+    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."""
+        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
+        fixed_itk = _sitk_to_itk(fixed)
+        moving_itk = _sitk_to_itk(moving)
+
+        # Optional linear pre-align: resample the moving onto the fixed grid so the deformable stage starts close.
+        affine = self._linear_align(fixed_itk, moving_itk) if self._linear else itk.AffineTransform[itk.D, DIM].New()
+        resampler = itk.ResampleImageFilter[_IMAGE_F, _IMAGE_F].New(
+            Input=moving_itk, ReferenceImage=fixed_itk, Transform=affine
+        )
+        resampler.UseReferenceImageOn()
+        resampler.SetInterpolator(itk.LinearInterpolateImageFunction[_IMAGE_F, itk.D].New())
+        resampler.Update()
+        moving_linear = resampler.GetOutput()
+
+        field = self._run_stages(fixed_itk, moving_linear, device)
+
+        # One transform on the fixed grid = affine then deformable, so the returned DVF/transform warps the
+        # ORIGINAL moving. SimpleITK applies the last-added transform first, so [affine, deformable] gives
+        # moved(p) = moving(affine(deformable(p))). A linear-only chain (field is None) yields the affine alone.
+        chain = [_itk_affine_to_sitk(affine)]
+        if field is not None:
+            chain.append(_itk_field_to_sitk_transform(field, fixed))
+        composite = sitk.CompositeTransform(chain)
+        moved = sitk.Resample(moving, fixed, composite, sitk.sitkLinear, 0.0, moving.GetPixelID())
+        dvf = sitk.TransformToDisplacementField(
+            composite,
+            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
+
+
+class ConvexAdamRegistration(torch.nn.Module):
+    """Graph module: (fixed, moving) tensors + their geometry -> moved image + DVF on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving the per-branch ``Attribute`` list alongside
+    the tensors (same convention as ``CriterionWithAttribute``); registration needs the physical geometry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: ConvexAdamEngine) -> None:
+        super().__init__()
+        self._engine = engine
+
+    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 a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field (DIM
+        # channels); downstream ChannelSelect modules split them. Masks are ignored by the ConvexAdam engine.
+        fixed_attrs, moving_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])
+            moved_np, dvf_np = self._engine.register(fixed_img, moving_img, device_index)
+            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 ConvexAdam registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1;
+    the mask branches 2/3 are accepted but unused by this engine).
+
+    Outputs on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField`` (the
+    DIM-component displacement field, in mm). 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()},
+        models: list[str] = [],
+        voxel_size: list[float] = [3.0, 3.0, 3.0],
+        num_channels: int = 1,
+        overlap: int = 2,
+        layers_mask: list[bool] = [True],
+        mixed_precision: bool = False,
+        grid_spacing: int = 4,
+        displacement_half_width: int = 6,
+        iterations: int = 150,
+        learning_rate: float = 0.2,
+        regularization_weight: float = 1.0,
+        grid_shrink: int = 4,
+        distance: list[str] = ["L1"],
+        layers_weight: list[float] = [1.0],
+        subset_features: list[int] = [32],
+        pca: list[int] = [0],
+        stages: list[str] = ["coarse", "fine"],
+        linear: bool = True,
+        linear_iterations: int = 200,
+        seed: int = 42,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        engine = ConvexAdamEngine(
+            models,
+            voxel_size,
+            num_channels,
+            overlap,
+            layers_mask,
+            mixed_precision,
+            grid_spacing,
+            displacement_half_width,
+            iterations,
+            learning_rate,
+            regularization_weight,
+            grid_shrink,
+            distance,
+            layers_weight,
+            subset_features,
+            pca,
+            stages,
+            linear,
+            linear_iterations,
+            seed,
+        )
+        self.add_module(
+            "Registration", ConvexAdamRegistration(engine), 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"])

ConvexAdam_Coarse/Prediction.yml

@@ -0,0 +1,140 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      models:
+      - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt
+      voxel_size:
+      - 3.0
+      - 3.0
+      - 3.0
+      num_channels: 1
+      overlap: 2
+      layers_mask:
+      - true
+      mixed_precision: false
+      grid_spacing: 4
+      displacement_half_width: 6
+      iterations: 150
+      learning_rate: 0.2
+      regularization_weight: 1.0
+      grid_shrink: 4
+      distance:
+      - L1
+      layers_weight:
+      - 1.0
+      subset_features:
+      - 32
+      pca:
+      - 0
+      linear: true
+      linear_iterations: 200
+      seed: 42
+      outputs_criterions: None
+      stages:
+      - coarse
+  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:
+      DataAugmentation_0:
+        data_augmentations:
+          Flip:
+            f_prob:
+            - 0
+            - 0.5
+            - 0.5
+            vector_field: true
+            prob: 1
+        nb: 2
+    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-ConvexAdam-Coarse
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

ConvexAdam_Coarse/Uncertainty.yml

@@ -0,0 +1,24 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean:
+              name: Uncertainty
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Norm: {}
+              StandardDeviation: {}
+              Save:
+                dataset: ./Uncertainties/ImpactReg/Output:mha
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

ConvexAdam_Coarse/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "ConvexAdam Coarse (MIND)",
+    "short_description": "Global coarse ConvexAdam initialization (whole-volume, IMPACT/MIND features).",
+    "description": "First stage of the ConvexAdam pipeline: an optional moments+affine linear pre-align followed by the ConvexAdam coarse coupled-convex initialization on the whole volume, driven by the IMPACT metric on MIND features. Produces a robust low-resolution displacement field on the fixed grid, meant to warm-start (and pre-resample the moving for) the fine stage.",
+    "task": "registration",
+    "tta": 3,
+    "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
+                }
+            }
+        }
+    }
+}

ConvexAdam_Coarse/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de99fbc36331ce674639acc774f52b4a2d0027f2f312d9d28669e831a0c4fd7e
+size 1249

ConvexAdam_Coarse/requirements.txt

@@ -0,0 +1 @@
+itk-impact

ConvexAdam_Composite/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

ConvexAdam_Composite/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

ConvexAdam_Composite/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

ConvexAdam_Composite/Model.py

@@ -0,0 +1,471 @@
+# 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
+
+"""ConvexAdam (itk-impact) registration as a self-contained KonfAI model.
+
+Same idiomatic ``add_module`` graph and the same output contract as the elastix preset
+(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``),
+so the orchestrator / app.json / ensemble / uncertainty are unchanged. The engine here is
+the native, in-memory itk-impact ConvexAdam pipeline (``pip install itk-impact``) instead of
+the elastix binary:
+
+    (optional) moments + affine Mattes-MI          [ITKv4 linear pre-align]
+      -> ImpactCoarseRegistration                   [coupled-convex init, IMPACT features]
+      -> ImpactFineRegistration                     [Adam instance optimisation, IMPACT features]
+
+The IMPACT feature models (e.g. MIND) are TorchScript ``.pt`` files fetched from Hugging Face
+and wrapped as ``itk.ModelConfiguration`` — the same models the elastix presets use.
+
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
+"""
+
+import itk
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+DIM = 3
+_IMAGE_F = itk.Image[itk.F, DIM]
+
+
+def _coarse_registration_type():
+    """The coupled-convex initializer, tolerant to the two names the wrapping has shipped under."""
+    cls = getattr(itk, "ImpactCoarseRegistration", None) or getattr(itk, "ImpactConvexAdamInitializer", None)
+    if cls is None:
+        raise RuntimeError(
+            "itk-impact does not expose ImpactCoarseRegistration / ImpactConvexAdamInitializer; "
+            "install a build with the ConvexAdam registration filters."
+        )
+    return cls[_IMAGE_F, _IMAGE_F]
+
+
+def _fine_registration_type():
+    """The Adam instance-optimisation stage, tolerant to the two names the wrapping has shipped under."""
+    cls = getattr(itk, "ImpactFineRegistration", None) or getattr(itk, "ImpactTorchAdamRegistration", None)
+    if cls is None:
+        raise RuntimeError(
+            "itk-impact does not expose ImpactFineRegistration / ImpactTorchAdamRegistration; "
+            "install a build with the ConvexAdam registration filters."
+        )
+    return cls[_IMAGE_F, _IMAGE_F]
+
+
+def _sitk_to_itk(image: sitk.Image) -> "itk.Image":
+    """Copy a scalar SimpleITK image (with its geometry) into an ``itk.Image[F, 3]``."""
+    itk_image = itk.image_from_array(sitk.GetArrayFromImage(image).astype(np.float32))
+    itk_image.SetOrigin([float(v) for v in image.GetOrigin()])
+    itk_image.SetSpacing([float(v) for v in image.GetSpacing()])
+    itk_image.SetDirection(itk.matrix_from_array(np.asarray(image.GetDirection(), dtype=float).reshape(DIM, DIM)))
+    return itk_image
+
+
+def _itk_field_to_sitk_transform(field: "itk.Image", reference: sitk.Image) -> sitk.Transform:
+    """Wrap an itk displacement field (on the fixed grid) as a SimpleITK ``DisplacementFieldTransform``."""
+    array = itk.array_from_image(field).astype(np.float64)  # [Z, Y, X, 3]
+    sitk_field = sitk.GetImageFromArray(array, isVector=True)
+    sitk_field.CopyInformation(reference)
+    return sitk.DisplacementFieldTransform(sitk.Cast(sitk_field, sitk.sitkVectorFloat64))
+
+
+def _itk_affine_to_sitk(affine: "itk.AffineTransform") -> sitk.AffineTransform:
+    """Convert an ``itk.AffineTransform[D, 3]`` into a SimpleITK ``AffineTransform`` (same LPS convention)."""
+    sitk_affine = sitk.AffineTransform(DIM)
+    sitk_affine.SetMatrix([float(v) for v in itk.array_from_matrix(affine.GetMatrix()).flatten()])
+    sitk_affine.SetTranslation([float(v) for v in affine.GetTranslation()])
+    sitk_affine.SetCenter([float(v) for v in affine.GetCenter()])
+    return sitk_affine
+
+
+class ConvexAdamEngine:
+    """Register a fixed/moving pair with the itk-impact ConvexAdam pipeline; return (moved, dvf) on the fixed grid.
+
+    The IMPACT feature models are downloaded once (``repo:filename`` on Hugging Face) and reused across cases.
+    Masks are accepted for signature compatibility with the elastix engine but ignored: the ConvexAdam
+    filters optimise over the whole image (no mask API is exposed by the coarse/fine stages).
+    """
+
+    def __init__(
+        self,
+        models: list[str],
+        voxel_size: list[float],
+        num_channels: int,
+        overlap: int,
+        layers_mask: list[bool],
+        mixed_precision: bool,
+        grid_spacing: int,
+        displacement_half_width: int,
+        iterations: int,
+        learning_rate: float,
+        regularization_weight: float,
+        grid_shrink: int,
+        distance: list[str],
+        layers_weight: list[float],
+        subset_features: list[int],
+        pca: list[int],
+        stages: list[str],
+        linear: bool,
+        linear_iterations: int,
+        seed: int,
+    ) -> None:
+        self._stages = stages
+        self._model_paths = self._download_models(models)
+        # Built lazily and cached: constructing an itk.ModelConfiguration loads the TorchScript model
+        # from disk in C++, so build the list once and reuse it across both stages and every case.
+        self._configurations: "list[itk.ModelConfiguration] | None" = None
+        self._voxel_size = voxel_size
+        self._num_channels = num_channels
+        self._overlap = overlap
+        self._layers_mask = layers_mask
+        self._mixed_precision = mixed_precision
+        self._grid_spacing = grid_spacing
+        self._displacement_half_width = displacement_half_width
+        self._iterations = iterations
+        self._learning_rate = learning_rate
+        self._regularization_weight = regularization_weight
+        self._grid_shrink = grid_shrink
+        self._distance = distance
+        self._layers_weight = layers_weight
+        self._subset_features = subset_features
+        self._pca = pca
+        self._linear = linear
+        self._linear_iterations = linear_iterations
+        self._seed = seed
+
+    @staticmethod
+    def _download_models(models: list[str]) -> list[str]:
+        """Fetch the TorchScript feature models (``repo:filename``); return their local paths."""
+        paths = []
+        for ref in models:
+            repo, filename = ref.split(":", 1)
+            paths.append(str(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")))  # nosec B615
+        return paths
+
+    def _model_configurations(self) -> list["itk.ModelConfiguration"]:
+        """Build one ``ModelConfiguration`` per feature model once, then reuse it across stages and cases.
+
+        Constructing an ``itk.ModelConfiguration`` loads the TorchScript module from disk on the C++ side, so
+        it is built lazily and cached. The coarse/fine filters copy each configuration by value in
+        ``AddModelConfiguration`` and the copy shares the loaded module through the configuration's internal
+        ``shared_ptr`` — so a single build is reused everywhere without any reload.
+        """
+        if self._configurations is None:
+            self._configurations = [
+                itk.ModelConfiguration(
+                    path,
+                    DIM,
+                    self._num_channels,
+                    [0, 0, 0],
+                    [float(v) for v in self._voxel_size],
+                    self._overlap,
+                    list(self._layers_mask),
+                    self._mixed_precision,
+                )
+                for path in self._model_paths
+            ]
+        return self._configurations
+
+    def _linear_align(self, fixed: "itk.Image", moving: "itk.Image") -> "itk.AffineTransform":
+        """Moments-initialised rigid + affine (Mattes MI), mapping fixed -> moving physical points."""
+        rigid = itk.VersorRigid3DTransform[itk.D].New()
+        initializer = itk.CenteredTransformInitializer[itk.VersorRigid3DTransform[itk.D], _IMAGE_F, _IMAGE_F].New(
+            Transform=rigid, FixedImage=fixed, MovingImage=moving
+        )
+        initializer.MomentsOn()
+        initializer.InitializeTransform()
+
+        affine = itk.AffineTransform[itk.D, DIM].New()
+        affine.SetCenter(rigid.GetCenter())
+        affine.SetMatrix(rigid.GetMatrix())
+        affine.SetOffset(rigid.GetOffset())
+        levels = 3
+        metric_type = itk.MattesMutualInformationImageToImageMetricv4[_IMAGE_F, _IMAGE_F]
+        metric = metric_type.New()
+        metric.SetNumberOfHistogramBins(32)
+        optimizer = itk.RegularStepGradientDescentOptimizerv4[itk.D].New()
+        optimizer.SetNumberOfIterations(self._linear_iterations)
+        optimizer.SetLearningRate(1.0)
+        optimizer.SetMinimumStepLength(1e-5)
+        optimizer.SetRelaxationFactor(0.6)
+        scales = itk.RegistrationParameterScalesFromPhysicalShift[metric_type].New()
+        scales.SetMetric(metric)
+        optimizer.SetScalesEstimator(scales)
+        registration = itk.ImageRegistrationMethodv4[_IMAGE_F, _IMAGE_F].New(
+            FixedImage=fixed, MovingImage=moving, Metric=metric, Optimizer=optimizer, InitialTransform=affine
+        )
+        registration.SetNumberOfLevels(levels)
+        registration.SetShrinkFactorsPerLevel([2 ** (levels - 1 - i) for i in range(levels)])
+        registration.SetSmoothingSigmasPerLevel([float(levels - 1 - i) for i in range(levels)])
+        registration.InPlaceOn()
+        registration.Update()
+        return affine
+
+    def _coarse(self, fixed: "itk.Image", moving: "itk.Image", device: str) -> "itk.Image":
+        """ConvexAdam coarse coupled-convex initializer -> robust low-resolution field on the fixed grid."""
+        coarse = _coarse_registration_type().New()
+        coarse.SetFixedImage(fixed)
+        coarse.SetMovingImage(moving)
+        for configuration in self._model_configurations():
+            coarse.AddModelConfiguration(configuration)
+        coarse.SetGridSpacing(self._grid_spacing)
+        coarse.SetDisplacementHalfWidth(self._displacement_half_width)
+        coarse.SetDevice(device)
+        coarse.SetSeed(self._seed)
+        coarse.Update()
+        field = coarse.GetOutput()
+        field.DisconnectPipeline()
+        return field
+
+    def _fine(
+        self, fixed: "itk.Image", moving: "itk.Image", initial_field: "itk.Image | None", device: str
+    ) -> "itk.Image":
+        """Adam instance-optimisation refinement, warm-started from ``initial_field`` (zero if none)."""
+        fine = _fine_registration_type().New()
+        fine.SetFixedImage(fixed)
+        fine.SetMovingImage(moving)
+        fine.SetInitialDisplacementField(initial_field if initial_field is not None else self._zero_field(fixed))
+        for configuration in self._model_configurations():
+            fine.AddModelConfiguration(configuration)
+        fine.SetDistance(list(self._distance))
+        fine.SetLayersWeight([float(v) for v in self._layers_weight])
+        fine.SetSubsetFeatures([int(v) for v in self._subset_features])
+        fine.SetPCA([int(v) for v in self._pca])
+        fine.SetNumberOfIterations(self._iterations)
+        fine.SetLearningRate(self._learning_rate)
+        fine.SetRegularizationWeight(self._regularization_weight)
+        fine.SetGridShrinkFactor(self._grid_shrink)
+        fine.SetDevice(device)
+        fine.SetSeed(self._seed)
+
+        # Mirror KonfAI's informative bars: drive a tqdm over the Adam iterations from the metric trace so
+        # SlicerKonfAI (which parses the "N% done/total" progress line) shows real progress. The observer is
+        # best-effort — if the filter does not emit IterationEvent the bar simply fills on completion.
+        progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+
+        def _update(*_: object) -> None:
+            values = list(fine.GetMetricValuesPerIteration())
+            progress.n = min(len(values), self._iterations)
+            if values:
+                progress.set_description(f"Registration : iter {len(values)} | metric {float(values[-1]):.4f}")
+            progress.refresh()
+
+        try:
+            fine.AddObserver(itk.IterationEvent(), _update)
+        except Exception:  # nosec B110 - progress is best-effort; never fail a run over the bar
+            pass
+        fine.Update()
+        progress.close()
+        field = fine.GetDisplacementField()
+        field.DisconnectPipeline()
+        return field
+
+    @staticmethod
+    def _zero_field(reference: "itk.Image") -> "itk.Image":
+        """An all-zero displacement field on ``reference``'s grid (identity warm-start for a lone fine stage)."""
+        field = itk.Image[itk.Vector[itk.F, DIM], DIM].New()
+        field.CopyInformation(reference)
+        field.SetRegions(reference.GetLargestPossibleRegion())
+        field.Allocate()
+        zero = itk.Vector[itk.F, DIM]()
+        zero.Fill(0)  # itk::Vector default ctor does not zero-initialise
+        field.FillBuffer(zero)
+        return field
+
+    def _run_stages(self, fixed: "itk.Image", moving: "itk.Image", device: str) -> "itk.Image | None":
+        """Run the configured coarse/fine chain; each fine warm-starts from the running field.
+
+        ``coarse`` produces a field from scratch; ``fine`` refines the running field. So ``['coarse']`` is a
+        coarse-only app, ``['fine']`` a fine-only app (zero warm-start), and ``['coarse', 'fine']`` chains both
+        (the composite, as before). Returns None when no deformable stage runs (e.g. a linear-only chain).
+        """
+        field: "itk.Image | None" = None
+        for stage in self._stages:
+            if stage == "coarse":
+                field = self._coarse(fixed, moving, device)
+            elif stage == "fine":
+                field = self._fine(fixed, moving, field, device)
+            else:
+                raise ValueError(f"Unknown registration stage '{stage}' (expected 'coarse' or 'fine').")
+        return field
+
+    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."""
+        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
+        fixed_itk = _sitk_to_itk(fixed)
+        moving_itk = _sitk_to_itk(moving)
+
+        # Optional linear pre-align: resample the moving onto the fixed grid so the deformable stage starts close.
+        affine = self._linear_align(fixed_itk, moving_itk) if self._linear else itk.AffineTransform[itk.D, DIM].New()
+        resampler = itk.ResampleImageFilter[_IMAGE_F, _IMAGE_F].New(
+            Input=moving_itk, ReferenceImage=fixed_itk, Transform=affine
+        )
+        resampler.UseReferenceImageOn()
+        resampler.SetInterpolator(itk.LinearInterpolateImageFunction[_IMAGE_F, itk.D].New())
+        resampler.Update()
+        moving_linear = resampler.GetOutput()
+
+        field = self._run_stages(fixed_itk, moving_linear, device)
+
+        # One transform on the fixed grid = affine then deformable, so the returned DVF/transform warps the
+        # ORIGINAL moving. SimpleITK applies the last-added transform first, so [affine, deformable] gives
+        # moved(p) = moving(affine(deformable(p))). A linear-only chain (field is None) yields the affine alone.
+        chain = [_itk_affine_to_sitk(affine)]
+        if field is not None:
+            chain.append(_itk_field_to_sitk_transform(field, fixed))
+        composite = sitk.CompositeTransform(chain)
+        moved = sitk.Resample(moving, fixed, composite, sitk.sitkLinear, 0.0, moving.GetPixelID())
+        dvf = sitk.TransformToDisplacementField(
+            composite,
+            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
+
+
+class ConvexAdamRegistration(torch.nn.Module):
+    """Graph module: (fixed, moving) tensors + their geometry -> moved image + DVF on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving the per-branch ``Attribute`` list alongside
+    the tensors (same convention as ``CriterionWithAttribute``); registration needs the physical geometry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: ConvexAdamEngine) -> None:
+        super().__init__()
+        self._engine = engine
+
+    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 a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field (DIM
+        # channels); downstream ChannelSelect modules split them. Masks are ignored by the ConvexAdam engine.
+        fixed_attrs, moving_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])
+            moved_np, dvf_np = self._engine.register(fixed_img, moving_img, device_index)
+            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 ConvexAdam registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1;
+    the mask branches 2/3 are accepted but unused by this engine).
+
+    Outputs on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField`` (the
+    DIM-component displacement field, in mm). 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()},
+        models: list[str] = [],
+        voxel_size: list[float] = [3.0, 3.0, 3.0],
+        num_channels: int = 1,
+        overlap: int = 2,
+        layers_mask: list[bool] = [True],
+        mixed_precision: bool = False,
+        grid_spacing: int = 4,
+        displacement_half_width: int = 6,
+        iterations: int = 150,
+        learning_rate: float = 0.2,
+        regularization_weight: float = 1.0,
+        grid_shrink: int = 4,
+        distance: list[str] = ["L1"],
+        layers_weight: list[float] = [1.0],
+        subset_features: list[int] = [32],
+        pca: list[int] = [0],
+        stages: list[str] = ["coarse", "fine"],
+        linear: bool = True,
+        linear_iterations: int = 200,
+        seed: int = 42,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        engine = ConvexAdamEngine(
+            models,
+            voxel_size,
+            num_channels,
+            overlap,
+            layers_mask,
+            mixed_precision,
+            grid_spacing,
+            displacement_half_width,
+            iterations,
+            learning_rate,
+            regularization_weight,
+            grid_shrink,
+            distance,
+            layers_weight,
+            subset_features,
+            pca,
+            stages,
+            linear,
+            linear_iterations,
+            seed,
+        )
+        self.add_module(
+            "Registration", ConvexAdamRegistration(engine), 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"])

ConvexAdam_Composite/Prediction.yml

@@ -0,0 +1,141 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      models:
+      - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt
+      voxel_size:
+      - 3.0
+      - 3.0
+      - 3.0
+      num_channels: 1
+      overlap: 2
+      layers_mask:
+      - true
+      mixed_precision: false
+      grid_spacing: 4
+      displacement_half_width: 6
+      iterations: 150
+      learning_rate: 0.2
+      regularization_weight: 1.0
+      grid_shrink: 4
+      distance:
+      - L1
+      layers_weight:
+      - 1.0
+      subset_features:
+      - 32
+      pca:
+      - 0
+      linear: true
+      linear_iterations: 200
+      seed: 42
+      outputs_criterions: None
+      stages:
+      - coarse
+      - fine
+  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:
+      DataAugmentation_0:
+        data_augmentations:
+          Flip:
+            f_prob:
+            - 0
+            - 0.5
+            - 0.5
+            vector_field: true
+            prob: 1
+        nb: 2
+    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-ConvexAdam-Composite
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

ConvexAdam_Composite/Uncertainty.yml

@@ -0,0 +1,24 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean:
+              name: Uncertainty
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Norm: {}
+              StandardDeviation: {}
+              Save:
+                dataset: ./Uncertainties/ImpactReg/Output:mha
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

ConvexAdam_Composite/app.json

@@ -0,0 +1,94 @@
+{
+    "display_name": "ConvexAdam (MIND)",
+    "short_description": "Coarse + fine ConvexAdam registration in one app (IMPACT/MIND features).",
+    "description": "The full ConvexAdam pipeline chained in one app: optional moments+affine linear pre-align, then a ConvexAdam coarse coupled-convex initialization, then an Adam instance-optimisation refinement, all driven by the IMPACT metric on MIND features. Produces the moved image and the displacement field on the fixed grid.",
+    "task": "registration",
+    "tta": 3,
+    "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
+                }
+            }
+        }
+    }
+}

ConvexAdam_Composite/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de99fbc36331ce674639acc774f52b4a2d0027f2f312d9d28669e831a0c4fd7e
+size 1249

ConvexAdam_Composite/requirements.txt

@@ -0,0 +1 @@
+itk-impact

ConvexAdam_Fine/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

ConvexAdam_Fine/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

ConvexAdam_Fine/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

ConvexAdam_Fine/Model.py

@@ -0,0 +1,471 @@
+# 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
+
+"""ConvexAdam (itk-impact) registration as a self-contained KonfAI model.
+
+Same idiomatic ``add_module`` graph and the same output contract as the elastix preset
+(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``),
+so the orchestrator / app.json / ensemble / uncertainty are unchanged. The engine here is
+the native, in-memory itk-impact ConvexAdam pipeline (``pip install itk-impact``) instead of
+the elastix binary:
+
+    (optional) moments + affine Mattes-MI          [ITKv4 linear pre-align]
+      -> ImpactCoarseRegistration                   [coupled-convex init, IMPACT features]
+      -> ImpactFineRegistration                     [Adam instance optimisation, IMPACT features]
+
+The IMPACT feature models (e.g. MIND) are TorchScript ``.pt`` files fetched from Hugging Face
+and wrapped as ``itk.ModelConfiguration`` — the same models the elastix presets use.
+
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
+"""
+
+import itk
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from konfai.network import network
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+DIM = 3
+_IMAGE_F = itk.Image[itk.F, DIM]
+
+
+def _coarse_registration_type():
+    """The coupled-convex initializer, tolerant to the two names the wrapping has shipped under."""
+    cls = getattr(itk, "ImpactCoarseRegistration", None) or getattr(itk, "ImpactConvexAdamInitializer", None)
+    if cls is None:
+        raise RuntimeError(
+            "itk-impact does not expose ImpactCoarseRegistration / ImpactConvexAdamInitializer; "
+            "install a build with the ConvexAdam registration filters."
+        )
+    return cls[_IMAGE_F, _IMAGE_F]
+
+
+def _fine_registration_type():
+    """The Adam instance-optimisation stage, tolerant to the two names the wrapping has shipped under."""
+    cls = getattr(itk, "ImpactFineRegistration", None) or getattr(itk, "ImpactTorchAdamRegistration", None)
+    if cls is None:
+        raise RuntimeError(
+            "itk-impact does not expose ImpactFineRegistration / ImpactTorchAdamRegistration; "
+            "install a build with the ConvexAdam registration filters."
+        )
+    return cls[_IMAGE_F, _IMAGE_F]
+
+
+def _sitk_to_itk(image: sitk.Image) -> "itk.Image":
+    """Copy a scalar SimpleITK image (with its geometry) into an ``itk.Image[F, 3]``."""
+    itk_image = itk.image_from_array(sitk.GetArrayFromImage(image).astype(np.float32))
+    itk_image.SetOrigin([float(v) for v in image.GetOrigin()])
+    itk_image.SetSpacing([float(v) for v in image.GetSpacing()])
+    itk_image.SetDirection(itk.matrix_from_array(np.asarray(image.GetDirection(), dtype=float).reshape(DIM, DIM)))
+    return itk_image
+
+
+def _itk_field_to_sitk_transform(field: "itk.Image", reference: sitk.Image) -> sitk.Transform:
+    """Wrap an itk displacement field (on the fixed grid) as a SimpleITK ``DisplacementFieldTransform``."""
+    array = itk.array_from_image(field).astype(np.float64)  # [Z, Y, X, 3]
+    sitk_field = sitk.GetImageFromArray(array, isVector=True)
+    sitk_field.CopyInformation(reference)
+    return sitk.DisplacementFieldTransform(sitk.Cast(sitk_field, sitk.sitkVectorFloat64))
+
+
+def _itk_affine_to_sitk(affine: "itk.AffineTransform") -> sitk.AffineTransform:
+    """Convert an ``itk.AffineTransform[D, 3]`` into a SimpleITK ``AffineTransform`` (same LPS convention)."""
+    sitk_affine = sitk.AffineTransform(DIM)
+    sitk_affine.SetMatrix([float(v) for v in itk.array_from_matrix(affine.GetMatrix()).flatten()])
+    sitk_affine.SetTranslation([float(v) for v in affine.GetTranslation()])
+    sitk_affine.SetCenter([float(v) for v in affine.GetCenter()])
+    return sitk_affine
+
+
+class ConvexAdamEngine:
+    """Register a fixed/moving pair with the itk-impact ConvexAdam pipeline; return (moved, dvf) on the fixed grid.
+
+    The IMPACT feature models are downloaded once (``repo:filename`` on Hugging Face) and reused across cases.
+    Masks are accepted for signature compatibility with the elastix engine but ignored: the ConvexAdam
+    filters optimise over the whole image (no mask API is exposed by the coarse/fine stages).
+    """
+
+    def __init__(
+        self,
+        models: list[str],
+        voxel_size: list[float],
+        num_channels: int,
+        overlap: int,
+        layers_mask: list[bool],
+        mixed_precision: bool,
+        grid_spacing: int,
+        displacement_half_width: int,
+        iterations: int,
+        learning_rate: float,
+        regularization_weight: float,
+        grid_shrink: int,
+        distance: list[str],
+        layers_weight: list[float],
+        subset_features: list[int],
+        pca: list[int],
+        stages: list[str],
+        linear: bool,
+        linear_iterations: int,
+        seed: int,
+    ) -> None:
+        self._stages = stages
+        self._model_paths = self._download_models(models)
+        # Built lazily and cached: constructing an itk.ModelConfiguration loads the TorchScript model
+        # from disk in C++, so build the list once and reuse it across both stages and every case.
+        self._configurations: "list[itk.ModelConfiguration] | None" = None
+        self._voxel_size = voxel_size
+        self._num_channels = num_channels
+        self._overlap = overlap
+        self._layers_mask = layers_mask
+        self._mixed_precision = mixed_precision
+        self._grid_spacing = grid_spacing
+        self._displacement_half_width = displacement_half_width
+        self._iterations = iterations
+        self._learning_rate = learning_rate
+        self._regularization_weight = regularization_weight
+        self._grid_shrink = grid_shrink
+        self._distance = distance
+        self._layers_weight = layers_weight
+        self._subset_features = subset_features
+        self._pca = pca
+        self._linear = linear
+        self._linear_iterations = linear_iterations
+        self._seed = seed
+
+    @staticmethod
+    def _download_models(models: list[str]) -> list[str]:
+        """Fetch the TorchScript feature models (``repo:filename``); return their local paths."""
+        paths = []
+        for ref in models:
+            repo, filename = ref.split(":", 1)
+            paths.append(str(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")))  # nosec B615
+        return paths
+
+    def _model_configurations(self) -> list["itk.ModelConfiguration"]:
+        """Build one ``ModelConfiguration`` per feature model once, then reuse it across stages and cases.
+
+        Constructing an ``itk.ModelConfiguration`` loads the TorchScript module from disk on the C++ side, so
+        it is built lazily and cached. The coarse/fine filters copy each configuration by value in
+        ``AddModelConfiguration`` and the copy shares the loaded module through the configuration's internal
+        ``shared_ptr`` — so a single build is reused everywhere without any reload.
+        """
+        if self._configurations is None:
+            self._configurations = [
+                itk.ModelConfiguration(
+                    path,
+                    DIM,
+                    self._num_channels,
+                    [0, 0, 0],
+                    [float(v) for v in self._voxel_size],
+                    self._overlap,
+                    list(self._layers_mask),
+                    self._mixed_precision,
+                )
+                for path in self._model_paths
+            ]
+        return self._configurations
+
+    def _linear_align(self, fixed: "itk.Image", moving: "itk.Image") -> "itk.AffineTransform":
+        """Moments-initialised rigid + affine (Mattes MI), mapping fixed -> moving physical points."""
+        rigid = itk.VersorRigid3DTransform[itk.D].New()
+        initializer = itk.CenteredTransformInitializer[itk.VersorRigid3DTransform[itk.D], _IMAGE_F, _IMAGE_F].New(
+            Transform=rigid, FixedImage=fixed, MovingImage=moving
+        )
+        initializer.MomentsOn()
+        initializer.InitializeTransform()
+
+        affine = itk.AffineTransform[itk.D, DIM].New()
+        affine.SetCenter(rigid.GetCenter())
+        affine.SetMatrix(rigid.GetMatrix())
+        affine.SetOffset(rigid.GetOffset())
+        levels = 3
+        metric_type = itk.MattesMutualInformationImageToImageMetricv4[_IMAGE_F, _IMAGE_F]
+        metric = metric_type.New()
+        metric.SetNumberOfHistogramBins(32)
+        optimizer = itk.RegularStepGradientDescentOptimizerv4[itk.D].New()
+        optimizer.SetNumberOfIterations(self._linear_iterations)
+        optimizer.SetLearningRate(1.0)
+        optimizer.SetMinimumStepLength(1e-5)
+        optimizer.SetRelaxationFactor(0.6)
+        scales = itk.RegistrationParameterScalesFromPhysicalShift[metric_type].New()
+        scales.SetMetric(metric)
+        optimizer.SetScalesEstimator(scales)
+        registration = itk.ImageRegistrationMethodv4[_IMAGE_F, _IMAGE_F].New(
+            FixedImage=fixed, MovingImage=moving, Metric=metric, Optimizer=optimizer, InitialTransform=affine
+        )
+        registration.SetNumberOfLevels(levels)
+        registration.SetShrinkFactorsPerLevel([2 ** (levels - 1 - i) for i in range(levels)])
+        registration.SetSmoothingSigmasPerLevel([float(levels - 1 - i) for i in range(levels)])
+        registration.InPlaceOn()
+        registration.Update()
+        return affine
+
+    def _coarse(self, fixed: "itk.Image", moving: "itk.Image", device: str) -> "itk.Image":
+        """ConvexAdam coarse coupled-convex initializer -> robust low-resolution field on the fixed grid."""
+        coarse = _coarse_registration_type().New()
+        coarse.SetFixedImage(fixed)
+        coarse.SetMovingImage(moving)
+        for configuration in self._model_configurations():
+            coarse.AddModelConfiguration(configuration)
+        coarse.SetGridSpacing(self._grid_spacing)
+        coarse.SetDisplacementHalfWidth(self._displacement_half_width)
+        coarse.SetDevice(device)
+        coarse.SetSeed(self._seed)
+        coarse.Update()
+        field = coarse.GetOutput()
+        field.DisconnectPipeline()
+        return field
+
+    def _fine(
+        self, fixed: "itk.Image", moving: "itk.Image", initial_field: "itk.Image | None", device: str
+    ) -> "itk.Image":
+        """Adam instance-optimisation refinement, warm-started from ``initial_field`` (zero if none)."""
+        fine = _fine_registration_type().New()
+        fine.SetFixedImage(fixed)
+        fine.SetMovingImage(moving)
+        fine.SetInitialDisplacementField(initial_field if initial_field is not None else self._zero_field(fixed))
+        for configuration in self._model_configurations():
+            fine.AddModelConfiguration(configuration)
+        fine.SetDistance(list(self._distance))
+        fine.SetLayersWeight([float(v) for v in self._layers_weight])
+        fine.SetSubsetFeatures([int(v) for v in self._subset_features])
+        fine.SetPCA([int(v) for v in self._pca])
+        fine.SetNumberOfIterations(self._iterations)
+        fine.SetLearningRate(self._learning_rate)
+        fine.SetRegularizationWeight(self._regularization_weight)
+        fine.SetGridShrinkFactor(self._grid_shrink)
+        fine.SetDevice(device)
+        fine.SetSeed(self._seed)
+
+        # Mirror KonfAI's informative bars: drive a tqdm over the Adam iterations from the metric trace so
+        # SlicerKonfAI (which parses the "N% done/total" progress line) shows real progress. The observer is
+        # best-effort — if the filter does not emit IterationEvent the bar simply fills on completion.
+        progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True)
+
+        def _update(*_: object) -> None:
+            values = list(fine.GetMetricValuesPerIteration())
+            progress.n = min(len(values), self._iterations)
+            if values:
+                progress.set_description(f"Registration : iter {len(values)} | metric {float(values[-1]):.4f}")
+            progress.refresh()
+
+        try:
+            fine.AddObserver(itk.IterationEvent(), _update)
+        except Exception:  # nosec B110 - progress is best-effort; never fail a run over the bar
+            pass
+        fine.Update()
+        progress.close()
+        field = fine.GetDisplacementField()
+        field.DisconnectPipeline()
+        return field
+
+    @staticmethod
+    def _zero_field(reference: "itk.Image") -> "itk.Image":
+        """An all-zero displacement field on ``reference``'s grid (identity warm-start for a lone fine stage)."""
+        field = itk.Image[itk.Vector[itk.F, DIM], DIM].New()
+        field.CopyInformation(reference)
+        field.SetRegions(reference.GetLargestPossibleRegion())
+        field.Allocate()
+        zero = itk.Vector[itk.F, DIM]()
+        zero.Fill(0)  # itk::Vector default ctor does not zero-initialise
+        field.FillBuffer(zero)
+        return field
+
+    def _run_stages(self, fixed: "itk.Image", moving: "itk.Image", device: str) -> "itk.Image | None":
+        """Run the configured coarse/fine chain; each fine warm-starts from the running field.
+
+        ``coarse`` produces a field from scratch; ``fine`` refines the running field. So ``['coarse']`` is a
+        coarse-only app, ``['fine']`` a fine-only app (zero warm-start), and ``['coarse', 'fine']`` chains both
+        (the composite, as before). Returns None when no deformable stage runs (e.g. a linear-only chain).
+        """
+        field: "itk.Image | None" = None
+        for stage in self._stages:
+            if stage == "coarse":
+                field = self._coarse(fixed, moving, device)
+            elif stage == "fine":
+                field = self._fine(fixed, moving, field, device)
+            else:
+                raise ValueError(f"Unknown registration stage '{stage}' (expected 'coarse' or 'fine').")
+        return field
+
+    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."""
+        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
+        fixed_itk = _sitk_to_itk(fixed)
+        moving_itk = _sitk_to_itk(moving)
+
+        # Optional linear pre-align: resample the moving onto the fixed grid so the deformable stage starts close.
+        affine = self._linear_align(fixed_itk, moving_itk) if self._linear else itk.AffineTransform[itk.D, DIM].New()
+        resampler = itk.ResampleImageFilter[_IMAGE_F, _IMAGE_F].New(
+            Input=moving_itk, ReferenceImage=fixed_itk, Transform=affine
+        )
+        resampler.UseReferenceImageOn()
+        resampler.SetInterpolator(itk.LinearInterpolateImageFunction[_IMAGE_F, itk.D].New())
+        resampler.Update()
+        moving_linear = resampler.GetOutput()
+
+        field = self._run_stages(fixed_itk, moving_linear, device)
+
+        # One transform on the fixed grid = affine then deformable, so the returned DVF/transform warps the
+        # ORIGINAL moving. SimpleITK applies the last-added transform first, so [affine, deformable] gives
+        # moved(p) = moving(affine(deformable(p))). A linear-only chain (field is None) yields the affine alone.
+        chain = [_itk_affine_to_sitk(affine)]
+        if field is not None:
+            chain.append(_itk_field_to_sitk_transform(field, fixed))
+        composite = sitk.CompositeTransform(chain)
+        moved = sitk.Resample(moving, fixed, composite, sitk.sitkLinear, 0.0, moving.GetPixelID())
+        dvf = sitk.TransformToDisplacementField(
+            composite,
+            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
+
+
+class ConvexAdamRegistration(torch.nn.Module):
+    """Graph module: (fixed, moving) tensors + their geometry -> moved image + DVF on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving the per-branch ``Attribute`` list alongside
+    the tensors (same convention as ``CriterionWithAttribute``); registration needs the physical geometry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: ConvexAdamEngine) -> None:
+        super().__init__()
+        self._engine = engine
+
+    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 a list[Attribute] over the batch.
+        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field (DIM
+        # channels); downstream ChannelSelect modules split them. Masks are ignored by the ConvexAdam engine.
+        fixed_attrs, moving_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])
+            moved_np, dvf_np = self._engine.register(fixed_img, moving_img, device_index)
+            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 ConvexAdam registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1;
+    the mask branches 2/3 are accepted but unused by this engine).
+
+    Outputs on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField`` (the
+    DIM-component displacement field, in mm). 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()},
+        models: list[str] = [],
+        voxel_size: list[float] = [3.0, 3.0, 3.0],
+        num_channels: int = 1,
+        overlap: int = 2,
+        layers_mask: list[bool] = [True],
+        mixed_precision: bool = False,
+        grid_spacing: int = 4,
+        displacement_half_width: int = 6,
+        iterations: int = 150,
+        learning_rate: float = 0.2,
+        regularization_weight: float = 1.0,
+        grid_shrink: int = 4,
+        distance: list[str] = ["L1"],
+        layers_weight: list[float] = [1.0],
+        subset_features: list[int] = [32],
+        pca: list[int] = [0],
+        stages: list[str] = ["coarse", "fine"],
+        linear: bool = True,
+        linear_iterations: int = 200,
+        seed: int = 42,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        engine = ConvexAdamEngine(
+            models,
+            voxel_size,
+            num_channels,
+            overlap,
+            layers_mask,
+            mixed_precision,
+            grid_spacing,
+            displacement_half_width,
+            iterations,
+            learning_rate,
+            regularization_weight,
+            grid_shrink,
+            distance,
+            layers_weight,
+            subset_features,
+            pca,
+            stages,
+            linear,
+            linear_iterations,
+            seed,
+        )
+        self.add_module(
+            "Registration", ConvexAdamRegistration(engine), 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"])

ConvexAdam_Fine/Prediction.yml

@@ -0,0 +1,140 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      models:
+      - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt
+      voxel_size:
+      - 3.0
+      - 3.0
+      - 3.0
+      num_channels: 1
+      overlap: 2
+      layers_mask:
+      - true
+      mixed_precision: false
+      grid_spacing: 4
+      displacement_half_width: 6
+      iterations: 150
+      learning_rate: 0.2
+      regularization_weight: 1.0
+      grid_shrink: 4
+      distance:
+      - L1
+      layers_weight:
+      - 1.0
+      subset_features:
+      - 32
+      pca:
+      - 0
+      linear: false
+      linear_iterations: 200
+      seed: 42
+      outputs_criterions: None
+      stages:
+      - fine
+  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:
+      DataAugmentation_0:
+        data_augmentations:
+          Flip:
+            f_prob:
+            - 0
+            - 0.5
+            - 0.5
+            vector_field: true
+            prob: 1
+        nb: 2
+    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-ConvexAdam-Fine
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

ConvexAdam_Fine/Uncertainty.yml

@@ -0,0 +1,24 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean:
+              name: Uncertainty
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Norm: {}
+              StandardDeviation: {}
+              Save:
+                dataset: ./Uncertainties/ImpactReg/Output:mha
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

ConvexAdam_Fine/app.json

@@ -0,0 +1,100 @@
+{
+    "display_name": "ConvexAdam Fine (MIND)",
+    "short_description": "Fine Adam refinement (tileable; expects a coarse/pre-aligned start).",
+    "description": "Second stage of the ConvexAdam pipeline: the Adam instance-optimisation refinement driven by the IMPACT metric on MIND features, with no linear pre-align (it assumes the moving is already on the fixed grid, e.g. resampled by the coarse stage). Runs from a zero warm-start and tiles naturally for large images (set a patch size).",
+    "task": "registration",
+    "tta": 3,
+    "mc_dropout": 0,
+    "vram_plan": {
+        "8": {"patch_size": [128, 128, 128], "batch_size": 1},
+        "16": {"patch_size": [192, 192, 192], "batch_size": 1},
+        "24": {"patch_size": [256, 256, 256], "batch_size": 1},
+        "40": {"patch_size": [320, 320, 320], "batch_size": 1}
+    },
+    "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
+                }
+            }
+        }
+    }
+}

ConvexAdam_Fine/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de99fbc36331ce674639acc774f52b4a2d0027f2f312d9d28669e831a0c4fd7e
+size 1249

ConvexAdam_Fine/requirements.txt

@@ -0,0 +1 @@
+itk-impact