diff --git a/CBCT_CT_HeadNeck/Model.py b/CBCT_CT_HeadNeck/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/CBCT_CT_HeadNeck/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/CBCT_CT_HeadNeck/Prediction.yml b/CBCT_CT_HeadNeck/Prediction.yml index d8201fa2b55de9a35aeafac563973353f4e7e070..614766ffd9f1cf6353e57b0ae22411b3de916183 100644 --- a/CBCT_CT_HeadNeck/Prediction.yml +++ b/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: diff --git a/CBCT_CT_HeadNeck/elastix_engine.py b/CBCT_CT_HeadNeck/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/CBCT_CT_HeadNeck/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/CBCT_CT_HeadNeck/install.py b/CBCT_CT_HeadNeck/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/CBCT_CT_HeadNeck/install.py +++ /dev/null @@ -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() diff --git a/CBCT_CT_HeadNeck/requirements.txt b/CBCT_CT_HeadNeck/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/CBCT_CT_HeadNeck/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1 diff --git a/CBCT_CT_MRSeg/Model.py b/CBCT_CT_MRSeg/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/CBCT_CT_MRSeg/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/CBCT_CT_MRSeg/Prediction.yml b/CBCT_CT_MRSeg/Prediction.yml index 64b50e707b9abb141b574c9164826b3f8c65a56d..3c39b582b90db366bac2d1fea6aceee15f5aeeae 100644 --- a/CBCT_CT_MRSeg/Prediction.yml +++ b/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: diff --git a/CBCT_CT_MRSeg/elastix_engine.py b/CBCT_CT_MRSeg/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/CBCT_CT_MRSeg/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/CBCT_CT_MRSeg/install.py b/CBCT_CT_MRSeg/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/CBCT_CT_MRSeg/install.py +++ /dev/null @@ -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() diff --git a/CBCT_CT_MRSeg/requirements.txt b/CBCT_CT_MRSeg/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/CBCT_CT_MRSeg/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1 diff --git a/ConvexAdam_Coarse/Model.py b/ConvexAdam_Coarse/Model.py deleted file mode 100644 index e46aff0b20500d813ea3d0667f04f14d9e9c72e6..0000000000000000000000000000000000000000 --- a/ConvexAdam_Coarse/Model.py +++ /dev/null @@ -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"]) diff --git a/ConvexAdam_Coarse/Prediction.yml b/ConvexAdam_Coarse/Prediction.yml index 8b1f6071905fdcbda24f1e093222080bdf5d6eef..5c27813a51156f74bda6c728cea18e6ff832c21a 100644 --- a/ConvexAdam_Coarse/Prediction.yml +++ b/ConvexAdam_Coarse/Prediction.yml @@ -1,6 +1,6 @@ Predictor: Model: - classpath: Model:RegistrationNet + classpath: impact_reg_konfai.models.convexadam:RegistrationNet RegistrationNet: models: '0': diff --git a/ConvexAdam_Coarse/requirements.txt b/ConvexAdam_Coarse/requirements.txt index ea9845e3595c2d936f0873883e3ed47732de9fae..d7831826ab8c4e8453fea3883f12daa93d68cd6e 100644 --- a/ConvexAdam_Coarse/requirements.txt +++ b/ConvexAdam_Coarse/requirements.txt @@ -1 +1,2 @@ itk-impact +impact_reg_konfai>=1.6.1 diff --git a/ConvexAdam_Composite/Model.py b/ConvexAdam_Composite/Model.py deleted file mode 100644 index e46aff0b20500d813ea3d0667f04f14d9e9c72e6..0000000000000000000000000000000000000000 --- a/ConvexAdam_Composite/Model.py +++ /dev/null @@ -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"]) diff --git a/ConvexAdam_Composite/Prediction.yml b/ConvexAdam_Composite/Prediction.yml index d4920be1cdf94f03306bf87e08ea56f385e31023..f773eb39330820d438d0977b3c4c706b06e6d0fb 100644 --- a/ConvexAdam_Composite/Prediction.yml +++ b/ConvexAdam_Composite/Prediction.yml @@ -1,6 +1,6 @@ Predictor: Model: - classpath: Model:RegistrationNet + classpath: impact_reg_konfai.models.convexadam:RegistrationNet RegistrationNet: models: '0': diff --git a/ConvexAdam_Composite/requirements.txt b/ConvexAdam_Composite/requirements.txt index ea9845e3595c2d936f0873883e3ed47732de9fae..d7831826ab8c4e8453fea3883f12daa93d68cd6e 100644 --- a/ConvexAdam_Composite/requirements.txt +++ b/ConvexAdam_Composite/requirements.txt @@ -1 +1,2 @@ itk-impact +impact_reg_konfai>=1.6.1 diff --git a/ConvexAdam_Fine/Model.py b/ConvexAdam_Fine/Model.py deleted file mode 100644 index e46aff0b20500d813ea3d0667f04f14d9e9c72e6..0000000000000000000000000000000000000000 --- a/ConvexAdam_Fine/Model.py +++ /dev/null @@ -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"]) diff --git a/ConvexAdam_Fine/Prediction.yml b/ConvexAdam_Fine/Prediction.yml index a8daba61e4c4b68e61bb3f761fd4da8761d55885..74e5a82450ccdc03c1ecec46d2b5074ecc4dadd4 100644 --- a/ConvexAdam_Fine/Prediction.yml +++ b/ConvexAdam_Fine/Prediction.yml @@ -1,6 +1,6 @@ Predictor: Model: - classpath: Model:RegistrationNet + classpath: impact_reg_konfai.models.convexadam:RegistrationNet RegistrationNet: models: '0': diff --git a/ConvexAdam_Fine/requirements.txt b/ConvexAdam_Fine/requirements.txt index ea9845e3595c2d936f0873883e3ed47732de9fae..d7831826ab8c4e8453fea3883f12daa93d68cd6e 100644 --- a/ConvexAdam_Fine/requirements.txt +++ b/ConvexAdam_Fine/requirements.txt @@ -1 +1,2 @@ itk-impact +impact_reg_konfai>=1.6.1 diff --git a/FireANTs_Affine/Model.py b/FireANTs_Affine/Model.py deleted file mode 100644 index 22f3c7aaae4da1d7cc07a35fd71c5a20d641f9a3..0000000000000000000000000000000000000000 --- a/FireANTs_Affine/Model.py +++ /dev/null @@ -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"]) diff --git a/FireANTs_Affine/Prediction.yml b/FireANTs_Affine/Prediction.yml index f6431014ebbfc2d5af0a85844b7e067ac520baa3..92da59dd818ae5101a44f8cfefcf88487a2e7a51 100644 --- a/FireANTs_Affine/Prediction.yml +++ b/FireANTs_Affine/Prediction.yml @@ -1,6 +1,6 @@ Predictor: Model: - classpath: Model:RegistrationNet + classpath: impact_reg_konfai.models.fireants:RegistrationNet RegistrationNet: scales: - 4 diff --git a/FireANTs_Affine/requirements.txt b/FireANTs_Affine/requirements.txt index 9bc74d0f2b0f12aebf323b583e62d92ef9a43667..29c53d99c78140da47f608f307d740fe14a23c20 100644 --- a/FireANTs_Affine/requirements.txt +++ b/FireANTs_Affine/requirements.txt @@ -1,3 +1,4 @@ hydra-core nibabel pandas +impact_reg_konfai>=1.6.1 diff --git a/FireANTs_IMPACT/Model.py b/FireANTs_IMPACT/Model.py deleted file mode 100644 index 22f3c7aaae4da1d7cc07a35fd71c5a20d641f9a3..0000000000000000000000000000000000000000 --- a/FireANTs_IMPACT/Model.py +++ /dev/null @@ -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"]) diff --git a/FireANTs_IMPACT/Prediction.yml b/FireANTs_IMPACT/Prediction.yml index 6ff588261d1026c3b8e053267a095da2802f5319..f77daceeb767d4f1b7eab6b8531e0528771444a3 100644 --- a/FireANTs_IMPACT/Prediction.yml +++ b/FireANTs_IMPACT/Prediction.yml @@ -1,6 +1,6 @@ Predictor: Model: - classpath: Model:RegistrationNet + classpath: impact_reg_konfai.models.fireants:RegistrationNet RegistrationNet: scales: - 4 diff --git a/FireANTs_IMPACT/requirements.txt b/FireANTs_IMPACT/requirements.txt index 9bc74d0f2b0f12aebf323b583e62d92ef9a43667..29c53d99c78140da47f608f307d740fe14a23c20 100644 --- a/FireANTs_IMPACT/requirements.txt +++ b/FireANTs_IMPACT/requirements.txt @@ -1,3 +1,4 @@ hydra-core nibabel pandas +impact_reg_konfai>=1.6.1 diff --git a/FireANTs_SyN/Model.py b/FireANTs_SyN/Model.py deleted file mode 100644 index 22f3c7aaae4da1d7cc07a35fd71c5a20d641f9a3..0000000000000000000000000000000000000000 --- a/FireANTs_SyN/Model.py +++ /dev/null @@ -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"]) diff --git a/FireANTs_SyN/Prediction.yml b/FireANTs_SyN/Prediction.yml index 3f0cbdbc3cdcbed31ae6cd6633a93d3afbd404ed..ed1cbea82216172c8ef6e3c75976bd5b0e6f2f9f 100644 --- a/FireANTs_SyN/Prediction.yml +++ b/FireANTs_SyN/Prediction.yml @@ -1,6 +1,6 @@ Predictor: Model: - classpath: Model:RegistrationNet + classpath: impact_reg_konfai.models.fireants:RegistrationNet RegistrationNet: scales: - 4 diff --git a/FireANTs_SyN/requirements.txt b/FireANTs_SyN/requirements.txt index 9bc74d0f2b0f12aebf323b583e62d92ef9a43667..29c53d99c78140da47f608f307d740fe14a23c20 100644 --- a/FireANTs_SyN/requirements.txt +++ b/FireANTs_SyN/requirements.txt @@ -1,3 +1,4 @@ hydra-core nibabel pandas +impact_reg_konfai>=1.6.1 diff --git a/Generic_Rigid/Model.py b/Generic_Rigid/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/Generic_Rigid/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/Generic_Rigid/Prediction.yml b/Generic_Rigid/Prediction.yml index 78525d1774dd1c4a04fb0034969ac2a28def64f6..0e60d4521b9a536bccd42800af98df80f47f9d29 100644 --- a/Generic_Rigid/Prediction.yml +++ b/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: diff --git a/Generic_Rigid/elastix_engine.py b/Generic_Rigid/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/Generic_Rigid/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/Generic_Rigid/install.py b/Generic_Rigid/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/Generic_Rigid/install.py +++ /dev/null @@ -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() diff --git a/Generic_Rigid/requirements.txt b/Generic_Rigid/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/Generic_Rigid/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1 diff --git a/Generic_Rigid_BSpline/Model.py b/Generic_Rigid_BSpline/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/Generic_Rigid_BSpline/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/Generic_Rigid_BSpline/Prediction.yml b/Generic_Rigid_BSpline/Prediction.yml index 910ddc93a47344592ddd23e334792a325a779a51..3ae3ab9e41f402c5fa1a86457147d20127a49101 100644 --- a/Generic_Rigid_BSpline/Prediction.yml +++ b/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: diff --git a/Generic_Rigid_BSpline/elastix_engine.py b/Generic_Rigid_BSpline/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/Generic_Rigid_BSpline/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/Generic_Rigid_BSpline/install.py b/Generic_Rigid_BSpline/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/Generic_Rigid_BSpline/install.py +++ /dev/null @@ -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() diff --git a/Generic_Rigid_BSpline/requirements.txt b/Generic_Rigid_BSpline/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/Generic_Rigid_BSpline/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1 diff --git a/MR_CT_HeadNeck/Model.py b/MR_CT_HeadNeck/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/MR_CT_HeadNeck/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/MR_CT_HeadNeck/Prediction.yml b/MR_CT_HeadNeck/Prediction.yml index 82df0a4d90e120c0d2a656bbc0e3d2cac7be46d0..3c9811e9806ccada5acf2b6fc62d1a373a7beaf7 100644 --- a/MR_CT_HeadNeck/Prediction.yml +++ b/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: diff --git a/MR_CT_HeadNeck/elastix_engine.py b/MR_CT_HeadNeck/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/MR_CT_HeadNeck/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/MR_CT_HeadNeck/install.py b/MR_CT_HeadNeck/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/MR_CT_HeadNeck/install.py +++ /dev/null @@ -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() diff --git a/MR_CT_HeadNeck/requirements.txt b/MR_CT_HeadNeck/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/MR_CT_HeadNeck/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1 diff --git a/MR_CT_MRSeg/Model.py b/MR_CT_MRSeg/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/MR_CT_MRSeg/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/MR_CT_MRSeg/Prediction.yml b/MR_CT_MRSeg/Prediction.yml index ed615fec6dc91948c1a4402d5dd29d1da4074097..6b04a868712f862c8a5097cb47e9e754fec808ab 100644 --- a/MR_CT_MRSeg/Prediction.yml +++ b/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: diff --git a/MR_CT_MRSeg/elastix_engine.py b/MR_CT_MRSeg/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/MR_CT_MRSeg/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/MR_CT_MRSeg/install.py b/MR_CT_MRSeg/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/MR_CT_MRSeg/install.py +++ /dev/null @@ -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() diff --git a/MR_CT_MRSeg/requirements.txt b/MR_CT_MRSeg/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/MR_CT_MRSeg/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1 diff --git a/MR_CT_TS/Model.py b/MR_CT_TS/Model.py deleted file mode 100644 index 31308481f92071b74fad76ab6d0dfbaa443fd020..0000000000000000000000000000000000000000 --- a/MR_CT_TS/Model.py +++ /dev/null @@ -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`` 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"]) diff --git a/MR_CT_TS/Prediction.yml b/MR_CT_TS/Prediction.yml index 0e8dc8ffce9cd551d9f76f1251fbe1230b67b7c9..f2bc94fe619c479ca01bb7dc0b496f8ba35a6dc1 100644 --- a/MR_CT_TS/Prediction.yml +++ b/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: diff --git a/MR_CT_TS/elastix_engine.py b/MR_CT_TS/elastix_engine.py deleted file mode 100644 index cb966ed8aa887e8247c92e02b7da944622e051a9..0000000000000000000000000000000000000000 --- a/MR_CT_TS/elastix_engine.py +++ /dev/null @@ -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 /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) diff --git a/MR_CT_TS/install.py b/MR_CT_TS/install.py deleted file mode 100644 index 9af2b88d5553b2262a8c622def90be1cb141e224..0000000000000000000000000000000000000000 --- a/MR_CT_TS/install.py +++ /dev/null @@ -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() diff --git a/MR_CT_TS/requirements.txt b/MR_CT_TS/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e5761794428e9ee6edf61be5bda279992de3672 --- /dev/null +++ b/MR_CT_TS/requirements.txt @@ -0,0 +1 @@ +impact_reg_konfai>=1.6.1