Move engine code to impact_reg_konfai.models; keep only config + weights

Each preset's Prediction.yml now points its classpath at
impact_reg_konfai.models.<engine>:RegistrationNet instead of a bundled Model.py,
and declares impact_reg_konfai>=1.6.1. The duplicated engine code (Model.py, plus
elastix_engine.py/install.py) is removed; the elastix parameter-map .txt stay, as
they are per-preset config. Requires impact_reg_konfai>=1.6.1 (ships the models
subpackage).

CBCT_CT_HeadNeck/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            in_branch=[0, 1, 2, 3],
-            out_branch=["registration"],
-        )
-        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
-        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

CBCT_CT_HeadNeck/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps:

CBCT_CT_HeadNeck/elastix_engine.py

@@ -1,386 +0,0 @@
-# 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
-
-"""Elastix-IMPACT runtime for the registration bundle.
-
-``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
-parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
-``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
-images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
-"""
-
-import os
-import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-import tqdm
-from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
-
-
-def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
-    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
-    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
-    the whole image) instead of paying for a mask that restricts nothing."""
-    if mask is None:
-        return False
-    arr = sitk.GetArrayViewFromImage(mask)
-    return bool((arr > 0).any()) and bool((arr == 0).any())
-
-
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
-    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random FOV-sized patches each iteration. One mode per preset.
-        self._mode = mode
-        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
-        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry() if resolutions else {}
-        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
-        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
-        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
-        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
-        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
-        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
-        # surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into ``work``.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
-        applies every per-token / exact override. Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
-        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
-            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if _is_partial_mask(mask):
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
-            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
-            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
-    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            mode,
-        )
-
-    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) stacked with the DVF (dim channels), both on
-        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)

CBCT_CT_HeadNeck/install.py

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

CBCT_CT_HeadNeck/requirements.txt

@@ -0,0 +1 @@
+impact_reg_konfai>=1.6.1

CBCT_CT_MRSeg/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            in_branch=[0, 1, 2, 3],
-            out_branch=["registration"],
-        )
-        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
-        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

CBCT_CT_MRSeg/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps:

CBCT_CT_MRSeg/elastix_engine.py

@@ -1,386 +0,0 @@
-# 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
-
-"""Elastix-IMPACT runtime for the registration bundle.
-
-``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
-parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
-``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
-images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
-"""
-
-import os
-import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-import tqdm
-from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
-
-
-def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
-    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
-    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
-    the whole image) instead of paying for a mask that restricts nothing."""
-    if mask is None:
-        return False
-    arr = sitk.GetArrayViewFromImage(mask)
-    return bool((arr > 0).any()) and bool((arr == 0).any())
-
-
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
-    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random FOV-sized patches each iteration. One mode per preset.
-        self._mode = mode
-        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
-        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry() if resolutions else {}
-        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
-        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
-        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
-        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
-        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
-        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
-        # surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into ``work``.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
-        applies every per-token / exact override. Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
-        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
-            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if _is_partial_mask(mask):
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
-            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
-            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
-    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            mode,
-        )
-
-    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) stacked with the DVF (dim channels), both on
-        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)

CBCT_CT_MRSeg/install.py

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

CBCT_CT_MRSeg/requirements.txt

@@ -0,0 +1 @@
+impact_reg_konfai>=1.6.1

ConvexAdam_Coarse/Model.py

@@ -1,549 +0,0 @@
-# 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 contextlib
-import json
-import os
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-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.config import Choices, Range
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-DIM = 3
-# The feature model's input channel count is an intrinsic property of the pretrained model (grayscale
-# medical images), not a tunable — so it's fixed here, never a config/signature parameter.
-NUM_CHANNELS = 1
-
-# A UI reads the tuning knobs straight from the TYPES on ``RegistrationNet.__init__`` and ``ModelSpec``:
-# ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` / ``Annotated[str, Choices]`` a dropdown.
-# ``models`` is a dict-of-objects (one ``ModelSpec`` per feature model) — the same shape as the elastix presets,
-# so SlicerKonfAI renders each model as a repeatable block with a ``ref`` / ``distance`` combo box.
-_IMAGE_F = itk.Image[itk.F, DIM]
-
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-
-def registry_choices() -> list[str]:
-    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry the
-    engine already fetches (offline-first). A user may still point ``ref`` at a local model path."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load ``models.json`` (the available feature models) from the model repo on Hugging Face. The registry is
-    NOT bundled: ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for dev/offline; otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference."""
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (fetched from HF, not "
-            "bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-@dataclass
-class ModelSpec:
-    """One feature model in the ConvexAdam multi-feature fusion. ``ref`` picks the model; the rest are its
-    per-model knobs (all models are compared jointly by the IMPACT metric). Same idea as the elastix ``ModelSpec``."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=lambda: [3.0, 3.0, 3.0])
-    layers_mask: str = "1"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
-    layers_weight: float = 1.0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-
-
-@contextlib.contextmanager
-def _no_texpr_fuser():
-    """Disable ONLY torch's TensorExpr (NNC) fuser for the block. The TS feature-model graphs have a direction/
-    orientation branch whose shape ops (aten::dim / aten::size) crash the fuser's alias analysis under itk-impact
-    ("INTERNAL ASSERT ... We don't have an op for aten::size" in FuseTensorExprs). It cannot be fixed model-side:
-    itk-impact passes an *undefined* direction, so guarding it needs a shape op — the very fuser trigger. The
-    modern profiling executor stays on (NOT the legacy executor); measured cost ~1% of a registration (the fuser
-    only touches the few feature forwards, not the C++ optimisation loop). The caller's setting is restored."""
-    prev = torch._C._jit_texpr_fuser_enabled()
-    torch._C._jit_set_texpr_fuser_enabled(False)
-    try:
-        yield
-    finally:
-        torch._C._jit_set_texpr_fuser_enabled(prev)
-
-
-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_sizes: list[list[float]],
-        overlap: int,
-        layers_masks: list[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_sizes = voxel_sizes
-        self._overlap = overlap
-        self._layers_masks = layers_masks
-        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,
-                    NUM_CHANNELS,
-                    [0, 0, 0],
-                    [float(v) for v in voxel_size],
-                    self._overlap,
-                    list(layers_mask),
-                    self._mixed_precision,
-                )
-                for path, voxel_size, layers_mask in zip(self._model_paths, self._voxel_sizes, self._layers_masks)
-            ]
-        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)
-
-        # Optional terminal progress over the Adam iterations, driven from the metric trace. ``disable=None``
-        # auto-hides it when stderr is not a TTY (e.g. under KonfAI/Slicer, where the outer "Prediction" bar
-        # already reports progress), so captured logs stay clean; ``leave=False`` avoids stacking one bar per
-        # patch. The observer is best-effort — if the filter emits no IterationEvent the bar just fills at the end.
-        progress = tqdm.tqdm(
-            total=self._iterations or None, desc="Registration", ncols=0, leave=False, disable=None
-        )
-
-        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.n = progress.total or self._iterations  # show completion even if no IterationEvent fired
-        progress.refresh()
-        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 = []
-        # ConvexAdam runs a gradient-based instance optimisation (Adam over the field) inside itk-impact's
-        # .Update(); the predictor calls forward under torch.inference_mode(), which forbids autograd. The
-        # image tensors have already crossed to numpy/ITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad(), _no_texpr_fuser():
-            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: dict[str, ModelSpec] = {},
-        overlap: Annotated[int, Range(1, 128)] = 2,
-        mixed_precision: bool = False,
-        grid_spacing: Annotated[int, Range(1, 512)] = 4,
-        displacement_half_width: Annotated[int, Range(1, 512)] = 6,
-        iterations: Annotated[int, Range(0, 100000)] = 150,
-        learning_rate: Annotated[float, Range(0.0, 100.0)] = 0.2,
-        regularization_weight: Annotated[float, Range(0.0, 1000.0)] = 1.0,
-        grid_shrink: Annotated[int, Range(1, 128)] = 4,
-        subset_features: list[int] = [],  # feature-channel indices to keep (empty = all); NOT a count
-        stages: list[str] = ["coarse", "fine"],
-        linear: bool = True,
-        linear_iterations: Annotated[int, Range(0, 100000)] = 200,
-        seed: int = 42,
-    ) -> None:
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        specs = _sorted_specs(models)
-        engine = ConvexAdamEngine(
-            [spec.ref for spec in specs],
-            [list(spec.voxel_size) for spec in specs],
-            overlap,
-            [[c == "1" for c in spec.layers_mask] for spec in specs],
-            mixed_precision,
-            grid_spacing,
-            displacement_half_width,
-            iterations,
-            learning_rate,
-            regularization_weight,
-            grid_shrink,
-            [spec.distance for spec in specs],
-            [float(spec.layers_weight) for spec in specs],
-            subset_features,
-            [int(spec.pca) for spec in specs],
-            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

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.convexadam:RegistrationNet
     RegistrationNet:
       models:
         '0':

ConvexAdam_Coarse/requirements.txt

@@ -1 +1,2 @@
 itk-impact
+impact_reg_konfai>=1.6.1

ConvexAdam_Composite/Model.py

@@ -1,549 +0,0 @@
-# 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 contextlib
-import json
-import os
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-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.config import Choices, Range
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-DIM = 3
-# The feature model's input channel count is an intrinsic property of the pretrained model (grayscale
-# medical images), not a tunable — so it's fixed here, never a config/signature parameter.
-NUM_CHANNELS = 1
-
-# A UI reads the tuning knobs straight from the TYPES on ``RegistrationNet.__init__`` and ``ModelSpec``:
-# ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` / ``Annotated[str, Choices]`` a dropdown.
-# ``models`` is a dict-of-objects (one ``ModelSpec`` per feature model) — the same shape as the elastix presets,
-# so SlicerKonfAI renders each model as a repeatable block with a ``ref`` / ``distance`` combo box.
-_IMAGE_F = itk.Image[itk.F, DIM]
-
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-
-def registry_choices() -> list[str]:
-    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry the
-    engine already fetches (offline-first). A user may still point ``ref`` at a local model path."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load ``models.json`` (the available feature models) from the model repo on Hugging Face. The registry is
-    NOT bundled: ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for dev/offline; otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference."""
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (fetched from HF, not "
-            "bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-@dataclass
-class ModelSpec:
-    """One feature model in the ConvexAdam multi-feature fusion. ``ref`` picks the model; the rest are its
-    per-model knobs (all models are compared jointly by the IMPACT metric). Same idea as the elastix ``ModelSpec``."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=lambda: [3.0, 3.0, 3.0])
-    layers_mask: str = "1"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
-    layers_weight: float = 1.0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-
-
-@contextlib.contextmanager
-def _no_texpr_fuser():
-    """Disable ONLY torch's TensorExpr (NNC) fuser for the block. The TS feature-model graphs have a direction/
-    orientation branch whose shape ops (aten::dim / aten::size) crash the fuser's alias analysis under itk-impact
-    ("INTERNAL ASSERT ... We don't have an op for aten::size" in FuseTensorExprs). It cannot be fixed model-side:
-    itk-impact passes an *undefined* direction, so guarding it needs a shape op — the very fuser trigger. The
-    modern profiling executor stays on (NOT the legacy executor); measured cost ~1% of a registration (the fuser
-    only touches the few feature forwards, not the C++ optimisation loop). The caller's setting is restored."""
-    prev = torch._C._jit_texpr_fuser_enabled()
-    torch._C._jit_set_texpr_fuser_enabled(False)
-    try:
-        yield
-    finally:
-        torch._C._jit_set_texpr_fuser_enabled(prev)
-
-
-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_sizes: list[list[float]],
-        overlap: int,
-        layers_masks: list[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_sizes = voxel_sizes
-        self._overlap = overlap
-        self._layers_masks = layers_masks
-        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,
-                    NUM_CHANNELS,
-                    [0, 0, 0],
-                    [float(v) for v in voxel_size],
-                    self._overlap,
-                    list(layers_mask),
-                    self._mixed_precision,
-                )
-                for path, voxel_size, layers_mask in zip(self._model_paths, self._voxel_sizes, self._layers_masks)
-            ]
-        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)
-
-        # Optional terminal progress over the Adam iterations, driven from the metric trace. ``disable=None``
-        # auto-hides it when stderr is not a TTY (e.g. under KonfAI/Slicer, where the outer "Prediction" bar
-        # already reports progress), so captured logs stay clean; ``leave=False`` avoids stacking one bar per
-        # patch. The observer is best-effort — if the filter emits no IterationEvent the bar just fills at the end.
-        progress = tqdm.tqdm(
-            total=self._iterations or None, desc="Registration", ncols=0, leave=False, disable=None
-        )
-
-        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.n = progress.total or self._iterations  # show completion even if no IterationEvent fired
-        progress.refresh()
-        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 = []
-        # ConvexAdam runs a gradient-based instance optimisation (Adam over the field) inside itk-impact's
-        # .Update(); the predictor calls forward under torch.inference_mode(), which forbids autograd. The
-        # image tensors have already crossed to numpy/ITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad(), _no_texpr_fuser():
-            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: dict[str, ModelSpec] = {},
-        overlap: Annotated[int, Range(1, 128)] = 2,
-        mixed_precision: bool = False,
-        grid_spacing: Annotated[int, Range(1, 512)] = 4,
-        displacement_half_width: Annotated[int, Range(1, 512)] = 6,
-        iterations: Annotated[int, Range(0, 100000)] = 150,
-        learning_rate: Annotated[float, Range(0.0, 100.0)] = 0.2,
-        regularization_weight: Annotated[float, Range(0.0, 1000.0)] = 1.0,
-        grid_shrink: Annotated[int, Range(1, 128)] = 4,
-        subset_features: list[int] = [],  # feature-channel indices to keep (empty = all); NOT a count
-        stages: list[str] = ["coarse", "fine"],
-        linear: bool = True,
-        linear_iterations: Annotated[int, Range(0, 100000)] = 200,
-        seed: int = 42,
-    ) -> None:
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        specs = _sorted_specs(models)
-        engine = ConvexAdamEngine(
-            [spec.ref for spec in specs],
-            [list(spec.voxel_size) for spec in specs],
-            overlap,
-            [[c == "1" for c in spec.layers_mask] for spec in specs],
-            mixed_precision,
-            grid_spacing,
-            displacement_half_width,
-            iterations,
-            learning_rate,
-            regularization_weight,
-            grid_shrink,
-            [spec.distance for spec in specs],
-            [float(spec.layers_weight) for spec in specs],
-            subset_features,
-            [int(spec.pca) for spec in specs],
-            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

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.convexadam:RegistrationNet
     RegistrationNet:
       models:
         '0':

ConvexAdam_Composite/requirements.txt

@@ -1 +1,2 @@
 itk-impact
+impact_reg_konfai>=1.6.1

ConvexAdam_Fine/Model.py

@@ -1,549 +0,0 @@
-# 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 contextlib
-import json
-import os
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-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.config import Choices, Range
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-DIM = 3
-# The feature model's input channel count is an intrinsic property of the pretrained model (grayscale
-# medical images), not a tunable — so it's fixed here, never a config/signature parameter.
-NUM_CHANNELS = 1
-
-# A UI reads the tuning knobs straight from the TYPES on ``RegistrationNet.__init__`` and ``ModelSpec``:
-# ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` / ``Annotated[str, Choices]`` a dropdown.
-# ``models`` is a dict-of-objects (one ``ModelSpec`` per feature model) — the same shape as the elastix presets,
-# so SlicerKonfAI renders each model as a repeatable block with a ``ref`` / ``distance`` combo box.
-_IMAGE_F = itk.Image[itk.F, DIM]
-
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-
-def registry_choices() -> list[str]:
-    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry the
-    engine already fetches (offline-first). A user may still point ``ref`` at a local model path."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load ``models.json`` (the available feature models) from the model repo on Hugging Face. The registry is
-    NOT bundled: ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for dev/offline; otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference."""
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (fetched from HF, not "
-            "bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-@dataclass
-class ModelSpec:
-    """One feature model in the ConvexAdam multi-feature fusion. ``ref`` picks the model; the rest are its
-    per-model knobs (all models are compared jointly by the IMPACT metric). Same idea as the elastix ``ModelSpec``."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=lambda: [3.0, 3.0, 3.0])
-    layers_mask: str = "1"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
-    layers_weight: float = 1.0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-
-
-@contextlib.contextmanager
-def _no_texpr_fuser():
-    """Disable ONLY torch's TensorExpr (NNC) fuser for the block. The TS feature-model graphs have a direction/
-    orientation branch whose shape ops (aten::dim / aten::size) crash the fuser's alias analysis under itk-impact
-    ("INTERNAL ASSERT ... We don't have an op for aten::size" in FuseTensorExprs). It cannot be fixed model-side:
-    itk-impact passes an *undefined* direction, so guarding it needs a shape op — the very fuser trigger. The
-    modern profiling executor stays on (NOT the legacy executor); measured cost ~1% of a registration (the fuser
-    only touches the few feature forwards, not the C++ optimisation loop). The caller's setting is restored."""
-    prev = torch._C._jit_texpr_fuser_enabled()
-    torch._C._jit_set_texpr_fuser_enabled(False)
-    try:
-        yield
-    finally:
-        torch._C._jit_set_texpr_fuser_enabled(prev)
-
-
-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_sizes: list[list[float]],
-        overlap: int,
-        layers_masks: list[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_sizes = voxel_sizes
-        self._overlap = overlap
-        self._layers_masks = layers_masks
-        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,
-                    NUM_CHANNELS,
-                    [0, 0, 0],
-                    [float(v) for v in voxel_size],
-                    self._overlap,
-                    list(layers_mask),
-                    self._mixed_precision,
-                )
-                for path, voxel_size, layers_mask in zip(self._model_paths, self._voxel_sizes, self._layers_masks)
-            ]
-        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)
-
-        # Optional terminal progress over the Adam iterations, driven from the metric trace. ``disable=None``
-        # auto-hides it when stderr is not a TTY (e.g. under KonfAI/Slicer, where the outer "Prediction" bar
-        # already reports progress), so captured logs stay clean; ``leave=False`` avoids stacking one bar per
-        # patch. The observer is best-effort — if the filter emits no IterationEvent the bar just fills at the end.
-        progress = tqdm.tqdm(
-            total=self._iterations or None, desc="Registration", ncols=0, leave=False, disable=None
-        )
-
-        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.n = progress.total or self._iterations  # show completion even if no IterationEvent fired
-        progress.refresh()
-        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 = []
-        # ConvexAdam runs a gradient-based instance optimisation (Adam over the field) inside itk-impact's
-        # .Update(); the predictor calls forward under torch.inference_mode(), which forbids autograd. The
-        # image tensors have already crossed to numpy/ITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad(), _no_texpr_fuser():
-            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: dict[str, ModelSpec] = {},
-        overlap: Annotated[int, Range(1, 128)] = 2,
-        mixed_precision: bool = False,
-        grid_spacing: Annotated[int, Range(1, 512)] = 4,
-        displacement_half_width: Annotated[int, Range(1, 512)] = 6,
-        iterations: Annotated[int, Range(0, 100000)] = 150,
-        learning_rate: Annotated[float, Range(0.0, 100.0)] = 0.2,
-        regularization_weight: Annotated[float, Range(0.0, 1000.0)] = 1.0,
-        grid_shrink: Annotated[int, Range(1, 128)] = 4,
-        subset_features: list[int] = [],  # feature-channel indices to keep (empty = all); NOT a count
-        stages: list[str] = ["coarse", "fine"],
-        linear: bool = True,
-        linear_iterations: Annotated[int, Range(0, 100000)] = 200,
-        seed: int = 42,
-    ) -> None:
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        specs = _sorted_specs(models)
-        engine = ConvexAdamEngine(
-            [spec.ref for spec in specs],
-            [list(spec.voxel_size) for spec in specs],
-            overlap,
-            [[c == "1" for c in spec.layers_mask] for spec in specs],
-            mixed_precision,
-            grid_spacing,
-            displacement_half_width,
-            iterations,
-            learning_rate,
-            regularization_weight,
-            grid_shrink,
-            [spec.distance for spec in specs],
-            [float(spec.layers_weight) for spec in specs],
-            subset_features,
-            [int(spec.pca) for spec in specs],
-            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

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.convexadam:RegistrationNet
     RegistrationNet:
       models:
         '0':

ConvexAdam_Fine/requirements.txt

@@ -1 +1,2 @@
 itk-impact
+impact_reg_konfai>=1.6.1

FireANTs_Affine/Model.py

@@ -1,669 +0,0 @@
-# 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
-#
-# This wrapper does NOT copy any FireANTs source: it only calls the public FireANTs API of the
-# separately-installed ``fireants`` wheel (PyPI). FireANTs is distributed under the FireANTs License
-# v1.0 and must be cited — see the NOTICE file in this directory for the license, copyright and
-# bibliography that ship with this app.
-
-"""FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets).
-
-Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset
-(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the
-orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own
-composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages:
-
-    Rigid (MI, centre-of-mass init) -> Affine (MI, seeded by the rigid) -> deformable
-
-The deformable stage is selected by ``deformable_method`` — the ONE knob that specialises this shared
-Model.py into the different presets (exactly as ConvexAdam's shared Model.py is specialised by
-``stages``):
-
-    "syn"    symmetric diffeomorphic SyN (CC)   — invertible, higher quality, averages cleanly for ensembling
-    "greedy" greedy diffeomorphic (CC)          — one-directional, faster / lower VRAM
-    "none"   linear only                        — Rigid+Affine, no deformable (the FireANTs_Affine preset)
-
-Masks: the optional Fixed/Moving masks restrict the metric to a region. FireANTs implements this by
-carrying the mask as the last image channel and prefixing the metric with ``masked_``; a mask is only
-honoured when it actually restricts (some voxels in, some out), so the common mask-free path is
-unchanged (an absent optional mask arrives as a whole-image default and is treated as no mask).
-
-The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear
-pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly.
-``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK —
-the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an
-ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is
-round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here).
-
-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 contextlib
-import json
-import os
-import tempfile
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Annotated, Literal
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-from konfai.metric.measure import IMPACTReg
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-DIM = 3
-
-# Feature-model registry (models.json): the available IMPACT feature models, fetched from HF (NOT bundled).
-# Only consulted by the "impact" deformable metric; ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins
-# for dev/offline. Mirrors the ConvexAdam preset so the same 30-model catalogue and picker are shared.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-_DISTANCES: dict[str, type[torch.nn.Module]] = {"L1": torch.nn.L1Loss, "L2": torch.nn.MSELoss}
-
-
-def _fireants_git_ref() -> str:
-    """Best-effort FireANTs git ref whose ``fused_ops`` matches the installed ``fireants``.
-
-    Overridable with ``FIREANTS_FUSED_OPS_REF``; falls back to ``main`` if the version is unknown.
-    """
-    try:
-        import importlib.metadata
-
-        version = importlib.metadata.version("fireants").strip()
-        if version:
-            return version if version.startswith("v") else f"v{version}"
-    except Exception:
-        pass
-    return "main"
-
-
-def ensure_fireants_runtime(build_kernels: bool = True) -> None:
-    """Make the ``fireants`` runtime importable before a preset uses it — best-effort, never fatal.
-
-    A plain ``pip install fireants`` fails inside a host like 3D Slicer for two reasons, both handled
-    here with a clear one-line status at each step:
-
-    1. **fireants won't install.** It pins ``simpleitk==2.2.1`` (no wheel on modern Python), while the
-       host already ships a newer SimpleITK. We install it with ``--no-deps`` so that pin is ignored
-       and the host's SimpleITK/torch are reused; its light deps ship in ``requirements.txt``.
-    2. **The fused CUDA kernels** (``fireants_fused_ops``) that make registration fast and
-       memory-light are OPTIONAL. Without them fireants runs in pure PyTorch (correct, only slower).
-       We enable them only when a CUDA compiler (``nvcc``) is present, compiling from the upstream
-       FireANTs source at install time — nothing is vendored into this app. No compiler, or a failed
-       build, simply falls back to pure PyTorch.
-    """
-    import importlib
-    import shutil
-    import subprocess
-    import sys
-
-    def _log(message: str) -> None:
-        print(f"[FireANTs] {message}", flush=True)
-
-    # 1) fireants itself — install --no-deps to sidestep its unsatisfiable ``simpleitk==2.2.1`` pin.
-    try:
-        importlib.import_module("fireants")
-    except Exception:
-        _log("installing fireants (--no-deps, reusing the host's SimpleITK/torch)...")
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--no-deps", "fireants"])
-            importlib.invalidate_caches()
-            importlib.import_module("fireants")
-            _log("fireants installed.")
-        except Exception as exc:
-            raise RuntimeError(
-                "Could not install 'fireants'. Install it manually with:\n"
-                f"    {sys.executable} -m pip install --no-deps fireants\n"
-                "(its hydra-core/nibabel/pandas dependencies ship in this app's requirements.txt).\n"
-                f"Original error: {exc}"
-            ) from exc
-
-    if not build_kernels or os.environ.get("FIREANTS_SKIP_FUSED_OPS", "").strip().lower() in {"1", "true", "yes"}:
-        _log("skipping the fused CUDA kernels -> pure PyTorch (correct, slower and more memory).")
-        return
-
-    # 2) fused CUDA kernels — optional accelerator.
-    try:
-        importlib.import_module("fireants_fused_ops")
-        _log("fused CUDA kernels already available.")
-        return
-    except Exception:
-        pass
-
-    # 2a) a prebuilt wheel matching this torch's CUDA build, if one is ever published.
-    cuda_tag = ""
-    try:
-        cuda_tag = (torch.version.cuda or "").replace(".", "")
-    except Exception:
-        pass
-    if cuda_tag:
-        wheel = f"fireants-fused-ops-cu{cuda_tag}"
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", wheel])
-            importlib.invalidate_caches()
-            importlib.import_module("fireants_fused_ops")
-            _log(f"installed prebuilt CUDA kernels ({wheel}).")
-            return
-        except Exception:
-            pass  # no matching wheel -> try a local build
-
-    # 2b) local build is OPT-IN: compiling CUDA kernels is heavy and can exhaust RAM on the user's
-    # machine, so it NEVER runs by default. Set FIREANTS_BUILD_KERNELS=1 to enable it (devs with a
-    # CUDA toolkit); it then builds ONE file at a time (MAX_JOBS=1) to keep memory bounded. The clean
-    # path for end users is a prebuilt wheel (2a) — a local compile also needs Python dev headers
-    # (absent from some embedded Pythons, e.g. Slicer) and a CUDA-compatible host compiler.
-    if os.environ.get("FIREANTS_BUILD_KERNELS", "").strip().lower() not in ("1", "true", "yes"):
-        _log(
-            "no prebuilt kernels for this platform -> running FireANTs in pure PyTorch (correct, slower "
-            "and heavier). Set FIREANTS_BUILD_KERNELS=1 to compile them locally (needs a CUDA toolkit)."
-        )
-        return
-    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
-    nvcc = shutil.which("nvcc") or (os.path.join(cuda_home, "bin", "nvcc") if cuda_home else None)
-    if not nvcc or not os.path.exists(nvcc) or shutil.which("git") is None:
-        _log("FIREANTS_BUILD_KERNELS set but no CUDA toolkit (nvcc) / git found -> pure PyTorch.")
-        return
-
-    ref = os.environ.get("FIREANTS_FUSED_OPS_REF") or _fireants_git_ref()
-    _log(f"nvcc found ({nvcc}); compiling the FireANTs CUDA kernels one file at a time (ref '{ref}')...")
-    import tempfile
-
-    env = os.environ.copy()
-    env["MAX_JOBS"] = "1"  # one compile at a time -> bounded RAM (prevents OOM on large hosts)
-    env.setdefault("NVCC_APPEND_FLAGS", "-allow-unsupported-compiler")  # tolerate a newer host compiler
-    tmp = tempfile.mkdtemp(prefix="fireants_fused_ops_")
-    try:
-        # Shallow clone WITHOUT --recursive: FireANTs' submodules (an SSH-only 'cookbook') are unrelated
-        # to fused_ops and would otherwise abort the build with a public-key/permission error.
-        subprocess.check_call(
-            ["git", "clone", "--depth", "1", "--branch", ref, "https://github.com/rohitrango/FireANTs.git", tmp]
-        )
-        subprocess.check_call(
-            [sys.executable, "-m", "pip", "install", "-q", "--no-build-isolation", os.path.join(tmp, "fused_ops")],
-            env=env,
-        )
-        importlib.invalidate_caches()
-        importlib.import_module("fireants_fused_ops")
-        _log("compiled and installed the fast CUDA kernels.")
-    except Exception as exc:
-        _log(
-            "kernel build failed -> running in pure PyTorch (correct, only speed/memory affected). Set "
-            "FIREANTS_FUSED_OPS_REF to a compatible FireANTs tag, or install a prebuilt fireants-fused-ops "
-            f"wheel, to enable them. Details: {exc}"
-        )
-    finally:
-        shutil.rmtree(tmp, ignore_errors=True)
-
-
-def registry_choices() -> list[str]:
-    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load ``models.json`` (available feature models). ``KONFAI_IMPACT_MODELS_REGISTRY`` (local path) wins
-    for dev/offline; otherwise ``ref`` is a ``repo:file`` Hugging Face reference (fetched, not bundled)."""
-    from huggingface_hub import hf_hub_download
-
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference — or set "
-            "KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-@dataclass
-class ModelSpec:
-    """One IMPACT feature model in the deformable metric (several are fused). ``ref`` picks the model; the
-    rest are its per-model knobs — the same as the ConvexAdam / elastix ``ModelSpec`` except ``voxel_size``
-    (an itk-impact resampling knob) has no meaning for FireANTs' geometry-free torch ``custom_loss`` and is
-    intentionally absent."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    layers_mask: str = "01"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
-    layers_weight: float = 1.0  # this model's weight in the multi-model fusion
-    pca: Annotated[int, Range(0, 100)] = 0  # keep the top-K principal components of the features (0 = no PCA)
-    distance: Literal["L1", "L2"] = "L1"
-
-
-@contextlib.contextmanager
-def _no_texpr_fuser():
-    """Disable the TensorExpr JIT fuser while IMPACT's TorchScript feature model runs under autograd.
-
-    The IMPACT feature models are TorchScript; run under FireANTs' gradient optimisation the TensorExpr
-    fuser trips on shape ops (``aten::size`` INTERNAL ASSERT). Scoped and restored so no other torch/JIT
-    user is affected; the modern profiling executor stays on (this is NOT the legacy executor).
-    """
-    torch._C._jit_set_texpr_fuser_enabled(False)
-    try:
-        yield
-    finally:
-        torch._C._jit_set_texpr_fuser_enabled(True)
-
-
-class _ImpactCore(IMPACTReg):
-    """One IMPACT feature model, exposed as a FireANTs ``forward(moved, fixed)``.
-
-    Reuses ``IMPACTReg._compute`` / ``preprocessing`` verbatim — the stats-normalised feature extraction
-    (the model wants per-image ``[min, mean, max, std]``) and the per-layer weighted distance — so the
-    metric is exactly KonfAI's, not a re-derivation. Only KonfAI's config-binding ``__init__`` and its
-    ``Attribute``-based geometry are replaced: FireANTs passes raw tensors at the current pyramid scale, so
-    the intensity statistics are computed from those tensors directly. ``pca`` (absent from KonfAI's torch
-    ``IMPACTReg``) is added here as a per-layer feature-space reduction matching itk-impact.
-    """
-
-    def __init__(self, ref: str, in_channels: int, weights: list[float], distance: str, pca: int) -> None:
-        from huggingface_hub import hf_hub_download
-
-        torch.nn.Module.__init__(self)  # bypass IMPACTReg.__init__ (KONFAI_CONFIG_PATH / apply_config binding)
-        self.name = "Reg"
-        self.in_channels = int(in_channels)
-        self.weights = [float(w) for w in weights]
-        self.nb_layer = len(self.weights)
-        self.loss = _DISTANCES[distance]()
-        self.pca = int(pca)  # PCA lives in KonfAI's IMPACTReg._compute (same behaviour as itk-impact)
-        self.dim = DIM
-        self.shape = None  # score the whole (downsampled) tensor — no ModelPatch tiling
-        if ":" in ref:  # a "repo:path" HF reference; otherwise a local model file
-            repo, filename = ref.split(":", 1)
-            self.model_path = hf_hub_download(repo, filename, repo_type="model")  # nosec B615
-        else:
-            self.model_path = ref
-        self.model = None  # lazy-loaded on the first forward, like IMPACTReg
-
-    @staticmethod
-    def _stats(tensor: torch.Tensor) -> dict:
-        detached = tensor.detach()
-        return {
-            "ImageMin": float(detached.min()),
-            "ImageMean": float(detached.mean()),
-            "ImageMax": float(detached.max()),
-            "ImageStd": float(detached.std()),
-        }
-
-    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
-        if self.model is None:
-            self.model = torch.jit.load(self.model_path)  # nosec B614
-        self.model.to(moved.device).eval()
-        with _no_texpr_fuser():
-            loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None)
-        return loss / max(true_nb, 1)
-
-
-class ImpactFeatureLoss(torch.nn.Module):
-    """FireANTs ``custom_loss`` = the KonfAI IMPACT metric fused over several feature models.
-
-    ``forward(moved, fixed)`` sums each model's ``layers_weight * IMPACT(model)``. A model's per-layer
-    weights come from its ``layers_mask`` bitmask; its input channel count is read from the registry
-    (``models.json`` ``numberofchannels``) so it never has to be configured by hand.
-    """
-
-    def __init__(self, specs: list["ModelSpec"]) -> None:
-        super().__init__()
-        registry = load_models_registry()
-        self._cores = torch.nn.ModuleList()
-        self._model_weights: list[float] = []
-        for spec in specs:
-            in_channels = int(registry.get(spec.ref.split(":", 1)[-1], {}).get("numberofchannels", 1))
-            weights = [1.0 if char == "1" else 0.0 for char in spec.layers_mask]
-            self._cores.append(_ImpactCore(spec.ref, in_channels, weights, spec.distance, spec.pca))
-            self._model_weights.append(float(spec.layers_weight))
-
-    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:
-        total: torch.Tensor | None = None
-        for weight, core in zip(self._model_weights, self._cores, strict=True):
-            term = weight * core(moved, fixed)
-            total = term if total is None else total + term
-        return total
-
-
-class FireANTsEngine:
-    """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return
-    (moved, dvf) on the fixed grid.
-
-    ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config
-    /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet``
-    annotations) on a machine without a GPU or without FireANTs installed.
-    """
-
-    def __init__(
-        self,
-        scales: list[int],
-        affine_iterations: list[int],
-        deformable_iterations: list[int],
-        cc_kernel: int,
-        affine_metric: str,
-        affine_lr: float,
-        deformable_method: str,
-        deformable_metric: str,
-        deformable_lr: float,
-        integrator_n: int,
-        smooth_warp_sigma: float,
-        smooth_grad_sigma: float,
-        seed: int,
-        impact_specs: list["ModelSpec"],
-    ) -> None:
-        self._scales = [int(s) for s in scales]
-        self._affine_iterations = [int(i) for i in affine_iterations]
-        self._deformable_iterations = [int(i) for i in deformable_iterations]
-        self._cc_kernel = int(cc_kernel)
-        self._affine_metric = affine_metric
-        self._affine_lr = float(affine_lr)
-        self._deformable_method = deformable_method
-        self._deformable_metric = deformable_metric
-        self._deformable_lr = float(deformable_lr)
-        self._integrator_n = int(integrator_n)
-        self._smooth_warp_sigma = float(smooth_warp_sigma)
-        self._smooth_grad_sigma = float(smooth_grad_sigma)
-        self._seed = int(seed)
-        # IMPACT deformable metric (only used when deformable_metric == "impact"): KonfAI IMPACT feature
-        # models drive the SyN/greedy stage instead of the analytic CC/MI/MSE.
-        self._impact_specs = impact_specs
-
-    @staticmethod
-    def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-        """True only for a mask that actually restricts the region — some voxels in, some out. An absent
-        optional mask arrives as a whole-image (all-ones) default and an all-zero mask is degenerate; both
-        are treated as no mask so the plain (non-masked) metric path is used."""
-        if mask is None:
-            return False
-        arr = sitk.GetArrayViewFromImage(mask)
-        return bool((arr > 0).any()) and bool((arr == 0).any())
-
-    @staticmethod
-    def _affine_to_sitk(affine_matrix: "torch.Tensor") -> sitk.AffineTransform:
-        """FireANTs' physical (LPS) linear matrix -> SimpleITK AffineTransform (fixed -> moving points),
-        the same convention FireANTs writes into an ANTs ``0GenericAffine.mat``."""
-        matrix = affine_matrix.float().cpu().numpy()[0]
-        affine = sitk.AffineTransform(DIM)
-        affine.SetMatrix(matrix[:DIM, :DIM].flatten().astype(np.float64))
-        affine.SetTranslation(matrix[:DIM, DIM].astype(np.float64))
-        return affine
-
-    def _total_field_transform(self, reg) -> sitk.Transform:
-        """Optimise a deformable stage and return its TOTAL displacement (affine baked in) as a
-        SimpleITK ``DisplacementFieldTransform`` on the fixed grid.
-
-        FireANTs serialises the total field (ANTs convention, fixed grid) only to a file, so it is
-        round-tripped through a temporary NIfTI — its public API, no internals reimplemented."""
-        reg.optimize()
-        with tempfile.TemporaryDirectory() as tmp:
-            warp_path = os.path.join(tmp, "total_warp.nii.gz")
-            reg.save_as_ants_transforms(warp_path)
-            total_field = sitk.ReadImage(warp_path, sitk.sitkVectorFloat64)
-        return sitk.DisplacementFieldTransform(total_field)  # consumes total_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."""
-        ensure_fireants_runtime()
-        from fireants.io import BatchedImages, Image
-        from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones
-        from fireants.registration.affine import AffineRegistration
-        from fireants.registration.rigid import RigidRegistration
-
-        torch.manual_seed(self._seed)
-        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
-        # FireANTs' Image ctor accepts a SimpleITK image directly, so the fixed/moving cross into
-        # FireANTs in-memory (no file load) with their geometry preserved.
-        fixed_img = Image(fixed, device=device)
-        moving_img = Image(moving, device=device)
-
-        # Masked metric only when a mask genuinely restricts the region. FireANTs' masked mode wants the
-        # mask as the last channel of BOTH images (all-ones where one side has none) and a ``masked_``
-        # metric prefix; the plain path is untouched when no real mask is present.
-        use_fixed_mask = self._is_partial_mask(fixed_mask)
-        use_moving_mask = self._is_partial_mask(moving_mask)
-        masked = use_fixed_mask or use_moving_mask
-        if masked:
-            fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img)
-            mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img)
-            fixed_img = apply_mask_to_image(fixed_img, fmask)
-            moving_img = apply_mask_to_image(moving_img, mmask)
-
-        bf = BatchedImages([fixed_img])
-        bm = BatchedImages([moving_img])
-        affine_loss = f"masked_{self._affine_metric}" if masked else self._affine_metric
-        deformable_loss = f"masked_{self._deformable_metric}" if masked else self._deformable_metric
-
-        # Linear: Rigid(MI, COM init) -> Affine(MI, seeded by the rigid), mirroring ANTs. The affine
-        # seeds the deformable stage (or is the whole transform when deformable_method == "none").
-        rigid = RigidRegistration(
-            scales=self._scales,
-            iterations=self._affine_iterations,
-            fixed_images=bf,
-            moving_images=bm,
-            loss_type=affine_loss,
-            optimizer="Adam",
-            optimizer_lr=self._affine_lr,
-            cc_kernel_size=self._cc_kernel,
-            init_translation="cof",
-        )
-        rigid.optimize()
-        rigid_matrix = rigid.get_rigid_matrix().detach()
-
-        affine = AffineRegistration(
-            scales=self._scales,
-            iterations=self._affine_iterations,
-            fixed_images=bf,
-            moving_images=bm,
-            loss_type=affine_loss,
-            optimizer="Adam",
-            optimizer_lr=self._affine_lr,
-            cc_kernel_size=self._cc_kernel,
-            init_rigid=rigid_matrix,
-        )
-        affine.optimize()
-        affine_matrix = affine.get_affine_matrix().detach()
-
-        # Deformable stage (or none). SyN and Greedy share the same constructor surface; both warm-start
-        # from the affine so their TOTAL transform already bakes in the linear pre-align.
-        if self._deformable_method == "none":
-            transform: sitk.Transform = self._affine_to_sitk(affine_matrix)
-        else:
-            if self._deformable_method == "syn":
-                from fireants.registration.syn import SyNRegistration as Deformable
-            elif self._deformable_method == "greedy":
-                from fireants.registration.greedy import GreedyRegistration as Deformable
-            else:
-                raise ValueError(
-                    f"Unknown deformable_method '{self._deformable_method}' (expected 'syn', 'greedy' or 'none')."
-                )
-            # "impact" swaps the analytic metric for a KonfAI IMPACT feature loss on the deformable stage
-            # (the linear pre-align keeps its own affine_metric); masks do not restrict the IMPACT metric.
-            if self._deformable_metric == "impact":
-                loss_type: str = "custom"
-                custom_loss: torch.nn.Module | None = ImpactFeatureLoss(self._impact_specs)
-            else:
-                loss_type, custom_loss = deformable_loss, None
-            reg = Deformable(
-                scales=self._scales,
-                iterations=self._deformable_iterations,
-                fixed_images=bf,
-                moving_images=bm,
-                loss_type=loss_type,
-                custom_loss=custom_loss,
-                cc_kernel_size=self._cc_kernel,
-                deformation_type="compositive",
-                integrator_n=self._integrator_n,
-                smooth_warp_sigma=self._smooth_warp_sigma,
-                smooth_grad_sigma=self._smooth_grad_sigma,
-                optimizer="Adam",
-                optimizer_lr=self._deformable_lr,
-                init_affine=affine_matrix,
-            )
-            transform = self._total_field_transform(reg)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-
-        # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path,
-        # so every FireANTs preset emits identical-shaped results.
-        moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-        dvf = sitk.TransformToDisplacementField(
-            transform,
-            sitk.sitkVectorFloat64,
-            fixed.GetSize(),
-            fixed.GetOrigin(),
-            fixed.GetSpacing(),
-            fixed.GetDirection(),
-        )
-        moved_np, _ = image_to_data(moved)
-        dvf_np, _ = image_to_data(dvf)
-        return moved_np, dvf_np
-
-
-class FireANTsRegistration(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 the ConvexAdam / elastix engines); registration needs the
-    physical geometry, and the mask branches restrict the metric.
-    """
-
-    accepts_attributes = True
-
-    def __init__(self, engine: FireANTsEngine) -> 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. A whole-image
-        # mask (the default when none is supplied) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        # FireANTs runs a gradient-based instance optimisation (Riemannian Adam over the warp); the
-        # predictor calls forward under torch.inference_mode(), which forbids autograd. The image tensors
-        # have already crossed to numpy/SimpleITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad():
-            for b in range(fixed.shape[0]):
-                fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-                moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-                fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-                moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-                moved_np, dvf_np = self._engine.register(
-                    fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-                )
-                combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise FireANTs registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = 2, moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    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``. The knobs below are read straight from these annotations by the
-    UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method``
-    is the knob that specialises this shared model into each FireANTs preset.
-    """
-
-    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()},
-        scales: list[int] = [4, 2, 1],
-        affine_iterations: list[int] = [200, 100, 50],
-        deformable_iterations: list[int] = [200, 100, 50],
-        cc_kernel: Annotated[int, Range(1, 21)] = 5,
-        affine_metric: Literal["mi", "cc", "mse"] = "mi",
-        affine_lr: Annotated[float, Range(0.0, 10.0)] = 0.003,
-        deformable_method: Literal["none", "syn", "greedy"] = "syn",
-        deformable_metric: Literal["cc", "mi", "mse", "impact"] = "cc",
-        deformable_lr: Annotated[float, Range(0.0, 10.0)] = 0.25,
-        integrator_n: Annotated[int, Range(1, 100)] = 10,
-        smooth_warp_sigma: Annotated[float, Range(0.0, 100.0)] = 0.5,
-        smooth_grad_sigma: Annotated[float, Range(0.0, 100.0)] = 1.0,
-        seed: int = 42,
-        models: dict[str, ModelSpec] = {},
-    ) -> None:
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        engine = FireANTsEngine(
-            scales,
-            affine_iterations,
-            deformable_iterations,
-            cc_kernel,
-            affine_metric,
-            affine_lr,
-            deformable_method,
-            deformable_metric,
-            deformable_lr,
-            integrator_n,
-            smooth_warp_sigma,
-            smooth_grad_sigma,
-            seed,
-            _sorted_specs(models),
-        )
-        self.add_module(
-            "Registration", FireANTsRegistration(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"])

FireANTs_Affine/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.fireants:RegistrationNet
     RegistrationNet:
       scales:
       - 4

FireANTs_Affine/requirements.txt

@@ -1,3 +1,4 @@
 hydra-core
 nibabel
 pandas
+impact_reg_konfai>=1.6.1

FireANTs_IMPACT/Model.py

@@ -1,669 +0,0 @@
-# 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
-#
-# This wrapper does NOT copy any FireANTs source: it only calls the public FireANTs API of the
-# separately-installed ``fireants`` wheel (PyPI). FireANTs is distributed under the FireANTs License
-# v1.0 and must be cited — see the NOTICE file in this directory for the license, copyright and
-# bibliography that ship with this app.
-
-"""FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets).
-
-Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset
-(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the
-orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own
-composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages:
-
-    Rigid (MI, centre-of-mass init) -> Affine (MI, seeded by the rigid) -> deformable
-
-The deformable stage is selected by ``deformable_method`` — the ONE knob that specialises this shared
-Model.py into the different presets (exactly as ConvexAdam's shared Model.py is specialised by
-``stages``):
-
-    "syn"    symmetric diffeomorphic SyN (CC)   — invertible, higher quality, averages cleanly for ensembling
-    "greedy" greedy diffeomorphic (CC)          — one-directional, faster / lower VRAM
-    "none"   linear only                        — Rigid+Affine, no deformable (the FireANTs_Affine preset)
-
-Masks: the optional Fixed/Moving masks restrict the metric to a region. FireANTs implements this by
-carrying the mask as the last image channel and prefixing the metric with ``masked_``; a mask is only
-honoured when it actually restricts (some voxels in, some out), so the common mask-free path is
-unchanged (an absent optional mask arrives as a whole-image default and is treated as no mask).
-
-The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear
-pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly.
-``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK —
-the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an
-ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is
-round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here).
-
-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 contextlib
-import json
-import os
-import tempfile
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Annotated, Literal
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-from konfai.metric.measure import IMPACTReg
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-DIM = 3
-
-# Feature-model registry (models.json): the available IMPACT feature models, fetched from HF (NOT bundled).
-# Only consulted by the "impact" deformable metric; ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins
-# for dev/offline. Mirrors the ConvexAdam preset so the same 30-model catalogue and picker are shared.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-_DISTANCES: dict[str, type[torch.nn.Module]] = {"L1": torch.nn.L1Loss, "L2": torch.nn.MSELoss}
-
-
-def _fireants_git_ref() -> str:
-    """Best-effort FireANTs git ref whose ``fused_ops`` matches the installed ``fireants``.
-
-    Overridable with ``FIREANTS_FUSED_OPS_REF``; falls back to ``main`` if the version is unknown.
-    """
-    try:
-        import importlib.metadata
-
-        version = importlib.metadata.version("fireants").strip()
-        if version:
-            return version if version.startswith("v") else f"v{version}"
-    except Exception:
-        pass
-    return "main"
-
-
-def ensure_fireants_runtime(build_kernels: bool = True) -> None:
-    """Make the ``fireants`` runtime importable before a preset uses it — best-effort, never fatal.
-
-    A plain ``pip install fireants`` fails inside a host like 3D Slicer for two reasons, both handled
-    here with a clear one-line status at each step:
-
-    1. **fireants won't install.** It pins ``simpleitk==2.2.1`` (no wheel on modern Python), while the
-       host already ships a newer SimpleITK. We install it with ``--no-deps`` so that pin is ignored
-       and the host's SimpleITK/torch are reused; its light deps ship in ``requirements.txt``.
-    2. **The fused CUDA kernels** (``fireants_fused_ops``) that make registration fast and
-       memory-light are OPTIONAL. Without them fireants runs in pure PyTorch (correct, only slower).
-       We enable them only when a CUDA compiler (``nvcc``) is present, compiling from the upstream
-       FireANTs source at install time — nothing is vendored into this app. No compiler, or a failed
-       build, simply falls back to pure PyTorch.
-    """
-    import importlib
-    import shutil
-    import subprocess
-    import sys
-
-    def _log(message: str) -> None:
-        print(f"[FireANTs] {message}", flush=True)
-
-    # 1) fireants itself — install --no-deps to sidestep its unsatisfiable ``simpleitk==2.2.1`` pin.
-    try:
-        importlib.import_module("fireants")
-    except Exception:
-        _log("installing fireants (--no-deps, reusing the host's SimpleITK/torch)...")
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--no-deps", "fireants"])
-            importlib.invalidate_caches()
-            importlib.import_module("fireants")
-            _log("fireants installed.")
-        except Exception as exc:
-            raise RuntimeError(
-                "Could not install 'fireants'. Install it manually with:\n"
-                f"    {sys.executable} -m pip install --no-deps fireants\n"
-                "(its hydra-core/nibabel/pandas dependencies ship in this app's requirements.txt).\n"
-                f"Original error: {exc}"
-            ) from exc
-
-    if not build_kernels or os.environ.get("FIREANTS_SKIP_FUSED_OPS", "").strip().lower() in {"1", "true", "yes"}:
-        _log("skipping the fused CUDA kernels -> pure PyTorch (correct, slower and more memory).")
-        return
-
-    # 2) fused CUDA kernels — optional accelerator.
-    try:
-        importlib.import_module("fireants_fused_ops")
-        _log("fused CUDA kernels already available.")
-        return
-    except Exception:
-        pass
-
-    # 2a) a prebuilt wheel matching this torch's CUDA build, if one is ever published.
-    cuda_tag = ""
-    try:
-        cuda_tag = (torch.version.cuda or "").replace(".", "")
-    except Exception:
-        pass
-    if cuda_tag:
-        wheel = f"fireants-fused-ops-cu{cuda_tag}"
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", wheel])
-            importlib.invalidate_caches()
-            importlib.import_module("fireants_fused_ops")
-            _log(f"installed prebuilt CUDA kernels ({wheel}).")
-            return
-        except Exception:
-            pass  # no matching wheel -> try a local build
-
-    # 2b) local build is OPT-IN: compiling CUDA kernels is heavy and can exhaust RAM on the user's
-    # machine, so it NEVER runs by default. Set FIREANTS_BUILD_KERNELS=1 to enable it (devs with a
-    # CUDA toolkit); it then builds ONE file at a time (MAX_JOBS=1) to keep memory bounded. The clean
-    # path for end users is a prebuilt wheel (2a) — a local compile also needs Python dev headers
-    # (absent from some embedded Pythons, e.g. Slicer) and a CUDA-compatible host compiler.
-    if os.environ.get("FIREANTS_BUILD_KERNELS", "").strip().lower() not in ("1", "true", "yes"):
-        _log(
-            "no prebuilt kernels for this platform -> running FireANTs in pure PyTorch (correct, slower "
-            "and heavier). Set FIREANTS_BUILD_KERNELS=1 to compile them locally (needs a CUDA toolkit)."
-        )
-        return
-    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
-    nvcc = shutil.which("nvcc") or (os.path.join(cuda_home, "bin", "nvcc") if cuda_home else None)
-    if not nvcc or not os.path.exists(nvcc) or shutil.which("git") is None:
-        _log("FIREANTS_BUILD_KERNELS set but no CUDA toolkit (nvcc) / git found -> pure PyTorch.")
-        return
-
-    ref = os.environ.get("FIREANTS_FUSED_OPS_REF") or _fireants_git_ref()
-    _log(f"nvcc found ({nvcc}); compiling the FireANTs CUDA kernels one file at a time (ref '{ref}')...")
-    import tempfile
-
-    env = os.environ.copy()
-    env["MAX_JOBS"] = "1"  # one compile at a time -> bounded RAM (prevents OOM on large hosts)
-    env.setdefault("NVCC_APPEND_FLAGS", "-allow-unsupported-compiler")  # tolerate a newer host compiler
-    tmp = tempfile.mkdtemp(prefix="fireants_fused_ops_")
-    try:
-        # Shallow clone WITHOUT --recursive: FireANTs' submodules (an SSH-only 'cookbook') are unrelated
-        # to fused_ops and would otherwise abort the build with a public-key/permission error.
-        subprocess.check_call(
-            ["git", "clone", "--depth", "1", "--branch", ref, "https://github.com/rohitrango/FireANTs.git", tmp]
-        )
-        subprocess.check_call(
-            [sys.executable, "-m", "pip", "install", "-q", "--no-build-isolation", os.path.join(tmp, "fused_ops")],
-            env=env,
-        )
-        importlib.invalidate_caches()
-        importlib.import_module("fireants_fused_ops")
-        _log("compiled and installed the fast CUDA kernels.")
-    except Exception as exc:
-        _log(
-            "kernel build failed -> running in pure PyTorch (correct, only speed/memory affected). Set "
-            "FIREANTS_FUSED_OPS_REF to a compatible FireANTs tag, or install a prebuilt fireants-fused-ops "
-            f"wheel, to enable them. Details: {exc}"
-        )
-    finally:
-        shutil.rmtree(tmp, ignore_errors=True)
-
-
-def registry_choices() -> list[str]:
-    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load ``models.json`` (available feature models). ``KONFAI_IMPACT_MODELS_REGISTRY`` (local path) wins
-    for dev/offline; otherwise ``ref`` is a ``repo:file`` Hugging Face reference (fetched, not bundled)."""
-    from huggingface_hub import hf_hub_download
-
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference — or set "
-            "KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-@dataclass
-class ModelSpec:
-    """One IMPACT feature model in the deformable metric (several are fused). ``ref`` picks the model; the
-    rest are its per-model knobs — the same as the ConvexAdam / elastix ``ModelSpec`` except ``voxel_size``
-    (an itk-impact resampling knob) has no meaning for FireANTs' geometry-free torch ``custom_loss`` and is
-    intentionally absent."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    layers_mask: str = "01"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
-    layers_weight: float = 1.0  # this model's weight in the multi-model fusion
-    pca: Annotated[int, Range(0, 100)] = 0  # keep the top-K principal components of the features (0 = no PCA)
-    distance: Literal["L1", "L2"] = "L1"
-
-
-@contextlib.contextmanager
-def _no_texpr_fuser():
-    """Disable the TensorExpr JIT fuser while IMPACT's TorchScript feature model runs under autograd.
-
-    The IMPACT feature models are TorchScript; run under FireANTs' gradient optimisation the TensorExpr
-    fuser trips on shape ops (``aten::size`` INTERNAL ASSERT). Scoped and restored so no other torch/JIT
-    user is affected; the modern profiling executor stays on (this is NOT the legacy executor).
-    """
-    torch._C._jit_set_texpr_fuser_enabled(False)
-    try:
-        yield
-    finally:
-        torch._C._jit_set_texpr_fuser_enabled(True)
-
-
-class _ImpactCore(IMPACTReg):
-    """One IMPACT feature model, exposed as a FireANTs ``forward(moved, fixed)``.
-
-    Reuses ``IMPACTReg._compute`` / ``preprocessing`` verbatim — the stats-normalised feature extraction
-    (the model wants per-image ``[min, mean, max, std]``) and the per-layer weighted distance — so the
-    metric is exactly KonfAI's, not a re-derivation. Only KonfAI's config-binding ``__init__`` and its
-    ``Attribute``-based geometry are replaced: FireANTs passes raw tensors at the current pyramid scale, so
-    the intensity statistics are computed from those tensors directly. ``pca`` (absent from KonfAI's torch
-    ``IMPACTReg``) is added here as a per-layer feature-space reduction matching itk-impact.
-    """
-
-    def __init__(self, ref: str, in_channels: int, weights: list[float], distance: str, pca: int) -> None:
-        from huggingface_hub import hf_hub_download
-
-        torch.nn.Module.__init__(self)  # bypass IMPACTReg.__init__ (KONFAI_CONFIG_PATH / apply_config binding)
-        self.name = "Reg"
-        self.in_channels = int(in_channels)
-        self.weights = [float(w) for w in weights]
-        self.nb_layer = len(self.weights)
-        self.loss = _DISTANCES[distance]()
-        self.pca = int(pca)  # PCA lives in KonfAI's IMPACTReg._compute (same behaviour as itk-impact)
-        self.dim = DIM
-        self.shape = None  # score the whole (downsampled) tensor — no ModelPatch tiling
-        if ":" in ref:  # a "repo:path" HF reference; otherwise a local model file
-            repo, filename = ref.split(":", 1)
-            self.model_path = hf_hub_download(repo, filename, repo_type="model")  # nosec B615
-        else:
-            self.model_path = ref
-        self.model = None  # lazy-loaded on the first forward, like IMPACTReg
-
-    @staticmethod
-    def _stats(tensor: torch.Tensor) -> dict:
-        detached = tensor.detach()
-        return {
-            "ImageMin": float(detached.min()),
-            "ImageMean": float(detached.mean()),
-            "ImageMax": float(detached.max()),
-            "ImageStd": float(detached.std()),
-        }
-
-    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
-        if self.model is None:
-            self.model = torch.jit.load(self.model_path)  # nosec B614
-        self.model.to(moved.device).eval()
-        with _no_texpr_fuser():
-            loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None)
-        return loss / max(true_nb, 1)
-
-
-class ImpactFeatureLoss(torch.nn.Module):
-    """FireANTs ``custom_loss`` = the KonfAI IMPACT metric fused over several feature models.
-
-    ``forward(moved, fixed)`` sums each model's ``layers_weight * IMPACT(model)``. A model's per-layer
-    weights come from its ``layers_mask`` bitmask; its input channel count is read from the registry
-    (``models.json`` ``numberofchannels``) so it never has to be configured by hand.
-    """
-
-    def __init__(self, specs: list["ModelSpec"]) -> None:
-        super().__init__()
-        registry = load_models_registry()
-        self._cores = torch.nn.ModuleList()
-        self._model_weights: list[float] = []
-        for spec in specs:
-            in_channels = int(registry.get(spec.ref.split(":", 1)[-1], {}).get("numberofchannels", 1))
-            weights = [1.0 if char == "1" else 0.0 for char in spec.layers_mask]
-            self._cores.append(_ImpactCore(spec.ref, in_channels, weights, spec.distance, spec.pca))
-            self._model_weights.append(float(spec.layers_weight))
-
-    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:
-        total: torch.Tensor | None = None
-        for weight, core in zip(self._model_weights, self._cores, strict=True):
-            term = weight * core(moved, fixed)
-            total = term if total is None else total + term
-        return total
-
-
-class FireANTsEngine:
-    """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return
-    (moved, dvf) on the fixed grid.
-
-    ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config
-    /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet``
-    annotations) on a machine without a GPU or without FireANTs installed.
-    """
-
-    def __init__(
-        self,
-        scales: list[int],
-        affine_iterations: list[int],
-        deformable_iterations: list[int],
-        cc_kernel: int,
-        affine_metric: str,
-        affine_lr: float,
-        deformable_method: str,
-        deformable_metric: str,
-        deformable_lr: float,
-        integrator_n: int,
-        smooth_warp_sigma: float,
-        smooth_grad_sigma: float,
-        seed: int,
-        impact_specs: list["ModelSpec"],
-    ) -> None:
-        self._scales = [int(s) for s in scales]
-        self._affine_iterations = [int(i) for i in affine_iterations]
-        self._deformable_iterations = [int(i) for i in deformable_iterations]
-        self._cc_kernel = int(cc_kernel)
-        self._affine_metric = affine_metric
-        self._affine_lr = float(affine_lr)
-        self._deformable_method = deformable_method
-        self._deformable_metric = deformable_metric
-        self._deformable_lr = float(deformable_lr)
-        self._integrator_n = int(integrator_n)
-        self._smooth_warp_sigma = float(smooth_warp_sigma)
-        self._smooth_grad_sigma = float(smooth_grad_sigma)
-        self._seed = int(seed)
-        # IMPACT deformable metric (only used when deformable_metric == "impact"): KonfAI IMPACT feature
-        # models drive the SyN/greedy stage instead of the analytic CC/MI/MSE.
-        self._impact_specs = impact_specs
-
-    @staticmethod
-    def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-        """True only for a mask that actually restricts the region — some voxels in, some out. An absent
-        optional mask arrives as a whole-image (all-ones) default and an all-zero mask is degenerate; both
-        are treated as no mask so the plain (non-masked) metric path is used."""
-        if mask is None:
-            return False
-        arr = sitk.GetArrayViewFromImage(mask)
-        return bool((arr > 0).any()) and bool((arr == 0).any())
-
-    @staticmethod
-    def _affine_to_sitk(affine_matrix: "torch.Tensor") -> sitk.AffineTransform:
-        """FireANTs' physical (LPS) linear matrix -> SimpleITK AffineTransform (fixed -> moving points),
-        the same convention FireANTs writes into an ANTs ``0GenericAffine.mat``."""
-        matrix = affine_matrix.float().cpu().numpy()[0]
-        affine = sitk.AffineTransform(DIM)
-        affine.SetMatrix(matrix[:DIM, :DIM].flatten().astype(np.float64))
-        affine.SetTranslation(matrix[:DIM, DIM].astype(np.float64))
-        return affine
-
-    def _total_field_transform(self, reg) -> sitk.Transform:
-        """Optimise a deformable stage and return its TOTAL displacement (affine baked in) as a
-        SimpleITK ``DisplacementFieldTransform`` on the fixed grid.
-
-        FireANTs serialises the total field (ANTs convention, fixed grid) only to a file, so it is
-        round-tripped through a temporary NIfTI — its public API, no internals reimplemented."""
-        reg.optimize()
-        with tempfile.TemporaryDirectory() as tmp:
-            warp_path = os.path.join(tmp, "total_warp.nii.gz")
-            reg.save_as_ants_transforms(warp_path)
-            total_field = sitk.ReadImage(warp_path, sitk.sitkVectorFloat64)
-        return sitk.DisplacementFieldTransform(total_field)  # consumes total_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."""
-        ensure_fireants_runtime()
-        from fireants.io import BatchedImages, Image
-        from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones
-        from fireants.registration.affine import AffineRegistration
-        from fireants.registration.rigid import RigidRegistration
-
-        torch.manual_seed(self._seed)
-        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
-        # FireANTs' Image ctor accepts a SimpleITK image directly, so the fixed/moving cross into
-        # FireANTs in-memory (no file load) with their geometry preserved.
-        fixed_img = Image(fixed, device=device)
-        moving_img = Image(moving, device=device)
-
-        # Masked metric only when a mask genuinely restricts the region. FireANTs' masked mode wants the
-        # mask as the last channel of BOTH images (all-ones where one side has none) and a ``masked_``
-        # metric prefix; the plain path is untouched when no real mask is present.
-        use_fixed_mask = self._is_partial_mask(fixed_mask)
-        use_moving_mask = self._is_partial_mask(moving_mask)
-        masked = use_fixed_mask or use_moving_mask
-        if masked:
-            fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img)
-            mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img)
-            fixed_img = apply_mask_to_image(fixed_img, fmask)
-            moving_img = apply_mask_to_image(moving_img, mmask)
-
-        bf = BatchedImages([fixed_img])
-        bm = BatchedImages([moving_img])
-        affine_loss = f"masked_{self._affine_metric}" if masked else self._affine_metric
-        deformable_loss = f"masked_{self._deformable_metric}" if masked else self._deformable_metric
-
-        # Linear: Rigid(MI, COM init) -> Affine(MI, seeded by the rigid), mirroring ANTs. The affine
-        # seeds the deformable stage (or is the whole transform when deformable_method == "none").
-        rigid = RigidRegistration(
-            scales=self._scales,
-            iterations=self._affine_iterations,
-            fixed_images=bf,
-            moving_images=bm,
-            loss_type=affine_loss,
-            optimizer="Adam",
-            optimizer_lr=self._affine_lr,
-            cc_kernel_size=self._cc_kernel,
-            init_translation="cof",
-        )
-        rigid.optimize()
-        rigid_matrix = rigid.get_rigid_matrix().detach()
-
-        affine = AffineRegistration(
-            scales=self._scales,
-            iterations=self._affine_iterations,
-            fixed_images=bf,
-            moving_images=bm,
-            loss_type=affine_loss,
-            optimizer="Adam",
-            optimizer_lr=self._affine_lr,
-            cc_kernel_size=self._cc_kernel,
-            init_rigid=rigid_matrix,
-        )
-        affine.optimize()
-        affine_matrix = affine.get_affine_matrix().detach()
-
-        # Deformable stage (or none). SyN and Greedy share the same constructor surface; both warm-start
-        # from the affine so their TOTAL transform already bakes in the linear pre-align.
-        if self._deformable_method == "none":
-            transform: sitk.Transform = self._affine_to_sitk(affine_matrix)
-        else:
-            if self._deformable_method == "syn":
-                from fireants.registration.syn import SyNRegistration as Deformable
-            elif self._deformable_method == "greedy":
-                from fireants.registration.greedy import GreedyRegistration as Deformable
-            else:
-                raise ValueError(
-                    f"Unknown deformable_method '{self._deformable_method}' (expected 'syn', 'greedy' or 'none')."
-                )
-            # "impact" swaps the analytic metric for a KonfAI IMPACT feature loss on the deformable stage
-            # (the linear pre-align keeps its own affine_metric); masks do not restrict the IMPACT metric.
-            if self._deformable_metric == "impact":
-                loss_type: str = "custom"
-                custom_loss: torch.nn.Module | None = ImpactFeatureLoss(self._impact_specs)
-            else:
-                loss_type, custom_loss = deformable_loss, None
-            reg = Deformable(
-                scales=self._scales,
-                iterations=self._deformable_iterations,
-                fixed_images=bf,
-                moving_images=bm,
-                loss_type=loss_type,
-                custom_loss=custom_loss,
-                cc_kernel_size=self._cc_kernel,
-                deformation_type="compositive",
-                integrator_n=self._integrator_n,
-                smooth_warp_sigma=self._smooth_warp_sigma,
-                smooth_grad_sigma=self._smooth_grad_sigma,
-                optimizer="Adam",
-                optimizer_lr=self._deformable_lr,
-                init_affine=affine_matrix,
-            )
-            transform = self._total_field_transform(reg)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-
-        # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path,
-        # so every FireANTs preset emits identical-shaped results.
-        moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-        dvf = sitk.TransformToDisplacementField(
-            transform,
-            sitk.sitkVectorFloat64,
-            fixed.GetSize(),
-            fixed.GetOrigin(),
-            fixed.GetSpacing(),
-            fixed.GetDirection(),
-        )
-        moved_np, _ = image_to_data(moved)
-        dvf_np, _ = image_to_data(dvf)
-        return moved_np, dvf_np
-
-
-class FireANTsRegistration(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 the ConvexAdam / elastix engines); registration needs the
-    physical geometry, and the mask branches restrict the metric.
-    """
-
-    accepts_attributes = True
-
-    def __init__(self, engine: FireANTsEngine) -> 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. A whole-image
-        # mask (the default when none is supplied) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        # FireANTs runs a gradient-based instance optimisation (Riemannian Adam over the warp); the
-        # predictor calls forward under torch.inference_mode(), which forbids autograd. The image tensors
-        # have already crossed to numpy/SimpleITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad():
-            for b in range(fixed.shape[0]):
-                fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-                moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-                fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-                moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-                moved_np, dvf_np = self._engine.register(
-                    fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-                )
-                combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise FireANTs registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = 2, moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    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``. The knobs below are read straight from these annotations by the
-    UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method``
-    is the knob that specialises this shared model into each FireANTs preset.
-    """
-
-    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()},
-        scales: list[int] = [4, 2, 1],
-        affine_iterations: list[int] = [200, 100, 50],
-        deformable_iterations: list[int] = [200, 100, 50],
-        cc_kernel: Annotated[int, Range(1, 21)] = 5,
-        affine_metric: Literal["mi", "cc", "mse"] = "mi",
-        affine_lr: Annotated[float, Range(0.0, 10.0)] = 0.003,
-        deformable_method: Literal["none", "syn", "greedy"] = "syn",
-        deformable_metric: Literal["cc", "mi", "mse", "impact"] = "cc",
-        deformable_lr: Annotated[float, Range(0.0, 10.0)] = 0.25,
-        integrator_n: Annotated[int, Range(1, 100)] = 10,
-        smooth_warp_sigma: Annotated[float, Range(0.0, 100.0)] = 0.5,
-        smooth_grad_sigma: Annotated[float, Range(0.0, 100.0)] = 1.0,
-        seed: int = 42,
-        models: dict[str, ModelSpec] = {},
-    ) -> None:
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        engine = FireANTsEngine(
-            scales,
-            affine_iterations,
-            deformable_iterations,
-            cc_kernel,
-            affine_metric,
-            affine_lr,
-            deformable_method,
-            deformable_metric,
-            deformable_lr,
-            integrator_n,
-            smooth_warp_sigma,
-            smooth_grad_sigma,
-            seed,
-            _sorted_specs(models),
-        )
-        self.add_module(
-            "Registration", FireANTsRegistration(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"])

FireANTs_IMPACT/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.fireants:RegistrationNet
     RegistrationNet:
       scales:
       - 4

FireANTs_IMPACT/requirements.txt

@@ -1,3 +1,4 @@
 hydra-core
 nibabel
 pandas
+impact_reg_konfai>=1.6.1

FireANTs_SyN/Model.py

@@ -1,669 +0,0 @@
-# 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
-#
-# This wrapper does NOT copy any FireANTs source: it only calls the public FireANTs API of the
-# separately-installed ``fireants`` wheel (PyPI). FireANTs is distributed under the FireANTs License
-# v1.0 and must be cited — see the NOTICE file in this directory for the license, copyright and
-# bibliography that ship with this app.
-
-"""FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets).
-
-Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset
-(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the
-orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own
-composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages:
-
-    Rigid (MI, centre-of-mass init) -> Affine (MI, seeded by the rigid) -> deformable
-
-The deformable stage is selected by ``deformable_method`` — the ONE knob that specialises this shared
-Model.py into the different presets (exactly as ConvexAdam's shared Model.py is specialised by
-``stages``):
-
-    "syn"    symmetric diffeomorphic SyN (CC)   — invertible, higher quality, averages cleanly for ensembling
-    "greedy" greedy diffeomorphic (CC)          — one-directional, faster / lower VRAM
-    "none"   linear only                        — Rigid+Affine, no deformable (the FireANTs_Affine preset)
-
-Masks: the optional Fixed/Moving masks restrict the metric to a region. FireANTs implements this by
-carrying the mask as the last image channel and prefixing the metric with ``masked_``; a mask is only
-honoured when it actually restricts (some voxels in, some out), so the common mask-free path is
-unchanged (an absent optional mask arrives as a whole-image default and is treated as no mask).
-
-The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear
-pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly.
-``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK —
-the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an
-ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is
-round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here).
-
-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 contextlib
-import json
-import os
-import tempfile
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Annotated, Literal
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-from konfai.metric.measure import IMPACTReg
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-DIM = 3
-
-# Feature-model registry (models.json): the available IMPACT feature models, fetched from HF (NOT bundled).
-# Only consulted by the "impact" deformable metric; ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins
-# for dev/offline. Mirrors the ConvexAdam preset so the same 30-model catalogue and picker are shared.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-_DISTANCES: dict[str, type[torch.nn.Module]] = {"L1": torch.nn.L1Loss, "L2": torch.nn.MSELoss}
-
-
-def _fireants_git_ref() -> str:
-    """Best-effort FireANTs git ref whose ``fused_ops`` matches the installed ``fireants``.
-
-    Overridable with ``FIREANTS_FUSED_OPS_REF``; falls back to ``main`` if the version is unknown.
-    """
-    try:
-        import importlib.metadata
-
-        version = importlib.metadata.version("fireants").strip()
-        if version:
-            return version if version.startswith("v") else f"v{version}"
-    except Exception:
-        pass
-    return "main"
-
-
-def ensure_fireants_runtime(build_kernels: bool = True) -> None:
-    """Make the ``fireants`` runtime importable before a preset uses it — best-effort, never fatal.
-
-    A plain ``pip install fireants`` fails inside a host like 3D Slicer for two reasons, both handled
-    here with a clear one-line status at each step:
-
-    1. **fireants won't install.** It pins ``simpleitk==2.2.1`` (no wheel on modern Python), while the
-       host already ships a newer SimpleITK. We install it with ``--no-deps`` so that pin is ignored
-       and the host's SimpleITK/torch are reused; its light deps ship in ``requirements.txt``.
-    2. **The fused CUDA kernels** (``fireants_fused_ops``) that make registration fast and
-       memory-light are OPTIONAL. Without them fireants runs in pure PyTorch (correct, only slower).
-       We enable them only when a CUDA compiler (``nvcc``) is present, compiling from the upstream
-       FireANTs source at install time — nothing is vendored into this app. No compiler, or a failed
-       build, simply falls back to pure PyTorch.
-    """
-    import importlib
-    import shutil
-    import subprocess
-    import sys
-
-    def _log(message: str) -> None:
-        print(f"[FireANTs] {message}", flush=True)
-
-    # 1) fireants itself — install --no-deps to sidestep its unsatisfiable ``simpleitk==2.2.1`` pin.
-    try:
-        importlib.import_module("fireants")
-    except Exception:
-        _log("installing fireants (--no-deps, reusing the host's SimpleITK/torch)...")
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--no-deps", "fireants"])
-            importlib.invalidate_caches()
-            importlib.import_module("fireants")
-            _log("fireants installed.")
-        except Exception as exc:
-            raise RuntimeError(
-                "Could not install 'fireants'. Install it manually with:\n"
-                f"    {sys.executable} -m pip install --no-deps fireants\n"
-                "(its hydra-core/nibabel/pandas dependencies ship in this app's requirements.txt).\n"
-                f"Original error: {exc}"
-            ) from exc
-
-    if not build_kernels or os.environ.get("FIREANTS_SKIP_FUSED_OPS", "").strip().lower() in {"1", "true", "yes"}:
-        _log("skipping the fused CUDA kernels -> pure PyTorch (correct, slower and more memory).")
-        return
-
-    # 2) fused CUDA kernels — optional accelerator.
-    try:
-        importlib.import_module("fireants_fused_ops")
-        _log("fused CUDA kernels already available.")
-        return
-    except Exception:
-        pass
-
-    # 2a) a prebuilt wheel matching this torch's CUDA build, if one is ever published.
-    cuda_tag = ""
-    try:
-        cuda_tag = (torch.version.cuda or "").replace(".", "")
-    except Exception:
-        pass
-    if cuda_tag:
-        wheel = f"fireants-fused-ops-cu{cuda_tag}"
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", wheel])
-            importlib.invalidate_caches()
-            importlib.import_module("fireants_fused_ops")
-            _log(f"installed prebuilt CUDA kernels ({wheel}).")
-            return
-        except Exception:
-            pass  # no matching wheel -> try a local build
-
-    # 2b) local build is OPT-IN: compiling CUDA kernels is heavy and can exhaust RAM on the user's
-    # machine, so it NEVER runs by default. Set FIREANTS_BUILD_KERNELS=1 to enable it (devs with a
-    # CUDA toolkit); it then builds ONE file at a time (MAX_JOBS=1) to keep memory bounded. The clean
-    # path for end users is a prebuilt wheel (2a) — a local compile also needs Python dev headers
-    # (absent from some embedded Pythons, e.g. Slicer) and a CUDA-compatible host compiler.
-    if os.environ.get("FIREANTS_BUILD_KERNELS", "").strip().lower() not in ("1", "true", "yes"):
-        _log(
-            "no prebuilt kernels for this platform -> running FireANTs in pure PyTorch (correct, slower "
-            "and heavier). Set FIREANTS_BUILD_KERNELS=1 to compile them locally (needs a CUDA toolkit)."
-        )
-        return
-    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
-    nvcc = shutil.which("nvcc") or (os.path.join(cuda_home, "bin", "nvcc") if cuda_home else None)
-    if not nvcc or not os.path.exists(nvcc) or shutil.which("git") is None:
-        _log("FIREANTS_BUILD_KERNELS set but no CUDA toolkit (nvcc) / git found -> pure PyTorch.")
-        return
-
-    ref = os.environ.get("FIREANTS_FUSED_OPS_REF") or _fireants_git_ref()
-    _log(f"nvcc found ({nvcc}); compiling the FireANTs CUDA kernels one file at a time (ref '{ref}')...")
-    import tempfile
-
-    env = os.environ.copy()
-    env["MAX_JOBS"] = "1"  # one compile at a time -> bounded RAM (prevents OOM on large hosts)
-    env.setdefault("NVCC_APPEND_FLAGS", "-allow-unsupported-compiler")  # tolerate a newer host compiler
-    tmp = tempfile.mkdtemp(prefix="fireants_fused_ops_")
-    try:
-        # Shallow clone WITHOUT --recursive: FireANTs' submodules (an SSH-only 'cookbook') are unrelated
-        # to fused_ops and would otherwise abort the build with a public-key/permission error.
-        subprocess.check_call(
-            ["git", "clone", "--depth", "1", "--branch", ref, "https://github.com/rohitrango/FireANTs.git", tmp]
-        )
-        subprocess.check_call(
-            [sys.executable, "-m", "pip", "install", "-q", "--no-build-isolation", os.path.join(tmp, "fused_ops")],
-            env=env,
-        )
-        importlib.invalidate_caches()
-        importlib.import_module("fireants_fused_ops")
-        _log("compiled and installed the fast CUDA kernels.")
-    except Exception as exc:
-        _log(
-            "kernel build failed -> running in pure PyTorch (correct, only speed/memory affected). Set "
-            "FIREANTS_FUSED_OPS_REF to a compatible FireANTs tag, or install a prebuilt fireants-fused-ops "
-            f"wheel, to enable them. Details: {exc}"
-        )
-    finally:
-        shutil.rmtree(tmp, ignore_errors=True)
-
-
-def registry_choices() -> list[str]:
-    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load ``models.json`` (available feature models). ``KONFAI_IMPACT_MODELS_REGISTRY`` (local path) wins
-    for dev/offline; otherwise ``ref`` is a ``repo:file`` Hugging Face reference (fetched, not bundled)."""
-    from huggingface_hub import hf_hub_download
-
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference — or set "
-            "KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-@dataclass
-class ModelSpec:
-    """One IMPACT feature model in the deformable metric (several are fused). ``ref`` picks the model; the
-    rest are its per-model knobs — the same as the ConvexAdam / elastix ``ModelSpec`` except ``voxel_size``
-    (an itk-impact resampling knob) has no meaning for FireANTs' geometry-free torch ``custom_loss`` and is
-    intentionally absent."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    layers_mask: str = "01"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
-    layers_weight: float = 1.0  # this model's weight in the multi-model fusion
-    pca: Annotated[int, Range(0, 100)] = 0  # keep the top-K principal components of the features (0 = no PCA)
-    distance: Literal["L1", "L2"] = "L1"
-
-
-@contextlib.contextmanager
-def _no_texpr_fuser():
-    """Disable the TensorExpr JIT fuser while IMPACT's TorchScript feature model runs under autograd.
-
-    The IMPACT feature models are TorchScript; run under FireANTs' gradient optimisation the TensorExpr
-    fuser trips on shape ops (``aten::size`` INTERNAL ASSERT). Scoped and restored so no other torch/JIT
-    user is affected; the modern profiling executor stays on (this is NOT the legacy executor).
-    """
-    torch._C._jit_set_texpr_fuser_enabled(False)
-    try:
-        yield
-    finally:
-        torch._C._jit_set_texpr_fuser_enabled(True)
-
-
-class _ImpactCore(IMPACTReg):
-    """One IMPACT feature model, exposed as a FireANTs ``forward(moved, fixed)``.
-
-    Reuses ``IMPACTReg._compute`` / ``preprocessing`` verbatim — the stats-normalised feature extraction
-    (the model wants per-image ``[min, mean, max, std]``) and the per-layer weighted distance — so the
-    metric is exactly KonfAI's, not a re-derivation. Only KonfAI's config-binding ``__init__`` and its
-    ``Attribute``-based geometry are replaced: FireANTs passes raw tensors at the current pyramid scale, so
-    the intensity statistics are computed from those tensors directly. ``pca`` (absent from KonfAI's torch
-    ``IMPACTReg``) is added here as a per-layer feature-space reduction matching itk-impact.
-    """
-
-    def __init__(self, ref: str, in_channels: int, weights: list[float], distance: str, pca: int) -> None:
-        from huggingface_hub import hf_hub_download
-
-        torch.nn.Module.__init__(self)  # bypass IMPACTReg.__init__ (KONFAI_CONFIG_PATH / apply_config binding)
-        self.name = "Reg"
-        self.in_channels = int(in_channels)
-        self.weights = [float(w) for w in weights]
-        self.nb_layer = len(self.weights)
-        self.loss = _DISTANCES[distance]()
-        self.pca = int(pca)  # PCA lives in KonfAI's IMPACTReg._compute (same behaviour as itk-impact)
-        self.dim = DIM
-        self.shape = None  # score the whole (downsampled) tensor — no ModelPatch tiling
-        if ":" in ref:  # a "repo:path" HF reference; otherwise a local model file
-            repo, filename = ref.split(":", 1)
-            self.model_path = hf_hub_download(repo, filename, repo_type="model")  # nosec B615
-        else:
-            self.model_path = ref
-        self.model = None  # lazy-loaded on the first forward, like IMPACTReg
-
-    @staticmethod
-    def _stats(tensor: torch.Tensor) -> dict:
-        detached = tensor.detach()
-        return {
-            "ImageMin": float(detached.min()),
-            "ImageMean": float(detached.mean()),
-            "ImageMax": float(detached.max()),
-            "ImageStd": float(detached.std()),
-        }
-
-    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
-        if self.model is None:
-            self.model = torch.jit.load(self.model_path)  # nosec B614
-        self.model.to(moved.device).eval()
-        with _no_texpr_fuser():
-            loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None)
-        return loss / max(true_nb, 1)
-
-
-class ImpactFeatureLoss(torch.nn.Module):
-    """FireANTs ``custom_loss`` = the KonfAI IMPACT metric fused over several feature models.
-
-    ``forward(moved, fixed)`` sums each model's ``layers_weight * IMPACT(model)``. A model's per-layer
-    weights come from its ``layers_mask`` bitmask; its input channel count is read from the registry
-    (``models.json`` ``numberofchannels``) so it never has to be configured by hand.
-    """
-
-    def __init__(self, specs: list["ModelSpec"]) -> None:
-        super().__init__()
-        registry = load_models_registry()
-        self._cores = torch.nn.ModuleList()
-        self._model_weights: list[float] = []
-        for spec in specs:
-            in_channels = int(registry.get(spec.ref.split(":", 1)[-1], {}).get("numberofchannels", 1))
-            weights = [1.0 if char == "1" else 0.0 for char in spec.layers_mask]
-            self._cores.append(_ImpactCore(spec.ref, in_channels, weights, spec.distance, spec.pca))
-            self._model_weights.append(float(spec.layers_weight))
-
-    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:
-        total: torch.Tensor | None = None
-        for weight, core in zip(self._model_weights, self._cores, strict=True):
-            term = weight * core(moved, fixed)
-            total = term if total is None else total + term
-        return total
-
-
-class FireANTsEngine:
-    """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return
-    (moved, dvf) on the fixed grid.
-
-    ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config
-    /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet``
-    annotations) on a machine without a GPU or without FireANTs installed.
-    """
-
-    def __init__(
-        self,
-        scales: list[int],
-        affine_iterations: list[int],
-        deformable_iterations: list[int],
-        cc_kernel: int,
-        affine_metric: str,
-        affine_lr: float,
-        deformable_method: str,
-        deformable_metric: str,
-        deformable_lr: float,
-        integrator_n: int,
-        smooth_warp_sigma: float,
-        smooth_grad_sigma: float,
-        seed: int,
-        impact_specs: list["ModelSpec"],
-    ) -> None:
-        self._scales = [int(s) for s in scales]
-        self._affine_iterations = [int(i) for i in affine_iterations]
-        self._deformable_iterations = [int(i) for i in deformable_iterations]
-        self._cc_kernel = int(cc_kernel)
-        self._affine_metric = affine_metric
-        self._affine_lr = float(affine_lr)
-        self._deformable_method = deformable_method
-        self._deformable_metric = deformable_metric
-        self._deformable_lr = float(deformable_lr)
-        self._integrator_n = int(integrator_n)
-        self._smooth_warp_sigma = float(smooth_warp_sigma)
-        self._smooth_grad_sigma = float(smooth_grad_sigma)
-        self._seed = int(seed)
-        # IMPACT deformable metric (only used when deformable_metric == "impact"): KonfAI IMPACT feature
-        # models drive the SyN/greedy stage instead of the analytic CC/MI/MSE.
-        self._impact_specs = impact_specs
-
-    @staticmethod
-    def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-        """True only for a mask that actually restricts the region — some voxels in, some out. An absent
-        optional mask arrives as a whole-image (all-ones) default and an all-zero mask is degenerate; both
-        are treated as no mask so the plain (non-masked) metric path is used."""
-        if mask is None:
-            return False
-        arr = sitk.GetArrayViewFromImage(mask)
-        return bool((arr > 0).any()) and bool((arr == 0).any())
-
-    @staticmethod
-    def _affine_to_sitk(affine_matrix: "torch.Tensor") -> sitk.AffineTransform:
-        """FireANTs' physical (LPS) linear matrix -> SimpleITK AffineTransform (fixed -> moving points),
-        the same convention FireANTs writes into an ANTs ``0GenericAffine.mat``."""
-        matrix = affine_matrix.float().cpu().numpy()[0]
-        affine = sitk.AffineTransform(DIM)
-        affine.SetMatrix(matrix[:DIM, :DIM].flatten().astype(np.float64))
-        affine.SetTranslation(matrix[:DIM, DIM].astype(np.float64))
-        return affine
-
-    def _total_field_transform(self, reg) -> sitk.Transform:
-        """Optimise a deformable stage and return its TOTAL displacement (affine baked in) as a
-        SimpleITK ``DisplacementFieldTransform`` on the fixed grid.
-
-        FireANTs serialises the total field (ANTs convention, fixed grid) only to a file, so it is
-        round-tripped through a temporary NIfTI — its public API, no internals reimplemented."""
-        reg.optimize()
-        with tempfile.TemporaryDirectory() as tmp:
-            warp_path = os.path.join(tmp, "total_warp.nii.gz")
-            reg.save_as_ants_transforms(warp_path)
-            total_field = sitk.ReadImage(warp_path, sitk.sitkVectorFloat64)
-        return sitk.DisplacementFieldTransform(total_field)  # consumes total_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."""
-        ensure_fireants_runtime()
-        from fireants.io import BatchedImages, Image
-        from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones
-        from fireants.registration.affine import AffineRegistration
-        from fireants.registration.rigid import RigidRegistration
-
-        torch.manual_seed(self._seed)
-        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
-        # FireANTs' Image ctor accepts a SimpleITK image directly, so the fixed/moving cross into
-        # FireANTs in-memory (no file load) with their geometry preserved.
-        fixed_img = Image(fixed, device=device)
-        moving_img = Image(moving, device=device)
-
-        # Masked metric only when a mask genuinely restricts the region. FireANTs' masked mode wants the
-        # mask as the last channel of BOTH images (all-ones where one side has none) and a ``masked_``
-        # metric prefix; the plain path is untouched when no real mask is present.
-        use_fixed_mask = self._is_partial_mask(fixed_mask)
-        use_moving_mask = self._is_partial_mask(moving_mask)
-        masked = use_fixed_mask or use_moving_mask
-        if masked:
-            fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img)
-            mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img)
-            fixed_img = apply_mask_to_image(fixed_img, fmask)
-            moving_img = apply_mask_to_image(moving_img, mmask)
-
-        bf = BatchedImages([fixed_img])
-        bm = BatchedImages([moving_img])
-        affine_loss = f"masked_{self._affine_metric}" if masked else self._affine_metric
-        deformable_loss = f"masked_{self._deformable_metric}" if masked else self._deformable_metric
-
-        # Linear: Rigid(MI, COM init) -> Affine(MI, seeded by the rigid), mirroring ANTs. The affine
-        # seeds the deformable stage (or is the whole transform when deformable_method == "none").
-        rigid = RigidRegistration(
-            scales=self._scales,
-            iterations=self._affine_iterations,
-            fixed_images=bf,
-            moving_images=bm,
-            loss_type=affine_loss,
-            optimizer="Adam",
-            optimizer_lr=self._affine_lr,
-            cc_kernel_size=self._cc_kernel,
-            init_translation="cof",
-        )
-        rigid.optimize()
-        rigid_matrix = rigid.get_rigid_matrix().detach()
-
-        affine = AffineRegistration(
-            scales=self._scales,
-            iterations=self._affine_iterations,
-            fixed_images=bf,
-            moving_images=bm,
-            loss_type=affine_loss,
-            optimizer="Adam",
-            optimizer_lr=self._affine_lr,
-            cc_kernel_size=self._cc_kernel,
-            init_rigid=rigid_matrix,
-        )
-        affine.optimize()
-        affine_matrix = affine.get_affine_matrix().detach()
-
-        # Deformable stage (or none). SyN and Greedy share the same constructor surface; both warm-start
-        # from the affine so their TOTAL transform already bakes in the linear pre-align.
-        if self._deformable_method == "none":
-            transform: sitk.Transform = self._affine_to_sitk(affine_matrix)
-        else:
-            if self._deformable_method == "syn":
-                from fireants.registration.syn import SyNRegistration as Deformable
-            elif self._deformable_method == "greedy":
-                from fireants.registration.greedy import GreedyRegistration as Deformable
-            else:
-                raise ValueError(
-                    f"Unknown deformable_method '{self._deformable_method}' (expected 'syn', 'greedy' or 'none')."
-                )
-            # "impact" swaps the analytic metric for a KonfAI IMPACT feature loss on the deformable stage
-            # (the linear pre-align keeps its own affine_metric); masks do not restrict the IMPACT metric.
-            if self._deformable_metric == "impact":
-                loss_type: str = "custom"
-                custom_loss: torch.nn.Module | None = ImpactFeatureLoss(self._impact_specs)
-            else:
-                loss_type, custom_loss = deformable_loss, None
-            reg = Deformable(
-                scales=self._scales,
-                iterations=self._deformable_iterations,
-                fixed_images=bf,
-                moving_images=bm,
-                loss_type=loss_type,
-                custom_loss=custom_loss,
-                cc_kernel_size=self._cc_kernel,
-                deformation_type="compositive",
-                integrator_n=self._integrator_n,
-                smooth_warp_sigma=self._smooth_warp_sigma,
-                smooth_grad_sigma=self._smooth_grad_sigma,
-                optimizer="Adam",
-                optimizer_lr=self._deformable_lr,
-                init_affine=affine_matrix,
-            )
-            transform = self._total_field_transform(reg)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-
-        # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path,
-        # so every FireANTs preset emits identical-shaped results.
-        moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-        dvf = sitk.TransformToDisplacementField(
-            transform,
-            sitk.sitkVectorFloat64,
-            fixed.GetSize(),
-            fixed.GetOrigin(),
-            fixed.GetSpacing(),
-            fixed.GetDirection(),
-        )
-        moved_np, _ = image_to_data(moved)
-        dvf_np, _ = image_to_data(dvf)
-        return moved_np, dvf_np
-
-
-class FireANTsRegistration(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 the ConvexAdam / elastix engines); registration needs the
-    physical geometry, and the mask branches restrict the metric.
-    """
-
-    accepts_attributes = True
-
-    def __init__(self, engine: FireANTsEngine) -> 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. A whole-image
-        # mask (the default when none is supplied) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        # FireANTs runs a gradient-based instance optimisation (Riemannian Adam over the warp); the
-        # predictor calls forward under torch.inference_mode(), which forbids autograd. The image tensors
-        # have already crossed to numpy/SimpleITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad():
-            for b in range(fixed.shape[0]):
-                fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-                moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-                fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-                moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-                moved_np, dvf_np = self._engine.register(
-                    fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-                )
-                combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise FireANTs registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = 2, moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    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``. The knobs below are read straight from these annotations by the
-    UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method``
-    is the knob that specialises this shared model into each FireANTs preset.
-    """
-
-    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()},
-        scales: list[int] = [4, 2, 1],
-        affine_iterations: list[int] = [200, 100, 50],
-        deformable_iterations: list[int] = [200, 100, 50],
-        cc_kernel: Annotated[int, Range(1, 21)] = 5,
-        affine_metric: Literal["mi", "cc", "mse"] = "mi",
-        affine_lr: Annotated[float, Range(0.0, 10.0)] = 0.003,
-        deformable_method: Literal["none", "syn", "greedy"] = "syn",
-        deformable_metric: Literal["cc", "mi", "mse", "impact"] = "cc",
-        deformable_lr: Annotated[float, Range(0.0, 10.0)] = 0.25,
-        integrator_n: Annotated[int, Range(1, 100)] = 10,
-        smooth_warp_sigma: Annotated[float, Range(0.0, 100.0)] = 0.5,
-        smooth_grad_sigma: Annotated[float, Range(0.0, 100.0)] = 1.0,
-        seed: int = 42,
-        models: dict[str, ModelSpec] = {},
-    ) -> None:
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        engine = FireANTsEngine(
-            scales,
-            affine_iterations,
-            deformable_iterations,
-            cc_kernel,
-            affine_metric,
-            affine_lr,
-            deformable_method,
-            deformable_metric,
-            deformable_lr,
-            integrator_n,
-            smooth_warp_sigma,
-            smooth_grad_sigma,
-            seed,
-            _sorted_specs(models),
-        )
-        self.add_module(
-            "Registration", FireANTsRegistration(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"])

FireANTs_SyN/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.fireants:RegistrationNet
     RegistrationNet:
       scales:
       - 4

FireANTs_SyN/requirements.txt

@@ -1,3 +1,4 @@
 hydra-core
 nibabel
 pandas
+impact_reg_konfai>=1.6.1

Generic_Rigid/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            in_branch=[0, 1, 2, 3],
-            out_branch=["registration"],
-        )
-        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
-        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

Generic_Rigid/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps:

Generic_Rigid/elastix_engine.py

@@ -1,386 +0,0 @@
-# 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
-
-"""Elastix-IMPACT runtime for the registration bundle.
-
-``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
-parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
-``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
-images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
-"""
-
-import os
-import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-import tqdm
-from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
-
-
-def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
-    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
-    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
-    the whole image) instead of paying for a mask that restricts nothing."""
-    if mask is None:
-        return False
-    arr = sitk.GetArrayViewFromImage(mask)
-    return bool((arr > 0).any()) and bool((arr == 0).any())
-
-
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
-    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random FOV-sized patches each iteration. One mode per preset.
-        self._mode = mode
-        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
-        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry() if resolutions else {}
-        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
-        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
-        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
-        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
-        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
-        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
-        # surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into ``work``.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
-        applies every per-token / exact override. Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
-        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
-            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if _is_partial_mask(mask):
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
-            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
-            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
-    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            mode,
-        )
-
-    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) stacked with the DVF (dim channels), both on
-        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)

Generic_Rigid/install.py

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

Generic_Rigid/requirements.txt

@@ -0,0 +1 @@
+impact_reg_konfai>=1.6.1

Generic_Rigid_BSpline/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            in_branch=[0, 1, 2, 3],
-            out_branch=["registration"],
-        )
-        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
-        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

Generic_Rigid_BSpline/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps:

Generic_Rigid_BSpline/elastix_engine.py

@@ -1,386 +0,0 @@
-# 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
-
-"""Elastix-IMPACT runtime for the registration bundle.
-
-``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
-parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
-``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
-images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
-"""
-
-import os
-import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-import tqdm
-from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
-
-
-def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
-    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
-    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
-    the whole image) instead of paying for a mask that restricts nothing."""
-    if mask is None:
-        return False
-    arr = sitk.GetArrayViewFromImage(mask)
-    return bool((arr > 0).any()) and bool((arr == 0).any())
-
-
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
-    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random FOV-sized patches each iteration. One mode per preset.
-        self._mode = mode
-        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
-        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry() if resolutions else {}
-        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
-        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
-        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
-        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
-        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
-        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
-        # surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into ``work``.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
-        applies every per-token / exact override. Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
-        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
-            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if _is_partial_mask(mask):
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
-            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
-            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
-    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            mode,
-        )
-
-    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) stacked with the DVF (dim channels), both on
-        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)

Generic_Rigid_BSpline/install.py

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

Generic_Rigid_BSpline/requirements.txt

@@ -0,0 +1 @@
+impact_reg_konfai>=1.6.1

MR_CT_HeadNeck/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            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"])

MR_CT_HeadNeck/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps:

MR_CT_HeadNeck/elastix_engine.py

@@ -1,386 +0,0 @@
-# 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
-
-"""Elastix-IMPACT runtime for the registration bundle.
-
-``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
-parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
-``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
-images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
-"""
-
-import os
-import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-import tqdm
-from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
-
-
-def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
-    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
-    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
-    the whole image) instead of paying for a mask that restricts nothing."""
-    if mask is None:
-        return False
-    arr = sitk.GetArrayViewFromImage(mask)
-    return bool((arr > 0).any()) and bool((arr == 0).any())
-
-
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
-    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random FOV-sized patches each iteration. One mode per preset.
-        self._mode = mode
-        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
-        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry() if resolutions else {}
-        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
-        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
-        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
-        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
-        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
-        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
-        # surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into ``work``.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
-        applies every per-token / exact override. Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
-        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
-            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if _is_partial_mask(mask):
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
-            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
-            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
-    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            mode,
-        )
-
-    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) stacked with the DVF (dim channels), both on
-        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)

MR_CT_HeadNeck/install.py

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

MR_CT_HeadNeck/requirements.txt

@@ -0,0 +1 @@
+impact_reg_konfai>=1.6.1

MR_CT_MRSeg/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            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"])

MR_CT_MRSeg/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps:

MR_CT_MRSeg/elastix_engine.py

@@ -1,386 +0,0 @@
-# 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
-
-"""Elastix-IMPACT runtime for the registration bundle.
-
-``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
-parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
-``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
-images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
-"""
-
-import os
-import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import SimpleITK as sitk
-import torch
-import tqdm
-from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
-
-
-def _is_partial_mask(mask: "sitk.Image | None") -> bool:
-    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
-    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
-    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
-    the whole image) instead of paying for a mask that restricts nothing."""
-    if mask is None:
-        return False
-    arr = sitk.GetArrayViewFromImage(mask)
-    return bool((arr > 0).any()) and bool((arr == 0).any())
-
-
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
-    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random FOV-sized patches each iteration. One mode per preset.
-        self._mode = mode
-        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
-        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry() if resolutions else {}
-        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
-        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
-        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
-        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
-        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
-        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
-        # surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into ``work``.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
-        applies every per-token / exact override. Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
-        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
-            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if _is_partial_mask(mask):
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
-            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
-            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
-    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            mode,
-        )
-
-    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) stacked with the DVF (dim channels), both on
-        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)

MR_CT_MRSeg/install.py

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

MR_CT_MRSeg/requirements.txt

@@ -0,0 +1 @@
+impact_reg_konfai>=1.6.1

MR_CT_TS/Model.py

@@ -1,301 +0,0 @@
-# Copyright (c) 2025 Valentin Boussot
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# SPDX-License-Identifier: Apache-2.0
-
-"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
-
-``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
-2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
-the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
-the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
-subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
-
-A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
-``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
-
-NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
-(``get_origin``); PEP 563 stringized annotations break arg resolution.
-"""
-
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Annotated, Literal
-
-import torch
-from huggingface_hub import hf_hub_download
-from konfai.network import network
-from konfai.utils.config import Choices, Range
-
-# IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
-# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
-_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
-
-# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
-# Static anyway; in Jacobian we clamp ``l`` to this plateau.
-_FOV_RAMP_MAX_LAYER = 6
-
-
-def registry_choices() -> list[str]:
-    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
-    (offline-first). A user may still point ``ref`` at a local model."""
-    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
-    return [f"{repo}:{key}" for key in load_models_registry()]
-
-
-def _num(x: object) -> str:
-    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
-    return "%g" % float(x)
-
-
-@dataclass
-class ModelSpec:
-    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
-    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
-    (``models.json``) keyed by ``ref`` — never tuned."""
-
-    ref: Annotated[str, Choices(registry_choices)]
-    voxel_size: list[float] = field(default_factory=list)
-    layers_weight: list[float] = field(default_factory=lambda: [1.0])
-    subset_features: Annotated[int, Range(0, 1000)] = 0
-    pca: Annotated[int, Range(0, 100)] = 0
-    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
-    layers_mask: str = ""
-
-
-@dataclass
-class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
-
-    max_iterations: Annotated[int, Range(1, 100000)]
-    models: dict[str, ModelSpec]
-
-
-def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
-    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
-
-
-def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
-
-    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
-    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
-    """
-    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
-    if local:
-        path = Path(local)
-    elif ":" in ref:
-        repo, filename = ref.split(":", 1)
-        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-    else:
-        raise ValueError(
-            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference (the registry is fetched "
-            f"from HF, not bundled) — or set KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
-        )
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
-    return ref.split(":", 1)[1] if ":" in ref else ref
-
-
-def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
-
-    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
-    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
-    receptive field, so the FOV is governed by the rightmost ``'1'``.
-    """
-    mask = layers_mask.strip().strip('"')
-    active = [i for i, char in enumerate(mask) if char == "1"]
-    if not active:
-        raise ValueError(f"LayersMask '{layers_mask}' selects no layer; cannot derive the model FOV.")
-    return max(active)
-
-
-def _fov_value(fov: dict, layers_mask: str) -> int:
-    """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
-
-    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
-                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
-    """
-    formula = str(fov.get("formula", "")).strip()
-    key = re.sub(r"\s+", "", formula).lower()
-    if key.isdigit():
-        return int(key)
-    if key == "2*r*d+1":
-        return 2 * int(fov["r"]) * int(fov["d"]) + 1
-    if key == "2^l+3":
-        return 2 ** min(_deepest_active_layer(layers_mask), _FOV_RAMP_MAX_LAYER) + 3
-    if "global" in key:
-        raise ValueError(f"model FOV '{formula}' is whole-image only (Static); it has no Jacobian patch size.")
-    if fov.get("value") is not None:
-        return int(fov["value"])
-    raise ValueError(f"cannot evaluate model FOV formula '{formula}'.")
-
-
-def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
-    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
-    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
-    dim = int(entry.get("dimension", 3))
-    if mode.strip().strip('"').lower() != "jacobian":
-        return " ".join(["0"] * dim)
-    fov = _fov_value(entry.get("fov", {}), layers_mask)
-    return " ".join([str(fov)] * dim)
-
-
-def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
-    """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
-
-    Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
-    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
-    from the registry formula and the cell's ``layers_mask``.
-    """
-    res = _sorted_specs(resolutions)
-    n = len(res)
-    mode_clean = mode.strip().strip('"') or "Static"
-
-    impact: list[str] = []
-    for k, r in enumerate(res):
-        models = _sorted_specs(r.models)
-        entries = [registry[_model_key(m.ref)] for m in models]
-
-        def row(stem: str, values: list[str]) -> None:
-            impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
-
-        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
-        # model FOV); everything else is a per-model knob taken straight from the cell.
-        row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
-        row("Dimension", [e["dimension"] for e in entries])
-        row("NumberOfChannels", [e["numberofchannels"] for e in entries])
-        row("PatchSize", [_patch_size(mode_clean, e, m.layers_mask) for e, m in zip(entries, models)])
-        row("VoxelSize", [" ".join(_num(v) for v in m.voxel_size) for m in models])
-        row("LayersMask", [f'"{m.layers_mask}"' for m in models])
-        row("SubsetFeatures", [str(m.subset_features) for m in models])
-        row("PCA", [str(m.pca) for m in models])
-        row("Distance", [f'"{m.distance}"' for m in models])
-        row("LayersWeight", [" ".join(_num(w) for w in m.layers_weight) for m in models])
-        impact.append("")  # blank line between resolutions, mirroring the reference maps
-
-    # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
-    lines = template_text.splitlines()
-    indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
-    block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
-    block_lo, block_hi = (block_rows[0], block_rows[-1]) if block_rows else (-1, -2)
-
-    out: list[str] = []
-    for i, (m, line) in enumerate(indexed):
-        key = m.group(1) if m else None
-        if block_lo <= i <= block_hi:
-            if i == block_lo:  # replace the whole span at its first line, drop the rest (incl. inner blanks)
-                out.extend(impact[:-1])
-        elif key == "MaximumNumberOfIterations":
-            out.append("(MaximumNumberOfIterations " + " ".join(_num(r.max_iterations) for r in res) + ")")
-        elif key == "NumberOfResolutions":
-            out.append(f"(NumberOfResolutions {n})")
-        elif key in ("FixedImagePyramidRescaleSchedule", "MovingImagePyramidRescaleSchedule"):
-            out.append(f"({key} " + " ".join(["1"] * 3 * n) + ")")
-        elif key == "ImpactMode":
-            out.append(f'(ImpactMode "{mode_clean}")')
-        else:
-            out.append(line)
-    return "\n".join(out)
-
-
-class ChannelSelect(torch.nn.Module):
-    """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
-
-    def __init__(self, start: int, stop: int) -> None:
-        super().__init__()
-        self._start = start
-        self._stop = stop
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor[:, self._start : self._stop]
-
-
-class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
-    moving mask = 3; masks restrict the metric, whole-image = no restriction).
-
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
-    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
-    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
-    ``same_as_group: Volume_0:Fixed``.
-    """
-
-    def __init__(
-        self,
-        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-        schedulers: dict[str, network.LRSchedulersLoader] = {
-            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
-        },
-        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        engine: str = "elastix",
-        parameter_maps: list[str] = [],
-        max_iterations: Annotated[int, Range(0, 100000)] = 0,
-        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
-        subset_features: Annotated[int, Range(0, 1000)] = 0,
-        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict[str, ResolutionSpec] = {},
-        mode: Literal["Static", "Jacobian"] = "Static",
-    ) -> None:
-        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
-        # lists its self-configured models; the download list is derived from the cells. Global knobs override
-        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
-        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
-        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
-        from elastix_engine import ElastixRegistration
-
-        super().__init__(
-            in_channels=1,
-            optimizer=optimizer,
-            schedulers=schedulers,
-            outputs_criterions=outputs_criterions,
-            dim=3,
-        )
-        self.add_module(
-            "Registration",
-            ElastixRegistration(
-                engine,
-                parameter_maps,
-                max_iterations,
-                final_grid_spacing,
-                subset_features,
-                spatial_samples,
-                parameter_overrides,
-                resolutions,
-                mode,
-            ),
-            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"])

MR_CT_TS/Prediction.yml

@@ -1,6 +1,6 @@
 Predictor:
   Model:
-    classpath: Model:RegistrationNet
+    classpath: impact_reg_konfai.models.elastix:RegistrationNet
     RegistrationNet:
       engine: elastix
       parameter_maps: