ENH: Add the Torch-backed dense registration pipeline + examples

A fast, self-contained two-stage registration built on the IMPACT core, in the spirit of
ConvexAdam but driven by anatomical features rather than raw intensities:

  - itk::ImpactCoarseRegistration  stage 1 (coarse): a discrete SSD cost volume over a dense
                                   displacement window on a pooled coarse grid, coupled-convex
                                   global regularization, upsampled to a full-resolution
                                   initial field. Robust to large misalignments; 2D and 3D.
  - itk::ImpactFineRegistration    stage 2 (fine): holds the displacement field as a GPU leaf
                                   tensor, warps with grid_sample, and minimizes an IMPACT
                                   feature (or intensity) loss plus a diffusion regularizer
                                   with torch::optim::Adam — the whole warp + comparison +
                                   optimizer in a single autograd graph. Optional low-resolution
                                   control grid and PCA channel reduction.

Both emit a geometry-correct itk::DisplacementFieldTransform (physical millimetres,
fixed->moving). Adds the Python wrapping, the GoogleTest backend suite (with a tiny committed
TorchScript toy model), and C++/Python metric-registration examples.

Author: vboussot

examples/CMakeLists.txt

@@ -0,0 +1,36 @@
+#==========================================================================
+#
+#   Copyright NumFOCUS
+#
+#   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
+#
+#          https://www.apache.org/licenses/LICENSE-2.0.txt
+#
+#   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.
+#
+#==========================================================================
+
+cmake_minimum_required(VERSION 3.22.1...3.29.0 FATAL_ERROR)
+project(ImpactExamples)
+
+set(ExampleSpecificComponents
+  Impact
+  )
+
+if(NOT ITK_SOURCE_DIR)
+  find_package(ITK REQUIRED COMPONENTS ITKImageIO ITKTransformIO ${ExampleSpecificComponents})
+else()
+  # When being built as part of ITK, ITKImageIO and ITKTransformIO
+  # lists of modules are not yet ready, causing a configure error
+  find_package(ITK REQUIRED COMPONENTS ${ExampleSpecificComponents})
+endif()
+include(${ITK_USE_FILE})
+
+add_executable(ImpactMetricExample ImpactMetricExample.cxx )
+target_link_libraries(ImpactMetricExample ${ITK_LIBRARIES})

examples/ImpactMetricExample.cxx

@@ -0,0 +1,137 @@
+/*=========================================================================
+ *
+ *  Copyright NumFOCUS
+ *
+ *  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
+ *
+ *         https://www.apache.org/licenses/LICENSE-2.0.txt
+ *
+ *  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.
+ *
+ *=========================================================================*/
+
+// IMPACT: semantic similarity registration from pretrained TorchScript features.
+//
+// This example mirrors ImpactMetricExample.py and shows the two ITK-facing pieces of
+// the module. No torch type ever crosses the API: it speaks only ITK/STL types (images,
+// a model-path string, a device string).
+//
+//   1. the core:   itk::ImageToFeaturesMap extracts a feature map from an image;
+//   2. the metric: itk::ImpactImageToImageMetricv4 plugs into
+//                  itk::ImageRegistrationMethodv4 to register two images by comparing
+//                  those features instead of raw intensities.
+
+#include "itkImageToFeaturesMap.h"
+#include "itkImpactImageToImageMetricv4.h"
+#include "itkModelConfiguration.h"
+
+#include <itkImageFileReader.h>
+#include <itkImageFileWriter.h>
+#include <itkImageRegistrationMethodv4.h>
+#include <itkRegularStepGradientDescentOptimizerv4.h>
+#include <itkTranslationTransform.h>
+#include <itkResampleImageFilter.h>
+#include <itkBSplineInterpolateImageFunction.h>
+
+int
+main(int argc, char * argv[])
+{
+  if (argc < 5)
+  {
+    std::cerr << "IMPACT: semantic similarity registration from pretrained features.\n";
+    std::cerr << "Usage: " << argv[0] << " model.pt fixedImage movingImage outputWarpedImage [device]\n";
+    std::cerr << "  device: \"cpu\" (default), \"cuda\", \"cuda:0\", ...\n";
+    return EXIT_FAILURE;
+  }
+  const char * const  modelPath = argv[1];
+  const char * const  fixedPath = argv[2];
+  const char * const  movingPath = argv[3];
+  const char * const  outputPath = argv[4];
+  const std::string   device = (argc > 5) ? argv[5] : "cpu";
+
+  constexpr unsigned int Dimension = 3;
+  using PixelType = float;
+  using ImageType = itk::Image<PixelType, Dimension>;
+
+  const auto fixed = itk::ReadImage<ImageType>(fixedPath);
+  const auto moving = itk::ReadImage<ImageType>(movingPath);
+
+  // A TorchScript model configuration: (path, dimension, channels, patchSize, voxelSize,
+  // overlap, layersMask, mixedPrecision). Only POD/STL types cross the API.
+  const itk::ModelConfiguration config(
+    modelPath, Dimension, 1, { 0, 0, 0 }, { 1.0f, 1.0f, 1.0f }, 2, { true }, false);
+
+  // --- 1. core: extract a dense feature map from the fixed image --------------------
+  using InterpolatorType = itk::BSplineInterpolateImageFunction<ImageType, double>;
+  auto interpolator = InterpolatorType::New();
+  interpolator->SetSplineOrder(3);
+
+  auto features = itk::ImageToFeaturesMap<ImageType, InterpolatorType>::New();
+  features->SetModelConfiguration(config);
+  features->SetInterpolator(interpolator);
+  features->SetDevice(device);
+  features->AddInput(fixed);
+  features->Update();
+  const auto featureMap = features->GetOutput(0); // itk::VectorImage<float, 3>
+  std::cout << "feature map: " << featureMap->GetLargestPossibleRegion().GetSize()
+            << "  channels: " << featureMap->GetNumberOfComponentsPerPixel() << std::endl;
+
+  // --- 2. metric: register moving onto fixed by comparing anatomical features -------
+  using MetricType = itk::ImpactImageToImageMetricv4<ImageType, ImageType>;
+  auto                                 metric = MetricType::New();
+  std::vector<itk::ModelConfiguration> models{ config };
+  metric->SetModelsConfiguration(models);
+  metric->SetDistance({ "NCC" });      // per-layer loss: L1, L2, NCC, Cosine, Dice, ...
+  metric->SetLayersWeight({ 1.0f });
+  metric->SetSubsetFeatures({ 4 });    // random channel subset for speed (0 = all)
+  metric->SetPCA({ 0 });
+  metric->SetMode("Static");           // "Static" (precomputed features) or "Jacobian"
+  metric->SetDevice(device);
+
+  using TransformType = itk::TranslationTransform<double, Dimension>;
+  auto transform = TransformType::New();
+  transform->SetIdentity();
+
+  using OptimizerType = itk::RegularStepGradientDescentOptimizerv4<double>;
+  auto optimizer = OptimizerType::New();
+  optimizer->SetNumberOfIterations(200);
+  optimizer->SetLearningRate(2.0);
+  optimizer->SetMinimumStepLength(1e-4);
+
+  using RegistrationType = itk::ImageRegistrationMethodv4<ImageType, ImageType, TransformType>;
+  auto registration = RegistrationType::New();
+  registration->SetFixedImage(fixed);
+  registration->SetMovingImage(moving);
+  registration->SetMetric(metric);
+  registration->SetOptimizer(optimizer);
+  registration->SetInitialTransform(transform);
+
+  // Single resolution level (no shrink, no smoothing).
+  RegistrationType::ShrinkFactorsArrayType   shrinkFactors(1);
+  RegistrationType::SmoothingSigmasArrayType smoothingSigmas(1);
+  shrinkFactors[0] = 1;
+  smoothingSigmas[0] = 0;
+  registration->SetNumberOfLevels(1);
+  registration->SetShrinkFactorsPerLevel(shrinkFactors);
+  registration->SetSmoothingSigmasPerLevel(smoothingSigmas);
+
+  registration->Update();
+  std::cout << "recovered parameters: " << registration->GetTransform()->GetParameters() << std::endl;
+
+  // --- 3. resample the moving image with the recovered transform --------------------
+  auto resample = itk::ResampleImageFilter<ImageType, ImageType>::New();
+  resample->SetInput(moving);
+  resample->SetTransform(registration->GetTransform());
+  resample->SetUseReferenceImage(true);
+  resample->SetReferenceImage(fixed);
+  resample->SetDefaultPixelValue(0);
+  itk::WriteImage(resample->GetOutput(), outputPath);
+
+  return EXIT_SUCCESS;
+}

examples/ImpactMetricExample.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+#==========================================================================
+#
+#   Copyright NumFOCUS
+#
+#   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
+#
+#          https://www.apache.org/licenses/LICENSE-2.0.txt
+#
+#   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.
+#
+#==========================================================================
+
+"""IMPACT: semantic similarity registration from pretrained TorchScript features.
+
+This example shows the two Python-facing pieces of the module -- note that no
+``import torch`` is needed: all LibTorch work happens inside the C++ implementation,
+and the public API speaks only ITK/STL types (images, a model-path string, a device
+string).
+
+  1. the core: ``itk.ImageToFeaturesMap`` extracts a feature map (an ``itk.VectorImage``)
+     from an image with a TorchScript model;
+  2. the metric: ``itk.ImpactImageToImageMetricv4`` plugs into
+     ``itk.ImageRegistrationMethodv4`` to register two images by comparing those features.
+
+Run with:  ./ImpactMetricExample.py model.pt fixed.mha moving.mha
+"""
+
+import argparse
+import itk
+
+parser = argparse.ArgumentParser(description="IMPACT feature extraction + registration demo.")
+parser.add_argument("model", help="TorchScript model (.pt) returning a list of feature maps")
+parser.add_argument("fixed_image")
+parser.add_argument("moving_image")
+parser.add_argument("--device", default="cpu", help='"cpu", "cuda", "cuda:0", ...')
+args = parser.parse_args()
+
+Dimension = 3
+PixelType = itk.F
+ImageType = itk.Image[PixelType, Dimension]
+
+fixed = itk.imread(args.fixed_image, PixelType)
+moving = itk.imread(args.moving_image, PixelType)
+
+# A TorchScript model configuration: (path, dimension, channels, patchSize, voxelSize,
+# overlap, layersMask, mixedPrecision). All POD/STL -- no torch type crosses to Python.
+config = itk.ModelConfiguration(
+    args.model, Dimension, 1, [0, 0, 0], [1.0, 1.0, 1.0], 2, [True], False
+)
+
+# --- 1. core: extract a feature map -----------------------------------------------
+InterpolatorType = itk.BSplineInterpolateImageFunction[ImageType, itk.D]
+interpolator = InterpolatorType.New()
+interpolator.SetSplineOrder(3)
+
+features = itk.ImageToFeaturesMap[ImageType, InterpolatorType].New()
+features.SetModelConfiguration(config)
+features.SetInterpolator(interpolator)
+features.SetDevice(args.device)
+features.AddInput(fixed)
+features.Update()
+feature_map = features.GetOutput(0)  # itk.VectorImage[itk.F, 3]
+print("feature map:", feature_map.GetLargestPossibleRegion().GetSize(),
+      "channels:", feature_map.GetNumberOfComponentsPerPixel())
+
+# --- 2. metric: register moving onto fixed ----------------------------------------
+metric = itk.ImpactImageToImageMetricv4[ImageType, ImageType].New()
+metric.SetFixedImage(fixed)
+metric.SetMovingImage(moving)
+metric.SetFixedModelsConfiguration([config])
+metric.SetMovingModelsConfiguration([config])
+metric.SetDistance(["NCC"])
+metric.SetLayersWeight([1.0])
+metric.SetSubsetFeatures([4])
+metric.SetPCA([0])
+metric.SetMode("Static")
+metric.SetDevice(args.device)
+
+transform = itk.TranslationTransform[itk.D, Dimension].New()
+transform.SetIdentity()
+metric.SetMovingTransform(transform)
+metric.SetFixedTransform(itk.IdentityTransform[itk.D, Dimension].New())
+
+optimizer = itk.RegularStepGradientDescentOptimizerv4[itk.D].New()
+optimizer.SetNumberOfIterations(200)
+optimizer.SetLearningRate(2.0)
+optimizer.SetMinimumStepLength(1e-4)
+
+registration = itk.ImageRegistrationMethodv4[ImageType, ImageType].New()
+registration.SetFixedImage(fixed)
+registration.SetMovingImage(moving)
+registration.SetMetric(metric)
+registration.SetOptimizer(optimizer)
+registration.SetInitialTransform(transform)
+registration.SetNumberOfLevels(1)
+registration.Update()
+
+print("recovered transform parameters:", list(registration.GetTransform().GetParameters()))