diff --git a/ImpactReg/Evaluation_with_fid.yml b/CBCT_CT_HeadNeck/Evaluation_with_fid.yml similarity index 100% rename from ImpactReg/Evaluation_with_fid.yml rename to CBCT_CT_HeadNeck/Evaluation_with_fid.yml diff --git a/ImpactReg/Evaluation_with_images.yml b/CBCT_CT_HeadNeck/Evaluation_with_images.yml similarity index 100% rename from ImpactReg/Evaluation_with_images.yml rename to CBCT_CT_HeadNeck/Evaluation_with_images.yml diff --git a/ImpactReg/Evaluation_with_seg.yml b/CBCT_CT_HeadNeck/Evaluation_with_seg.yml similarity index 100% rename from ImpactReg/Evaluation_with_seg.yml rename to CBCT_CT_HeadNeck/Evaluation_with_seg.yml diff --git a/CBCT_CT_HeadNeck/Model.py b/CBCT_CT_HeadNeck/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/CBCT_CT_HeadNeck/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/ParameterMap_CBCT_HN.txt b/CBCT_CT_HeadNeck/ParameterMap_CBCT_HN.txt similarity index 100% rename from ParameterMap_CBCT_HN.txt rename to CBCT_CT_HeadNeck/ParameterMap_CBCT_HN.txt diff --git a/CBCT_CT_HeadNeck/Prediction.yml b/CBCT_CT_HeadNeck/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..0b170f24f1cf1d3605be729f0ca3939a3d58068e --- /dev/null +++ b/CBCT_CT_HeadNeck/Prediction.yml @@ -0,0 +1,108 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - ParameterMap_CBCT_HN.txt + models: + - VBoussot/impact-torchscript-models:TS/M732.pt + - VBoussot/impact-torchscript-models:TS/M731.pt + - VBoussot/impact-torchscript-models:TS/M730.pt + iterations: 1150 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-CBCT-CT-HeadNeck + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/CBCT_CT_HeadNeck/app.json b/CBCT_CT_HeadNeck/app.json new file mode 100644 index 0000000000000000000000000000000000000000..cf14005fcc63723d297680f7eddbc6714683ec93 --- /dev/null +++ b/CBCT_CT_HeadNeck/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "CBCT/CT Head&Neck", + "short_description": "Optimized preset for CBCT/CT registration on head & neck", + "description": "A five-level recursive B-spline deformable registration optimized for CBCT/CT head-and-neck alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 200, 200, and 150 iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm) using L1 distances on selected internal layers of the network. A composite objective (IMPACT + mutual information + bending energy penalty, with increased MI weight) ensures robust cross-modality alignment in complex head-and-neck anatomy while enforcing smooth, physically plausible deformations.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/CBCT_CT_HeadNeck/install.py b/CBCT_CT_HeadNeck/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/CBCT_CT_HeadNeck/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/CBCT_CT_HeadNeck/model.pt b/CBCT_CT_HeadNeck/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/CBCT_CT_HeadNeck/model.pt differ diff --git a/CBCT_CT_MRSeg/Evaluation_with_fid.yml b/CBCT_CT_MRSeg/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/CBCT_CT_MRSeg/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/CBCT_CT_MRSeg/Evaluation_with_images.yml b/CBCT_CT_MRSeg/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/CBCT_CT_MRSeg/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/CBCT_CT_MRSeg/Evaluation_with_seg.yml b/CBCT_CT_MRSeg/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/CBCT_CT_MRSeg/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/CBCT_CT_MRSeg/Model.py b/CBCT_CT_MRSeg/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/CBCT_CT_MRSeg/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/ParameterMap_CBCT_generic_MRSeg.txt b/CBCT_CT_MRSeg/ParameterMap_CBCT_generic_MRSeg.txt similarity index 100% rename from ParameterMap_CBCT_generic_MRSeg.txt rename to CBCT_CT_MRSeg/ParameterMap_CBCT_generic_MRSeg.txt diff --git a/CBCT_CT_MRSeg/Prediction.yml b/CBCT_CT_MRSeg/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..927619e47236cb272beb67333994c65af9caa84e --- /dev/null +++ b/CBCT_CT_MRSeg/Prediction.yml @@ -0,0 +1,107 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - ParameterMap_CBCT_generic_MRSeg.txt + models: + - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt + - VBoussot/impact-torchscript-models:MRSeg/MRSeg.pt + iterations: 1100 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-CBCT-CT-MRSeg + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/CBCT_CT_MRSeg/app.json b/CBCT_CT_MRSeg/app.json new file mode 100644 index 0000000000000000000000000000000000000000..552a747f6ce65ddc5a14099466090facabc84559 --- /dev/null +++ b/CBCT_CT_MRSeg/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "CBCT/CT preset with MRSegmentator", + "short_description": "Generic CBCT/CT deformable registration using MRSegmentator features", + "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from the pretrained MRSegmentator model. The scheme follows a multi-resolution strategy with up to 300, 300, 250, and 200 ASGD iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), with a level-dependent combination of Dice-based segmentation overlap and L1 feature distances on selected internal layers of MRSegmentator. Early levels rely on pure Dice supervision, while finer stages progressively integrate feature-level alignment with increasing L1 contribution (0.3/0.7, 0.5/0.5) and a final purely feature-based stage. The optimization minimizes a composite objective (IMPACT + mutual information + bending energy penalty), enabling robust cross-modality alignment between CBCT and CT while enforcing smooth, physically plausible deformations.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/CBCT_CT_MRSeg/install.py b/CBCT_CT_MRSeg/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/CBCT_CT_MRSeg/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/CBCT_CT_MRSeg/model.pt b/CBCT_CT_MRSeg/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/CBCT_CT_MRSeg/model.pt differ diff --git a/CBCT_CT_TS/Evaluation_with_fid.yml b/CBCT_CT_TS/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/CBCT_CT_TS/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/CBCT_CT_TS/Evaluation_with_images.yml b/CBCT_CT_TS/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/CBCT_CT_TS/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/CBCT_CT_TS/Evaluation_with_seg.yml b/CBCT_CT_TS/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/CBCT_CT_TS/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/CBCT_CT_TS/Model.py b/CBCT_CT_TS/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/CBCT_CT_TS/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/ParameterMap_CBCT_generic_TS.txt b/CBCT_CT_TS/ParameterMap_CBCT_generic_TS.txt similarity index 100% rename from ParameterMap_CBCT_generic_TS.txt rename to CBCT_CT_TS/ParameterMap_CBCT_generic_TS.txt diff --git a/CBCT_CT_TS/Prediction.yml b/CBCT_CT_TS/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..eca138748134e500d5ffe965049a6ea058691ab7 --- /dev/null +++ b/CBCT_CT_TS/Prediction.yml @@ -0,0 +1,107 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - ParameterMap_CBCT_generic_TS.txt + models: + - VBoussot/impact-torchscript-models:TS/M852.pt + - VBoussot/impact-torchscript-models:TS/M850.pt + iterations: 1050 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-CBCT-CT-TS + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/CBCT_CT_TS/app.json b/CBCT_CT_TS/app.json new file mode 100644 index 0000000000000000000000000000000000000000..25647b0bb1826e0b484ceaa4dc7537502e08a13f --- /dev/null +++ b/CBCT_CT_TS/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "CBCT/CT preset with TotalSegmentator", + "short_description": "Generic CBCT/CT deformable registration using TotalSegmentator features", + "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 250, and 200 iterations using 2000 random spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), starting with Dice-based overlap on segmentation outputs and progressively integrating feature-level alignment via L1 distances on selected internal layers (0.3/0.7 then 0.5/0.5 L1/Dice), ending with a final purely feature-based stage. A composite objective (IMPACT + mutual information + bending energy penalty) ensures robust cross-modality alignment while enforcing smooth, physically plausible deformations.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/CBCT_CT_TS/install.py b/CBCT_CT_TS/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/CBCT_CT_TS/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/CBCT_CT_TS/model.pt b/CBCT_CT_TS/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/CBCT_CT_TS/model.pt differ diff --git a/Generic_Rigid/Evaluation_with_fid.yml b/Generic_Rigid/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/Generic_Rigid/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/Generic_Rigid/Evaluation_with_images.yml b/Generic_Rigid/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/Generic_Rigid/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/Generic_Rigid/Evaluation_with_seg.yml b/Generic_Rigid/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/Generic_Rigid/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/Generic_Rigid/Model.py b/Generic_Rigid/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/Generic_Rigid/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/Parameters_Rigid.txt b/Generic_Rigid/Parameters_Rigid.txt similarity index 100% rename from Parameters_Rigid.txt rename to Generic_Rigid/Parameters_Rigid.txt diff --git a/Generic_Rigid/Prediction.yml b/Generic_Rigid/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..2184dd44051482372585f629716956e2edd5fbaf --- /dev/null +++ b/Generic_Rigid/Prediction.yml @@ -0,0 +1,105 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - Parameters_Rigid.txt + models: [] + iterations: 250 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-Generic-Rigid + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/Generic_Rigid/app.json b/Generic_Rigid/app.json new file mode 100644 index 0000000000000000000000000000000000000000..cb46b55f020ee19b491f7df88979dbb2b87f30ff --- /dev/null +++ b/Generic_Rigid/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "Generic Rigid", + "short_description": "Rigid registration using mutual information and a multi-resolution pyramid.", + "description": "This preset performs rigid alignment using an Euler transform optimized with Adaptive Stochastic Gradient Descent. It uses a 4-level multi-resolution strategy and Mattes mutual information as similarity metric. Initial alignment based on image centers are enabled to ensure robust convergence for multimodal images.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/Generic_Rigid/install.py b/Generic_Rigid/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/Generic_Rigid/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/Generic_Rigid/model.pt b/Generic_Rigid/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/Generic_Rigid/model.pt differ diff --git a/Generic_Rigid_BSpline/Evaluation_with_fid.yml b/Generic_Rigid_BSpline/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/Generic_Rigid_BSpline/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/Generic_Rigid_BSpline/Evaluation_with_images.yml b/Generic_Rigid_BSpline/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/Generic_Rigid_BSpline/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/Generic_Rigid_BSpline/Evaluation_with_seg.yml b/Generic_Rigid_BSpline/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/Generic_Rigid_BSpline/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/Generic_Rigid_BSpline/Model.py b/Generic_Rigid_BSpline/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/Generic_Rigid_BSpline/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/Parameters_BSpline.txt b/Generic_Rigid_BSpline/Parameters_BSpline.txt similarity index 100% rename from Parameters_BSpline.txt rename to Generic_Rigid_BSpline/Parameters_BSpline.txt diff --git a/Generic_Rigid_BSpline/Parameters_Rigid.txt b/Generic_Rigid_BSpline/Parameters_Rigid.txt new file mode 100644 index 0000000000000000000000000000000000000000..cce67f5856961b9d1732d17795d60fa36a0677ba --- /dev/null +++ b/Generic_Rigid_BSpline/Parameters_Rigid.txt @@ -0,0 +1,147 @@ +// Example parameter file for rotation registration +// C-style comments: // + +// The internal pixel type, used for internal computations +// Leave to float in general. +// NB: this is not the type of the input images! The pixel +// type of the input images is automatically read from the +// images themselves. +// This setting can be changed to "short" to save some memory +// in case of very large 3D images. +(FixedInternalImagePixelType "float") +(MovingInternalImagePixelType "float") + +// **************** Main Components ************************** + +// The following components should usually be left as they are: +(Registration "MultiResolutionRegistration") +(Interpolator "BSplineInterpolator") +(ResampleInterpolator "FinalBSplineInterpolator") +(Resampler "DefaultResampler") + +// These may be changed to Fixed/MovingSmoothingImagePyramid. +// See the manual. +(FixedImagePyramid "FixedRecursiveImagePyramid") +(MovingImagePyramid "MovingRecursiveImagePyramid") + +// The following components are most important: +// The optimizer AdaptiveStochasticGradientDescent (ASGD) works +// quite ok in general. The Transform and Metric are important +// and need to be chosen careful for each application. See manual. +(Optimizer "AdaptiveStochasticGradientDescent") +(Transform "EulerTransform") +(Metric "AdvancedMattesMutualInformation") + +// ***************** Transformation ************************** + +// Scales the rotations compared to the translations, to make +// sure they are in the same range. In general, it's best to +// use automatic scales estimation: +(AutomaticScalesEstimation "true") + +// Automatically guess an initial translation by aligning the +// geometric centers of the fixed and moving. +(AutomaticTransformInitialization "true") + +// Whether transforms are combined by composition or by addition. +// In generally, Compose is the best option in most cases. +// It does not influence the results very much. +(HowToCombineTransforms "Compose") + +//Save composite ITK transform +(ITKTransformOutputFileNameExtension "itk.txt") +(WriteITKCompositeTransform "true") + +// ******************* Similarity measure ********************* + +// Number of grey level bins in each resolution level, +// for the mutual information. 16 or 32 usually works fine. +// You could also employ a hierarchical strategy: +//(NumberOfHistogramBins 16 32 64) +(NumberOfHistogramBins 32) + +// If you use a mask, this option is important. +// If the mask serves as region of interest, set it to false. +// If the mask indicates which pixels are valid, then set it to true. +// If you do not use a mask, the option doesn't matter. +(ErodeMask "false") + +// ******************** Multiresolution ********************** + +// The number of resolutions. 1 Is only enough if the expected +// deformations are small. 3 or 4 mostly works fine. For large +// images and large deformations, 5 or 6 may even be useful. +(NumberOfResolutions 4) + +// The downsampling/blurring factors for the image pyramids. +// By default, the images are downsampled by a factor of 2 +// compared to the next resolution. +// So, in 2D, with 4 resolutions, the following schedule is used: +//(ImagePyramidSchedule 8 8 4 4 2 2 1 1 ) +// And in 3D: +//(ImagePyramidSchedule 8 8 8 4 4 4 2 2 2 1 1 1 ) +// You can specify any schedule, for example: +//(ImagePyramidSchedule 4 4 4 3 2 1 1 1 ) +// Make sure that the number of elements equals the number +// of resolutions times the image dimension. + +// ******************* Optimizer **************************** + +// Maximum number of iterations in each resolution level: +// 200-500 works usually fine for rigid registration. +// For more robustness, you may increase this to 1000-2000. +(MaximumNumberOfIterations 250) + +// The step size of the optimizer, in mm. By default the voxel size is used. +// which usually works well. In case of unusual high-resolution images +// (eg histology) it is necessary to increase this value a bit, to the size +// of the "smallest visible structure" in the image: +//(MaximumStepLength 1.0) + +// **************** Image sampling ********************** + +// Number of spatial samples used to compute the mutual +// information (and its derivative) in each iteration. +// With an AdaptiveStochasticGradientDescent optimizer, +// in combination with the two options below, around 2000 +// samples may already suffice. +(NumberOfSpatialSamples 2048) + +// Refresh these spatial samples in every iteration, and select +// them randomly. See the manual for information on other sampling +// strategies. +(NewSamplesEveryIteration "true") +(ImageSampler "Random") + +// ************* Interpolation and Resampling **************** + +// Order of B-Spline interpolation used during registration/optimisation. +// It may improve accuracy if you set this to 3. Never use 0. +// An order of 1 gives linear interpolation. This is in most +// applications a good choice. +(BSplineInterpolationOrder 1) + +// Order of B-Spline interpolation used for applying the final +// deformation. +// 3 gives good accuracy; recommended in most cases. +// 1 gives worse accuracy (linear interpolation) +// 0 gives worst accuracy, but is appropriate for binary images +// (masks, segmentations); equivalent to nearest neighbor interpolation. +(FinalBSplineInterpolationOrder 3) + +//Default pixel value for pixels that come from outside the picture: +(DefaultPixelValue 0) + +// Choose whether to generate the deformed moving image. +// You can save some time by setting this to false, if you are +// only interested in the final (nonrigidly) deformed moving image +// for example. +(WriteResultImage "false") + +// The pixel type and format of the resulting deformed moving image +(ResultImagePixelType "short") +(ResultImageFormat "mha") +(WriteTransformParametersEachIteration "false") +(WriteTransformParametersEachResolution "false") +(WriteIterationInfo "false") + diff --git a/Generic_Rigid_BSpline/Prediction.yml b/Generic_Rigid_BSpline/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..a2887b51034260566ac579a7e5a62f58726d6a67 --- /dev/null +++ b/Generic_Rigid_BSpline/Prediction.yml @@ -0,0 +1,106 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - Parameters_Rigid.txt + - Parameters_BSpline.txt + models: [] + iterations: 3000 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-Generic-Rigid-BSpline + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/Generic_Rigid_BSpline/app.json b/Generic_Rigid_BSpline/app.json new file mode 100644 index 0000000000000000000000000000000000000000..d8de23977b1f2aaeec179e75a7312292391818fd --- /dev/null +++ b/Generic_Rigid_BSpline/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "Generic Rigid + BSpline", + "short_description": "Two-stage registration: rigid alignment followed by BSpline refinement.", + "description": "A combined registration strategy: first a rigid Euler transform corrects global misalignment, then a BSpline model captures localized anatomical deformations. Both stages use a multi-resolution pyramid, mutual information, and stochastic optimization for robust performance across a wide range of multimodal imaging scenarios.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/Generic_Rigid_BSpline/install.py b/Generic_Rigid_BSpline/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/Generic_Rigid_BSpline/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/Generic_Rigid_BSpline/model.pt b/Generic_Rigid_BSpline/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/Generic_Rigid_BSpline/model.pt differ diff --git a/ImpactReg/Uncertainty.yml b/ImpactReg/Uncertainty.yml deleted file mode 100644 index 9f83eff425eac9bd6396e0384c94bf4926fb2e35..0000000000000000000000000000000000000000 --- a/ImpactReg/Uncertainty.yml +++ /dev/null @@ -1,24 +0,0 @@ -Evaluator: - metrics: - Uncertainty: - targets_criterions: - None: - criterions_loader: - Mean: - name: Uncertainty - Dataset: - groups_src: - Volume_0: - groups_dest: - Uncertainty: - transforms: - UncertaintyDVF:DVFToNorm: {} - Variance: {} - Save: - dataset: ./Uncertainties/ImpactReg/Output:mha - group: None - subset: None - dataset_filenames: - - ./Dataset:mha - validation: None - train_name: ImpactReg diff --git a/ImpactReg/UncertaintyDVF.py b/ImpactReg/UncertaintyDVF.py deleted file mode 100644 index 873a9546d3d2dd99723db54f4d5cdd7ae5d73ba0..0000000000000000000000000000000000000000 --- a/ImpactReg/UncertaintyDVF.py +++ /dev/null @@ -1,15 +0,0 @@ -from konfai.data.transform import Transform -from konfai.utils.dataset import Attribute -import torch - -class DVFToNorm(Transform): - - def __init__(self) -> None: - super().__init__() - - def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - cache_attribute["Origin"] = cache_attribute.pop_np_array("Origin")[:tensors.dim()-2] - cache_attribute["Spacing"] = cache_attribute.pop_np_array("Spacing")[:tensors.dim()-2] - cache_attribute.pop("Direction") - cache_attribute["Direction"] = torch.eye(tensors.dim()-2).flatten() - return torch.linalg.norm(tensors, dim=-1) \ No newline at end of file diff --git a/ImpactReg/app.json b/ImpactReg/app.json deleted file mode 100644 index 5175717b492c9f4bdb6c17319e7f3bbcba645dcd..0000000000000000000000000000000000000000 --- a/ImpactReg/app.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "display_name": "ImpactReg", - "short_description": "", - "description": "", - "tta": 2, - "mc_dropout": false -} \ No newline at end of file diff --git a/MR_CT_HeadNeck/Evaluation_with_fid.yml b/MR_CT_HeadNeck/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/MR_CT_HeadNeck/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_HeadNeck/Evaluation_with_images.yml b/MR_CT_HeadNeck/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/MR_CT_HeadNeck/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_HeadNeck/Evaluation_with_seg.yml b/MR_CT_HeadNeck/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/MR_CT_HeadNeck/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_HeadNeck/Model.py b/MR_CT_HeadNeck/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/MR_CT_HeadNeck/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/ParameterMap_MRI_HN.txt b/MR_CT_HeadNeck/ParameterMap_MRI_HN.txt similarity index 100% rename from ParameterMap_MRI_HN.txt rename to MR_CT_HeadNeck/ParameterMap_MRI_HN.txt diff --git a/MR_CT_HeadNeck/Prediction.yml b/MR_CT_HeadNeck/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..537a1866e2a9694f6ad03c9a10e01126ae00a562 --- /dev/null +++ b/MR_CT_HeadNeck/Prediction.yml @@ -0,0 +1,106 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - ParameterMap_MRI_HN.txt + models: + - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt + iterations: 1050 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-MR-CT-HeadNeck + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/MR_CT_HeadNeck/__pycache__/Model.cpython-314.pyc b/MR_CT_HeadNeck/__pycache__/Model.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64b50b74f3d19d7b16c443f155b35a96367dae20 Binary files /dev/null and b/MR_CT_HeadNeck/__pycache__/Model.cpython-314.pyc differ diff --git a/MR_CT_HeadNeck/app.json b/MR_CT_HeadNeck/app.json new file mode 100644 index 0000000000000000000000000000000000000000..3033bfc5dff5d70a4911d02e83834203920c26db --- /dev/null +++ b/MR_CT_HeadNeck/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "MR/CT Head&Neck", + "short_description": "Optimized preset for MR/CT registration on head & neck", + "description": "A four-level recursive B-spline deformable registration optimized for MRI head-and-neck images, driven by the IMPACT metric and combining semantic features extracted from the pretrained MIND model at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm). The optimization follows a multi-resolution scheme with up to 300, 300, 250, and 200 ASGD iterations, using stochastic sampling and a composite metric (IMPACT + mutual information + bending energy) to achieve robust semantic alignment in complex head-and-neck anatomy.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/MR_CT_HeadNeck/install.py b/MR_CT_HeadNeck/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/MR_CT_HeadNeck/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/MR_CT_HeadNeck/model.pt b/MR_CT_HeadNeck/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/MR_CT_HeadNeck/model.pt differ diff --git a/MR_CT_MRSeg/Evaluation_with_fid.yml b/MR_CT_MRSeg/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/MR_CT_MRSeg/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_MRSeg/Evaluation_with_images.yml b/MR_CT_MRSeg/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/MR_CT_MRSeg/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_MRSeg/Evaluation_with_seg.yml b/MR_CT_MRSeg/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/MR_CT_MRSeg/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_MRSeg/Model.py b/MR_CT_MRSeg/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/MR_CT_MRSeg/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/ParameterMap_MRI_MRSeg.txt b/MR_CT_MRSeg/ParameterMap_MRI_MRSeg.txt similarity index 100% rename from ParameterMap_MRI_MRSeg.txt rename to MR_CT_MRSeg/ParameterMap_MRI_MRSeg.txt diff --git a/MR_CT_MRSeg/Prediction.yml b/MR_CT_MRSeg/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..6ef91686be41cf5d88a88aa199880d41398c928a --- /dev/null +++ b/MR_CT_MRSeg/Prediction.yml @@ -0,0 +1,107 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - ParameterMap_MRI_MRSeg.txt + models: + - VBoussot/impact-torchscript-models:MRSeg/MRSeg.pt + - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt + iterations: 1100 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-MR-CT-MRSeg + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/MR_CT_MRSeg/app.json b/MR_CT_MRSeg/app.json new file mode 100644 index 0000000000000000000000000000000000000000..f93d3b835eeb14e31c6baf0eb6c71f3dd829d57c --- /dev/null +++ b/MR_CT_MRSeg/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "MR/CT preset with MRSegmentator and MIND", + "short_description": "Generic MR/CT deformable registration using MIND + MRSegmentator features", + "description": "A four-level recursive B-spline deformable registration optimized for generic MR/CT alignment, driven by the IMPACT metric and combining semantic features from two pretrained models: MIND (L1 distance on a subset of 32 features) and MRSegmentator (Dice overlap on segmentation outputs with 64 features). Features are extracted at progressively finer voxel scales with level-dependent weighting between MIND and MRSegmentator. The optimization follows a multi-resolution ASGD scheme with a composite objective (IMPACT + mutual information + bending energy penalty) to ensure robust cross-modality semantic alignment and smooth deformations.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/MR_CT_MRSeg/install.py b/MR_CT_MRSeg/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/MR_CT_MRSeg/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/MR_CT_MRSeg/model.pt b/MR_CT_MRSeg/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/MR_CT_MRSeg/model.pt differ diff --git a/MR_CT_TS/Evaluation_with_fid.yml b/MR_CT_TS/Evaluation_with_fid.yml new file mode 100644 index 0000000000000000000000000000000000000000..267c860058b1bcbbfe85cc6ae314289089b9ed23 --- /dev/null +++ b/MR_CT_TS/Evaluation_with_fid.yml @@ -0,0 +1,22 @@ +Evaluator: + metrics: + FixedFid: + targets_criterions: + MovingFid: + criterions_loader: + TRE: {} + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedFid: + transforms: None + Reference_0: + groups_dest: + MovingFid: + transforms: None + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_TS/Evaluation_with_images.yml b/MR_CT_TS/Evaluation_with_images.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce145cfb0b192646149642d91f15e53da8a5bfe6 --- /dev/null +++ b/MR_CT_TS/Evaluation_with_images.yml @@ -0,0 +1,35 @@ +Evaluator: + metrics: + FixedImage: + targets_criterions: + MovingImage;Mask: + criterions_loader: + MAESaveMap: + reduction: mean + dataset: ./Evaluations/ImpactReg/Output:mha + group: MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedImage: + transforms: + TensorCast: + dtype: float32 + Reference_0: + groups_dest: + MovingImage: + transforms: + TensorCast: + dtype: float32 + Mask_0: + groups_dest: + Mask: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_TS/Evaluation_with_seg.yml b/MR_CT_TS/Evaluation_with_seg.yml new file mode 100644 index 0000000000000000000000000000000000000000..6221238dd64ccead3148eb85730feb73df4d8fde --- /dev/null +++ b/MR_CT_TS/Evaluation_with_seg.yml @@ -0,0 +1,29 @@ +Evaluator: + metrics: + FixedSeg: + targets_criterions: + MovingSeg: + criterions_loader: + DiceSaveMap: + labels: None + dataset: ./Evaluations/ImpactReg/Output:mha + group: Seg_MAE_map + Dataset: + groups_src: + Volume_0: + groups_dest: + FixedSeg: + transforms: + TensorCast: + dtype: uint8 + Reference_0: + groups_dest: + MovingSeg: + transforms: + TensorCast: + dtype: uint8 + subset: None + dataset_filenames: + - ./Dataset:mha + validation: None + train_name: ImpactReg diff --git a/MR_CT_TS/Model.py b/MR_CT_TS/Model.py new file mode 100644 index 0000000000000000000000000000000000000000..389d54ede2f15e52b531caeb9c17df9716ba888d --- /dev/null +++ b/MR_CT_TS/Model.py @@ -0,0 +1,291 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph). + +``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image +(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled +onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes`` +(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the +``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine +needs to register in physical space. + +Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode: + + Patch.patch_size: None # one patch per case (registration is global) + batch_size: 1 # one fixed/moving pair at a time + +NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on +runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution. +""" + +import os +import re +import shutil +import subprocess # nosec B404 +import tempfile +from pathlib import Path + +import numpy as np +import SimpleITK as sitk +import torch +import tqdm +from huggingface_hub import hf_hub_download +from install import get_elastix_bin, install_elastix_impact, try_elastix +from konfai.network import network +from konfai.utils.dataset import Attribute, data_to_image, image_to_data + +# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs. +# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download. +ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact" + + +class ElastixEngine: + """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + + NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix + does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. + """ + + def __init__(self, parameter_maps: list[str], models: list[str], iterations: int) -> None: + self._bundle_dir = Path(__file__).resolve().parent + self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps] + self._models = models + self._iterations = iterations + self._elastix_bin = self._ensure_binary() + self._local_models = self._download_models() + + def _ensure_binary(self) -> Path: + # Optional override: point at an existing elastix-IMPACT install (skips the download). + override = os.environ.get("KONFAI_ELASTIX_DIR", "") + if override: + try_elastix(Path(override)) + return get_elastix_bin(Path(override)).resolve() + ELASTIX_CACHE.mkdir(parents=True, exist_ok=True) + try: + try_elastix(ELASTIX_CACHE) + except Exception: + install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False) + try_elastix(ELASTIX_CACHE) + return get_elastix_bin(ELASTIX_CACHE).resolve() + + def _download_models(self) -> list[tuple[str, Path]]: + """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path).""" + models = [] + for ref in self._models: + repo, filename = ref.split(":", 1) + local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) # nosec B615 + models.append((filename, local)) + return models + + def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]: + """Copy parameter maps into the work dir, patching the ImpactGPU device field when present.""" + staged = [] + for src in self._parameter_maps: + dst = work / src.name + lines = [] + for line in src.read_text(encoding="utf-8").splitlines(): + if line.strip().startswith("(ImpactGPU"): + line = f"(ImpactGPU {device_index})" + lines.append(line) + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + staged.append(dst) + return staged + + def register( + self, + fixed: sitk.Image, + moving: sitk.Image, + device_index: int, + fixed_mask: sitk.Image | None = None, + moving_mask: sitk.Image | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + + Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region + (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. + """ + work = Path(tempfile.mkdtemp(prefix="konfai_reg_")) + try: + fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha" + sitk.WriteImage(fixed, str(fixed_path)) + sitk.WriteImage(moving, str(moving_path)) + + # Stage the feature models at the relative path the parameter maps reference + # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory. + for rel_name, model_path in self._local_models: + dst = work / rel_name + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(model_path) + + args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] + for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + if mask is not None: + mask_path = work / name + sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) + args += [flag, str(mask_path)] + args += ["-out", str(work)] + for pmap in self._stage_parameter_maps(work, device_index): + args += ["-p", str(pmap)] + + # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses + # the "N% done/total" progress line) shows real progress during the long registration. + # Make the elastix binary's own libs (bundled libtorch under /lib) and any extra + # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads. + env = os.environ.copy() + extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")] + env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p) + proc = subprocess.Popen( # nosec B603 + args, + cwd=str(work), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + captured: list[str] = [] + iteration_line = re.compile(r"^\d+\s") + # ``iterations`` is the total iteration budget declared for the preset (summed over the + # chained parameter maps), so the bar spans the whole chain of registration stages. + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=True) + assert proc.stdout is not None + for line in proc.stdout: + captured.append(line) + if iteration_line.match(line): + progress.update(1) + progress.close() + returncode = proc.wait() + if returncode != 0: + raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}") + + transforms = sorted( + work.glob("TransformParameters.*-Composite.itk.txt"), + key=lambda p: int(p.name.split(".")[1].split("-")[0]), + ) + if not transforms: + raise FileNotFoundError("elastix produced no composite transform file.") + transform = sitk.ReadTransform(str(transforms[-1])) + + moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) + dvf = sitk.TransformToDisplacementField( + transform, + sitk.sitkVectorFloat64, + fixed.GetSize(), + fixed.GetOrigin(), + fixed.GetSpacing(), + fixed.GetDirection(), + ) + moved_np, _ = image_to_data(moved) + dvf_np, _ = image_to_data(dvf) + return moved_np, dvf_np + finally: + shutil.rmtree(work, ignore_errors=True) + + +class ElastixRegistration(torch.nn.Module): + """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid. + + ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch + ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix + needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry. + """ + + accepts_attributes = True + + def __init__(self, engine: str, parameter_maps: list[str], models: list[str], iterations: int) -> None: + super().__init__() + if engine != "elastix": + raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.") + self._engine = ElastixEngine(parameter_maps, models, iterations) + + def forward( + self, + fixed: torch.Tensor, + moving: torch.Tensor, + fixed_mask: torch.Tensor, + moving_mask: torch.Tensor, + attributes: list[list[Attribute]], + ) -> torch.Tensor: + # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch. + # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field + # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering + # the whole image (the auto-filled default when the user supplies none) restricts nothing. + fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes + device_index = fixed.device.index if fixed.device.type == "cuda" else -1 + combined = [] + for b in range(fixed.shape[0]): + fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) + moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) + fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) + moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) + moved_np, dvf_np = self._engine.register( + fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img + ) + combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + return torch.stack(combined, dim=0).to(fixed.device) + + +class ChannelSelect(torch.nn.Module): + """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF).""" + + def __init__(self, start: int, stop: int) -> None: + super().__init__() + self._start = start + self._stop = stop + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor[:, self._start : self._stop] + + +class RegistrationNet(network.Network): + """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, + fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction). + + Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and + ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces + both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by + ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. + """ + + def __init__( + self, + optimizer: network.OptimizerLoader = network.OptimizerLoader(), + schedulers: dict[str, network.LRSchedulersLoader] = { + "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) + }, + outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, + engine: str = "elastix", + parameter_maps: list[str] = [], + models: list[str] = [], + iterations: int = 1000, + ) -> None: + super().__init__( + in_channels=1, + optimizer=optimizer, + schedulers=schedulers, + outputs_criterions=outputs_criterions, + dim=3, + ) + self.add_module( + "Registration", + ElastixRegistration(engine, parameter_maps, models, iterations), + in_branch=[0, 1, 2, 3], + out_branch=["registration"], + ) + self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) + self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) diff --git a/ParameterMap_MRI_TS.txt b/MR_CT_TS/ParameterMap_MRI_TS.txt similarity index 100% rename from ParameterMap_MRI_TS.txt rename to MR_CT_TS/ParameterMap_MRI_TS.txt diff --git a/MR_CT_TS/Prediction.yml b/MR_CT_TS/Prediction.yml new file mode 100644 index 0000000000000000000000000000000000000000..806bb17d03f12d61c943b246868e1e9dd7f46afb --- /dev/null +++ b/MR_CT_TS/Prediction.yml @@ -0,0 +1,108 @@ +Predictor: + Model: + classpath: Model:RegistrationNet + RegistrationNet: + engine: elastix + parameter_maps: + - ParameterMap_MRI_TS.txt + models: + - VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt + - VBoussot/impact-torchscript-models:TS/M852.pt + - VBoussot/impact-torchscript-models:TS/M850.pt + iterations: 1100 + outputs_criterions: None + Dataset: + groups_src: + Volume_0: + groups_dest: + Fixed: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_1: + groups_dest: + Moving: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_2: + groups_dest: + FixedMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + Volume_3: + groups_dest: + MovingMask: + transforms: + TensorCast: + dtype: float32 + inverse: false + patch_transforms: None + is_input: true + augmentations: None + Patch: + patch_size: None + overlap: None + mask: None + pad_value: None + extend_slice: 0 + subset: None + filter: None + dataset_filenames: + - ./Dataset/:mha + use_cache: false + batch_size: 1 + num_workers: None + pin_memory: false + prefetch_factor: None + persistent_workers: None + outputs_dataset: + MovedImage: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: Moved:mha + group: Moved + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + DisplacementField: + OutputDataset: + name_class: OutSameAsGroupDataset + before_reduction_transforms: None + after_reduction_transforms: None + final_transforms: + TensorCast: + dtype: float32 + inverse: false + dataset_filename: DVF:mha + group: DVF + same_as_group: Volume_0:Fixed + patch_combine: None + inverse_transform: false + reduction: Mean + Mean: {} + train_name: ImpactReg-MR-CT-TS + manual_seed: 42 + gpu_checkpoints: None + images_log: None + combine: Mean + autocast: false + data_log: None diff --git a/MR_CT_TS/app.json b/MR_CT_TS/app.json new file mode 100644 index 0000000000000000000000000000000000000000..3f9081939d8f0de64084db317a1a18ffd1a38dd1 --- /dev/null +++ b/MR_CT_TS/app.json @@ -0,0 +1,94 @@ +{ + "display_name": "MR/CT preset with TotalSegmentator and MIND", + "short_description": "Generic MR/CT deformable registration using MIND + Totalsegmentator features", + "description": "A four-level recursive B-spline deformable registration optimized for generic MR/CT alignment, driven by the IMPACT metric and combining semantic features from two pretrained models: MIND (L1 distance on a subset of 32 features) and Totalsegmentator (Dice overlap on segmentation outputs with 64 features). Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm) with level-dependent weighting between MIND and MRSegmentator (0.2/0.8, 0.3/0.7, 0.6/0.4, 0.7/0.3). The optimization follows a multi-resolution ASGD scheme with up to 400, 300, 200, and 200 iterations using 2000 random spatial samples, and a composite objective (IMPACT + mutual information + bending energy penalty) to ensure robust cross-modality semantic alignment and smooth deformations.", + "task": "registration", + "tta": 0, + "mc_dropout": 0, + "models": [ + "model.pt" + ], + "inputs": { + "Fixed": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "Moving": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "FixedMask": { + "display_name": "Fixed mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + }, + "MovingMask": { + "display_name": "Moving mask (optional)", + "volume_type": "SEGMENTATION", + "required": false + } + }, + "outputs": { + "MovedImage": { + "display_name": "Moved image", + "volume_type": "VOLUME", + "required": true + }, + "DisplacementField": { + "display_name": "Displacement field", + "volume_type": "VOLUME", + "required": false + } + }, + "inputs_evaluations": { + "Image": { + "Evaluation_with_images.yml": { + "FixedImage": { + "display_name": "Fixed image", + "volume_type": "VOLUME", + "required": true + }, + "MovingImage": { + "display_name": "Moving image", + "volume_type": "VOLUME", + "required": true + }, + "Mask": { + "display_name": "Evaluation mask", + "volume_type": "SEGMENTATION", + "required": false + } + } + }, + "Segmentation": { + "Evaluation_with_seg.yml": { + "FixedSeg": { + "display_name": "Fixed segmentation", + "volume_type": "SEGMENTATION", + "required": true + }, + "MovingSeg": { + "display_name": "Moving segmentation", + "volume_type": "SEGMENTATION", + "required": true + } + } + }, + "Landmarks": { + "Evaluation_with_fid.yml": { + "FixedFid": { + "display_name": "Fixed landmarks", + "volume_type": "FIDUCIALS", + "required": true + }, + "MovingFid": { + "display_name": "Moving landmarks", + "volume_type": "FIDUCIALS", + "required": true + } + } + } + } +} diff --git a/MR_CT_TS/install.py b/MR_CT_TS/install.py new file mode 100644 index 0000000000000000000000000000000000000000..9af2b88d5553b2262a8c622def90be1cb141e224 --- /dev/null +++ b/MR_CT_TS/install.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import platform +import re +import shutil +import stat +import subprocess # nosec B404 +import zipfile +from pathlib import Path + +import requests +from tqdm import tqdm + +# ----------------------------------------------------------------------------- +# Elastix + IMPACT binary assets hosted on GitHub Releases. +# +# Key format: (OS, ARCH, FLAVOR) +# - OS : platform.system() -> "Linux", "Windows", "Darwin" +# - ARCH : normalized architecture -> "x86_64" +# - FLAVOR : "cpu" or "cu128" +# +# CPU assets are standalone (LibTorch CPU bundled). +# CUDA assets do NOT bundle LibTorch (too large for GitHub limits). +# ----------------------------------------------------------------------------- +ELX_ASSET_TEMPLATE = { + ("Linux", "x86_64", "cpu"): "elastix-impact-linux-x86_64-cpu.zip", + ("Linux", "x86_64", "cu128"): "elastix-impact-linux-x86_64-cu128.zip", + ("Windows", "x86_64", "cpu"): "elastix-impact-windows-x86_64-cpu.zip", + ("Windows", "x86_64", "cu128"): "elastix-impact-windows-x86_64-cu128.zip", + ("Darwin", "x86_64", "cpu"): "elastix-impact-macos-14-x86_64-cpu.zip", +} + +# ----------------------------------------------------------------------------- +# Official LibTorch downloads (PyTorch). +# +# IMPORTANT: +# - Version MUST match the one used at build time (ABI compatibility). +# - Using "shared-with-deps" ensures CUDA runtime libraries are included +# (except the NVIDIA driver, which must be installed system-wide). +# ----------------------------------------------------------------------------- + +LIBTORCH_URL = { + ( + "Linux", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Windows", + "x86_64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Darwin", + "arm64", + "cpu", + ): "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.8.0%2Bcpu.zip", + ( + "Linux", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-2.8.0%2Bcu128.zip", + ( + "Windows", + "x86_64", + "cu128", + ): "https://download.pytorch.org/libtorch/cu128/libtorch-win-shared-with-deps-2.8.0%2Bcu128.zip", +} + +# ----------------------------------------------------------------------------- +# Minimum NVIDIA driver versions required for CUDA 12.8. +# +# The CUDA Toolkit itself is NOT required. +# Only a sufficiently recent NVIDIA driver must be installed. +# ----------------------------------------------------------------------------- +CUDA128_MIN_DRIVER_LINUX = (570, 26) +CUDA128_MIN_DRIVER_WINDOWS = (570, 65) + +GITHUB_OWNER = "vboussot" +GITHUB_REPO = "ImpactElastix" +GITHUB_TAG = "1.0.0" + + +DEFAULT_PREFIX = Path.cwd() / "elastix-impact" + + +def run_cmd(cmd: list[str]) -> str: + return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8", errors="replace") # nosec B603 + + +def detect_nvidia_driver() -> tuple[bool, tuple[int, int] | None]: + """ + Detect presence of an NVIDIA GPU and extract the driver version. + Returns: + (True, (major, minor)) if detected + (False, None) if nvidia-smi is not available + """ + try: + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi not found in PATH") + + out = ( + subprocess.check_output( + [nvidia_smi, "--query-gpu=driver_version", "--format=csv,noheader"], stderr=subprocess.DEVNULL + ) # nosec B603 + .decode("utf-8") + .strip() + ) + except Exception: + return (False, None) + + # Exemple: "575.64" + m = re.match(r"([0-9]+)\.([0-9]+)", out) + if not m: + return (True, None) + + return (True, (int(m.group(1)), int(m.group(2)))) + + +def driver_ok_for_cuda(os_name: str, drv: tuple[int, int] | None) -> bool: + """ + Check whether the detected NVIDIA driver satisfies the minimum + requirement for CUDA 12.8 on the given operating system. + """ + if drv is None: + return False + if os_name == "Linux": + return drv >= CUDA128_MIN_DRIVER_LINUX + if os_name == "Windows": + return drv >= CUDA128_MIN_DRIVER_WINDOWS + return False + + +def normalize_arch(machine: str) -> str: + """ + Normalize platform.machine() output across operating systems. + + Examples: + AMD64 -> x86_64 + aarch64 -> arm64 + """ + m = machine.lower() + if m in ("x86_64", "amd64"): + return "x86_64" + if m in ("aarch64", "arm64"): + return "arm64" + return machine + + +def download_file(url: str, dst: Path) -> None: + """ + Download a file from the given URL with a progress indicator. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading: {url}", flush=True) + print(f" to: {dst}", flush=True) + + try: + with requests.get(url, stream=True, timeout=60) as r: + r.raise_for_status() + total = int(r.headers.get("content-length", 0)) + with open(dst, "wb") as f: + with tqdm( + total=total, + unit="B", + unit_scale=True, + desc=f"Downloading {dst.name}", + ) as pbar: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + except Exception as e: + raise e + + +def extract_archive(archive: Path, dst_dir: Path) -> None: + """ + Extract a ZIP archive to the destination directory. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"Extracting: {archive} -> {dst_dir}", flush=True) + with zipfile.ZipFile(archive, "r") as z: + z.extractall(dst_dir) + archive.unlink() + + +def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool): + os_name = platform.system() + arch = normalize_arch(platform.machine()) + has_nvidia, drv = detect_nvidia_driver() + + if os_name not in ("Linux", "Windows", "Darwin"): + raise NameError(f"Unsupported OS: {os_name}") + + if arch not in ("x86_64", "arm64"): + raise NameError(f"Unsupported arch: {arch} (expected x86_64, arm64)") + + flavor = "cpu" + if force_cuda: + if not has_nvidia or not driver_ok_for_cuda(os_name, drv): + raise NameError( + "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}" + ) + + flavor = "cu128" + elif not force_cpu: + if has_nvidia and driver_ok_for_cuda(os_name, drv): + flavor = "cu128" + else: + flavor = "cpu" + + print(f"System: {os_name} {arch}", flush=True) + print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True) + print(f"Selected flavor: {flavor}", flush=True) + + key = (os_name, arch, flavor) + if key not in ELX_ASSET_TEMPLATE: + raise NameError(f"No elastix asset configured for {key}") + if flavor != "cpu" and key not in LIBTORCH_URL: + raise NameError(f"No libtorch url configured for {key}") + + install_path = install_path.resolve() + install_path.mkdir(parents=True, exist_ok=True) + + elx_asset = ELX_ASSET_TEMPLATE[key] + elx_url = f"https://github.com/{GITHUB_OWNER}/{GITHUB_REPO}/releases/download/{GITHUB_TAG}/{elx_asset}" + elx_archive = install_path / elx_asset + download_file(elx_url, elx_archive) + extract_archive(elx_archive, install_path) + + # ------------------------------------------------------------------------- + # ZIP archives may drop executable permissions. + # Ensure elastix and transformix are executable on Unix platforms. + # ------------------------------------------------------------------------- + if os_name in ("Linux", "Darwin"): + for exe in ("elastix", "transformix"): + p = install_path / "bin" / exe + if p.exists(): + p.chmod(p.stat().st_mode | stat.S_IEXEC) + + # --------------------------------------------------------------------- + # - Download matching LibTorch 12.8 with or without CUDA 12.8 + # - Extract it + # - Move runtime libraries to a location visible to the dynamic loader + # + # Linux / macOS : prefix/lib + # Windows : prefix (next to elastix.exe) + # --------------------------------------------------------------------- + lt_url = LIBTORCH_URL[key] + lt_archive = install_path / Path(lt_url).name + download_file(lt_url, lt_archive) + extract_archive(lt_archive, install_path) + for p in (install_path / "libtorch" / "lib").iterdir(): + if ".so" in p.name or ".dll" in p.name or ".dylib" in p.name: + shutil.move(p, (install_path if os_name == "Windows" else install_path / "lib") / p.name) + + shutil.rmtree(install_path / "libtorch") + if os_name == "Linux" and flavor == "cu128": + shutil.copy(install_path / "lib" / "libcudart-c3a75b33.so.12", install_path / "lib" / "libcudart.so.12") + + +def get_elastix_bin(install_path: Path) -> Path: + return install_path / ("elastix.exe" if platform.system() == "Windows" else (Path("bin") / "elastix")) + + +def try_elastix(install_path: Path) -> None: + try: + subprocess.run( + [str(get_elastix_bin(install_path)), "-h"], + capture_output=True, + text=True, + check=True, + ) # nosec B603 + except subprocess.CalledProcessError as e: + msg = "Elastix execution failed.\n\n" + + msg += f"Command:\n{' '.join(e.cmd)}\n" + msg += f"Return code: {e.returncode}\n\n" + + if e.stderr: + msg += "Error output:\n" + msg += e.stderr.strip() + raise NameError(msg) + + except OSError as e: + msg = ( + "Elastix could not be started.\n\n" + "This is usually caused by missing shared libraries " + "(e.g. LibTorch or CUDA runtime).\n\n" + f"System error:\n{str(e)}" + ) + + raise NameError(msg) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--install-path", type=Path, default=DEFAULT_PREFIX, help="Install directory (default: ./elastix-impact)" + ) + group = ap.add_mutually_exclusive_group() + group.add_argument("--force-cpu", action="store_true", help="Force CPU install even if NVIDIA present") + group.add_argument("--force-cuda", action="store_true", help="Force CUDA install (fails if no suitable driver)") + args = vars(ap.parse_args()) + install_elastix_impact(**args) + + +if __name__ == "__main__": + main() diff --git a/MR_CT_TS/model.pt b/MR_CT_TS/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..123acc2698166f388f0687f3cbce0daa252acc57 Binary files /dev/null and b/MR_CT_TS/model.pt differ diff --git a/PresetDatabase.json b/PresetDatabase.json deleted file mode 100644 index 764bbc754cd3a011fb623c45f73113a003f76a4c..0000000000000000000000000000000000000000 --- a/PresetDatabase.json +++ /dev/null @@ -1,119 +0,0 @@ -{ - "presets": [ - { - "type": "elastix", - "display_name": "Generic Rigid + BSpline", - "parameter_maps": [ - "Parameters_Rigid.txt", - "Parameters_BSpline.txt" - ], - "models": [], - "preprocess_function": "", - "iterations": 3000, - "short_description": "Two-stage registration: rigid alignment followed by BSpline refinement.", - "description": "A combined registration strategy: first a rigid Euler transform corrects global misalignment, then a BSpline model captures localized anatomical deformations. Both stages use a multi-resolution pyramid, mutual information, and stochastic optimization for robust performance across a wide range of multimodal imaging scenarios." - }, - { - "type": "elastix", - "display_name": "Generic Rigid", - "parameter_maps": [ - "Parameters_Rigid.txt" - ], - "models": [], - "preprocess_function": "", - "iterations": 1000, - "short_description": "Rigid registration using mutual information and a multi-resolution pyramid.", - "description": "This preset performs rigid alignment using an Euler transform optimized with Adaptive Stochastic Gradient Descent. It uses a 4-level multi-resolution strategy and Mattes mutual information as similarity metric. Initial alignment based on image centers are enabled to ensure robust convergence for multimodal images." - }, - { - "type": "elastix", - "display_name": "MR/CT Head&Neck", - "parameter_maps": [ - "ParameterMap_MRI_HN.txt" - ], - "models": [ - "VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt" - ], - "preprocess_function": "", - "iterations": 1050, - "short_description": "Optimized preset for MR/CT registration on head & neck", - "description": "A four-level recursive B-spline deformable registration optimized for MRI head-and-neck images, driven by the IMPACT metric and combining semantic features extracted from the pretrained MIND model at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm). The optimization follows a multi-resolution scheme with up to 300, 300, 250, and 200 ASGD iterations, using stochastic sampling and a composite metric (IMPACT + mutual information + bending energy) to achieve robust semantic alignment in complex head-and-neck anatomy." - }, - { - "type": "elastix", - "display_name": "MR/CT preset with MRSegmentator and MIND", - "parameter_maps": [ - "ParameterMap_MRI_MRSeg.txt" - ], - "models": [ - "VBoussot/impact-torchscript-models:MRSeg/MRSeg.pt", - "VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt" - ], - "preprocess_function": "", - "iterations": 1100, - "short_description": "Generic MR/CT deformable registration using MIND + MRSegmentator features", - "description": "A four-level recursive B-spline deformable registration optimized for generic MR/CT alignment, driven by the IMPACT metric and combining semantic features from two pretrained models: MIND (L1 distance on a subset of 32 features) and MRSegmentator (Dice overlap on segmentation outputs with 64 features). Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm) with level-dependent weighting between MIND and MRSegmentator (0.2/0.8, 0.3/0.7, 0.6/0.4, 0.7/0.3). The optimization follows a multi-resolution ASGD scheme with up to 400, 300, 200, and 200 iterations using 2000 random spatial samples, and a composite objective (IMPACT + mutual information + bending energy penalty) to ensure robust cross-modality semantic alignment and smooth deformations." - }, - { - "type": "elastix", - "display_name": "MR/CT preset with TotalSegmentator and MIND", - "parameter_maps": [ - "ParameterMap_MRI_TS.txt" - ], - "models": [ - "VBoussot/impact-torchscript-models:TS/M852.pt", - "VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt" - ], - "preprocess_function": "", - "iterations": 1100, - "short_description": "Generic MR/CT deformable registration using MIND + Totalsegmentator features", - "description": "A four-level recursive B-spline deformable registration optimized for generic MR/CT alignment, driven by the IMPACT metric and combining semantic features from two pretrained models: MIND (L1 distance on a subset of 32 features) and Totalsegmentator (Dice overlap on segmentation outputs with 64 features). Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm) with level-dependent weighting between MIND and MRSegmentator (0.2/0.8, 0.3/0.7, 0.6/0.4, 0.7/0.3). The optimization follows a multi-resolution ASGD scheme with up to 400, 300, 200, and 200 iterations using 2000 random spatial samples, and a composite objective (IMPACT + mutual information + bending energy penalty) to ensure robust cross-modality semantic alignment and smooth deformations." - }, - { - "type": "elastix", - "display_name": "CBCT/CT preset with MRSegmentator", - "parameter_maps": [ - "ParameterMap_CBCT_generic_MRSeg.txt" - ], - "models": [ - "VBoussot/impact-torchscript-models:MRSeg/MRSeg.pt", - "VBoussot/impact-torchscript-models:MIND/R1D2_3D.pt" - ], - "preprocess_function": "", - "iterations": 1050, - "short_description": "Generic CBCT/CT deformable registration using MRSegmentator features", - "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from the pretrained MRSegmentator model. The scheme follows a multi-resolution strategy with up to 300, 300, 250, and 200 ASGD iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), with a level-dependent combination of Dice-based segmentation overlap and L1 feature distances on selected internal layers of MRSegmentator. Early levels rely on pure Dice supervision, while finer stages progressively integrate feature-level alignment with increasing L1 contribution (0.3/0.7, 0.5/0.5) and a final purely feature-based stage. The optimization minimizes a composite objective (IMPACT + mutual information + bending energy penalty), enabling robust cross-modality alignment between CBCT and CT while enforcing smooth, physically plausible deformations." - }, - { - "type": "elastix", - "display_name": "CBCT/CT preset with TotalSegmentator", - "parameter_maps": [ - "ParameterMap_CBCT_generic_TS.txt" - ], - "models": [ - "VBoussot/impact-torchscript-models:TS/M852.pt", - "VBoussot/impact-torchscript-models:TS/M850.pt" - ], - "preprocess_function": "", - "iterations": 1050, - "short_description": "Generic CBCT/CT deformable registration using TotalSegmentator features", - "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 250, and 200 iterations using 2000 random spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), starting with Dice-based overlap on segmentation outputs and progressively integrating feature-level alignment via L1 distances on selected internal layers (0.3/0.7 then 0.5/0.5 L1/Dice), ending with a final purely feature-based stage. A composite objective (IMPACT + mutual information + bending energy penalty) ensures robust cross-modality alignment while enforcing smooth, physically plausible deformations." - }, - { - "type": "elastix", - "display_name": "CBCT/CT Head&Neck", - "parameter_maps": [ - "ParameterMap_CBCT_HN.txt" - ], - "models": [ - "VBoussot/impact-torchscript-models:TS/M732.pt", - "VBoussot/impact-torchscript-models:TS/M731.pt", - "VBoussot/impact-torchscript-models:TS/M730.pt" - ], - "preprocess_function": "", - "iterations": 1150, - "short_description": "Optimized preset for CBCT/CT registration on head & neck", - "description": "A five-level recursive B-spline deformable registration optimized for CBCT/CT head-and-neck alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 200, 200, and 150 iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm) using L1 distances on selected internal layers of the network. A composite objective (IMPACT + mutual information + bending energy penalty, with increased MI weight) ensures robust cross-modality alignment in complex head-and-neck anatomy while enforcing smooth, physically plausible deformations." - } - ] -}