ENH: Register a CT to an MR from Python, with the settings that decide the result

ImpactRigidBSplineExample.py drives both stages of ITK's own registration pipeline with the
same itk.ImpactImageToImageMetricv4: Euler3D with a scales estimator, then BSplineTransform
with LBFGSBOptimizerv4, as in Examples/RegistrationITKv4/DeformableRegistration12.cxx. The
composed transform can be written out to carry the alignment to a segmentation or another
sequence.

The rotation centre comes from the fixed image, the domain a moving transform maps from, and
goes through TransformIndexToPhysicalPoint so the direction cosines are applied. Computed as
origin + 0.5 * spacing * size it ignores them: on an image stored LPS rather than RAS it names
a point hundreds of millimetres outside the anatomy, and every degree of rotation then arrives
with that much lever arm.

The distance is L2 across the examples, NCC being an ablation rather than a setting to reach
for.

The mask a caller passes decides the result more than anything else these examples expose. A
whole-body mask is mostly wall and fat; those agree well between MR and CT and outvote the
organs, and the registration settles two centimetres away from the anatomy.
MakeRegistrationImages.py builds a mask around the organs with --mask-model, by segmenting the
fixed image with the model matching its modality, so the figures need no ground truth.

Both figures carry axial and coronal, before against after in each plane. One plane cannot
settle an alignment: a cranio-caudal offset leaves an axial slice looking aligned while every
organ has moved. The coronal panels drop the checkerboard, since this MR has a limited field of
view and the tiles there alternate between anatomy and black bands.

The figures are measured on the pair as delivered. The synthetic offset they used to apply is
off by default: the delivered pair is close to aligned but not aligned, and the registration
improves on it, so there is something to recover without inventing any.

Author: vboussot

examples/ImpactMetricExample.cxx

@@ -87,7 +87,7 @@ main(int argc, char * argv[])
   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->SetDistance({ "L2" });       // 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 });

examples/ImpactMetricExample.py

@@ -60,7 +60,7 @@
     metric.AddModelConfiguration(
         itk.ModelConfiguration(path, Dimension, 1, [0, 0, 0], voxel, 0, [True], False)
     )
-metric.SetDistance(["NCC"] * len(args.models))     # modality-invariant, smooth
+metric.SetDistance(["L2"] * len(args.models))      # the distance to reach for; NCC is an ablation
 metric.SetLayersWeight([1.0] * len(args.models))
 metric.SetSubsetFeatures([12] * len(args.models))  # channels compared per model
 metric.SetPCA([0] * len(args.models))
@@ -70,16 +70,24 @@
     # Air agrees between modalities everywhere, so a whole-image domain measures mostly
     # background: on an abdominal MR-CT pair the similarity varies by 1e-4 over a 2 cm offset,
     # which no optimizer can follow, against 5e-2 for the same offset inside the body.
+    #
+    # How tight the mask is decides the result more than any other setting here. A whole-body
+    # mask is mostly wall and fat, which agree well between the two modalities and outvote the
+    # organs: on the pair below it costs 0.14 of organ Dice against a mask around the organs
+    # themselves, and no distance, model or optimizer setting recovers that. A mask of the
+    # region you want aligned is worth more than a better similarity.
     mask = itk.ImageMaskSpatialObject[Dimension].New()
     mask.SetImage(itk.imread(args.mask, itk.UC))
     mask.Update()
     metric.SetFixedImageMask(mask)
 
-# Rigid transform, centred on the moving image.
-center = [
-    o + 0.5 * s * n
-    for o, s, n in zip(moving.GetOrigin(), moving.GetSpacing(), moving.GetLargestPossibleRegion().GetSize())
-]
+# Rigid transform, centred on the fixed image, which is the domain the moving transform maps
+# from. The centre goes through the direction cosines: computing it as origin + 0.5 * spacing *
+# size ignores them, and on an image stored LPS rather than RAS that lands hundreds of
+# millimetres outside the anatomy. Every degree of rotation then arrives with a lever arm, and
+# the registration settles on a transform that translates the organs instead of turning them.
+size = fixed.GetLargestPossibleRegion().GetSize()
+center = list(fixed.TransformIndexToPhysicalPoint([size[0] // 2, size[1] // 2, size[2] // 2]))
 transform = itk.Euler3DTransform[itk.D].New()
 transform.SetIdentity()
 transform.SetCenter(center)

examples/ImpactRigidBSplineExample.py

@@ -0,0 +1,229 @@
+#!/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 with the standard ITK registration pipeline: rigid, then B-spline.
+
+The same `itk.ImpactImageToImageMetricv4` drives both stages; only the transform and the
+optimizer change. The layout follows ITK's own examples: Euler3D with a scales estimator for the
+rigid stage, `BSplineTransform` with `LBFGSBOptimizerv4` for the deformable one, as in
+`Examples/RegistrationITKv4/DeformableRegistration12.cxx`.
+
+Three settings are worth knowing about, because each is easy to get wrong and each decides
+whether the result is an alignment or a wander.
+
+**Mask the region you want aligned**, not the whole body. Wall and fat agree well between
+modalities and outvote the organs, so a body mask sends the registration somewhere else
+entirely. No labels are needed: segment the fixed image with the model matching its modality and
+dilate what it does not call background.
+
+**The B-spline stage uses no scales estimator.** A B-spline coefficient is already a displacement
+in millimetres, so every coefficient has the same physical meaning and there is nothing to
+equalise. Asked anyway, the estimator returns about 6.7e-07 for each of them against 1 for a
+translation parameter, and a gradient descent divides its update by that. LBFGSB is quasi-Newton
+and builds its own conditioning, which is why ITK pairs it with the B-spline transform.
+
+**One resolution level.** ITK's shrink factors never reach the model: `ImageToFeaturesMap`
+resamples whatever it is handed to the `ModelConfiguration` voxel size, so a shrunk image is
+interpolated straight back up before the network sees it, leaving the smoothing and the lost
+detail. Coarsening the model's voxel size is the real knob, and for a local descriptor like MIND
+it describes different anatomy rather than the same anatomy more coarsely.
+
+A finer B-spline mesh is not reached by chaining two registrations: ITK cannot compose a B-spline
+as a moving initial transform (`ComputeJacobianWithRespectToPosition` is not implemented for it),
+which is what `itk::BSplineTransformParametersAdaptor` exists for. That class is not currently
+wrapped for Python, so a coarse-to-fine schedule needs C++.
+
+Run:  ./ImpactRigidBSplineExample.py fixed.nii.gz moving.nii.gz mind.pt \\
+          --mask organs.nii.gz --grid-spacing 40 --out moved.nii.gz --device cuda:0
+"""
+
+import argparse
+
+import itk
+
+parser = argparse.ArgumentParser(description="IMPACT rigid then B-spline registration.")
+parser.add_argument("fixed_image")
+parser.add_argument("moving_image")
+parser.add_argument("models", nargs="+", help="one or more TorchScript feature models (.pt)")
+parser.add_argument("--mask", default=None,
+                    help="fixed-image mask, around the organs rather than the whole body; "
+                         "see the docstring, it is the setting that matters most")
+parser.add_argument("--out", default="moved.nii.gz", help="warped moving image")
+parser.add_argument("--voxel", type=float, default=2.0, help="feature voxel size in mm")
+parser.add_argument("--grid-spacing", type=float, default=20.0,
+                    help="distance between B-spline control points, in mm. This is the knob that\n"
+                         "decides what the deformation can express: too coarse and one control\n"
+                         "point drags a whole region, which shows up as one structure improving\n"
+                         "while its neighbour is pulled apart. elastix targets 10 mm for this\n"
+                         "kind of abdominal case (FinalGridSpacingInPhysicalUnits)")
+parser.add_argument("--sampling", type=float, default=0.10, help="fraction of voxels sampled")
+parser.add_argument("--shrink-factors", type=int, nargs="+", default=[1],
+                    help="one shrink factor per resolution level, coarsest first")
+parser.add_argument("--smoothing-sigmas", type=float, nargs="+", default=[0.0],
+                    help="Gaussian sigma per level, in millimetres, matching --shrink-factors")
+parser.add_argument("--rigid-iterations", type=int, default=300)
+parser.add_argument("--bspline-iterations", type=int, default=200,
+                    help="also caps the LBFGSB function evaluations; the cost of the run is "
+                         "roughly this times the sampled fraction of the volume")
+parser.add_argument("--save-transform", default=None,
+                    help="write the composed rigid+B-spline transform (.tfm/.h5), so the same\n"
+                         "alignment can be applied to a segmentation or another sequence")
+parser.add_argument("--device", default="cuda:0", help='"cpu", "cuda", "cuda:0", ...')
+args = parser.parse_args()
+
+if len(args.shrink_factors) != len(args.smoothing_sigmas):
+    parser.error("--shrink-factors and --smoothing-sigmas carry one entry per level, "
+                 f"got {len(args.shrink_factors)} and {len(args.smoothing_sigmas)}")
+
+Dimension = 3
+SplineOrder = 3
+ImageType = itk.Image[itk.F, Dimension]
+
+fixed = itk.imread(args.fixed_image, itk.F)
+moving = itk.imread(args.moving_image, itk.F)
+
+# --- the metric, shared by both stages ---------------------------------------------------
+metric = itk.ImpactImageToImageMetricv4[ImageType, ImageType].New()
+for path in args.models:
+    metric.AddModelConfiguration(
+        itk.ModelConfiguration(path, Dimension, 1, [0, 0, 0], [args.voxel] * Dimension, 0, [True], False)
+    )
+metric.SetDistance(["L2"] * len(args.models))
+metric.SetLayersWeight([1.0] * len(args.models))
+metric.SetSubsetFeatures([12] * len(args.models))
+metric.SetPCA([0] * len(args.models))
+metric.SetMode("Static")
+metric.SetDevice(args.device)
+
+if args.mask:
+    # Air agrees between modalities everywhere, so a whole-image domain measures mostly
+    # background and the similarity barely varies with the transform.
+    mask = itk.ImageMaskSpatialObject[Dimension].New()
+    mask.SetImage(itk.imread(args.mask, itk.UC))
+    mask.Update()
+    metric.SetFixedImageMask(mask)
+
+identity = itk.IdentityTransform[itk.D, Dimension].New()
+
+
+def pyramid(registration, shrink_factors, smoothing_sigmas):
+    """Resolution levels and sampling: see the module docstring."""
+    registration.SetNumberOfLevels(len(shrink_factors))
+    registration.SetShrinkFactorsPerLevel(shrink_factors)
+    registration.SetSmoothingSigmasPerLevel(smoothing_sigmas)
+    registration.SetMetricSamplingStrategy(2)  # 0 NONE, 1 REGULAR, 2 RANDOM
+    registration.SetMetricSamplingPercentage(args.sampling)
+
+
+# --- 1. rigid ------------------------------------------------------------------------------
+size = fixed.GetLargestPossibleRegion().GetSize()
+centre = list(fixed.TransformIndexToPhysicalPoint([size[0] // 2, size[1] // 2, size[2] // 2]))
+
+rigid = itk.Euler3DTransform[itk.D].New()
+rigid.SetIdentity()
+rigid.SetCenter(centre)
+
+rigid_optimizer = itk.RegularStepGradientDescentOptimizerv4[itk.D].New()
+rigid_optimizer.SetLearningRate(2.0)
+rigid_optimizer.SetMinimumStepLength(1e-4)
+rigid_optimizer.SetRelaxationFactor(0.8)
+rigid_optimizer.SetNumberOfIterations(args.rigid_iterations)
+# A radian and a millimetre are not comparable; let ITK derive the ratio from the physical
+# shift each parameter causes, measured against this metric's own gradient.
+rigid_scales = itk.RegistrationParameterScalesFromPhysicalShift[type(metric)].New()
+rigid_scales.SetMetric(metric)
+rigid_optimizer.SetScalesEstimator(rigid_scales)
+
+rigid_registration = itk.ImageRegistrationMethodv4[ImageType, ImageType].New()
+rigid_registration.SetFixedImage(fixed)
+rigid_registration.SetMovingImage(moving)
+rigid_registration.SetMetric(metric)
+rigid_registration.SetOptimizer(rigid_optimizer)
+rigid_registration.SetInitialTransform(rigid)
+pyramid(rigid_registration, args.shrink_factors, args.smoothing_sigmas)
+rigid_registration.Update()
+print("rigid :", rigid_optimizer.GetStopConditionDescription())
+
+# --- 2. B-spline, on top of the rigid ------------------------------------------------------
+# The transform domain covers the fixed image, as ITK's example sets it out.
+physical_dimensions = [
+    fixed.GetSpacing()[i] * (fixed.GetLargestPossibleRegion().GetSize()[i] - 1) for i in range(Dimension)
+]
+bspline = itk.BSplineTransform[itk.D, Dimension, SplineOrder].New()
+bspline.SetTransformDomainOrigin(fixed.GetOrigin())
+bspline.SetTransformDomainPhysicalDimensions(physical_dimensions)
+bspline.SetTransformDomainDirection(fixed.GetDirection())
+# Mesh size from the requested physical spacing, so the grid does not silently change
+# meaning when the image geometry does.
+mesh_size = [max(1, int(round(physical_dimensions[i] / args.grid_spacing))) for i in range(Dimension)]
+bspline.SetTransformDomainMeshSize(mesh_size)
+print(f"bspline grid: mesh {mesh_size}, "
+      f"{[round(physical_dimensions[i] / mesh_size[i]) for i in range(Dimension)]} mm apart, "
+      f"{bspline.GetNumberOfParameters()} parameters")
+
+number_of_parameters = bspline.GetNumberOfParameters()
+bspline_optimizer = itk.LBFGSBOptimizerv4.New()
+# All bounds deselected, so the coefficients are free: LBFGSB is used here for its
+# quasi-Newton conditioning, not for its bounds.
+bspline_optimizer.SetBoundSelection(itk.Array[itk.SL]([0] * number_of_parameters))
+bspline_optimizer.SetLowerBound(itk.Array[itk.D]([0.0] * number_of_parameters))
+bspline_optimizer.SetUpperBound(itk.Array[itk.D]([0.0] * number_of_parameters))
+bspline_optimizer.SetCostFunctionConvergenceFactor(1e7)
+bspline_optimizer.SetGradientConvergenceTolerance(1e-35)
+bspline_optimizer.SetNumberOfIterations(args.bspline_iterations)
+bspline_optimizer.SetMaximumNumberOfFunctionEvaluations(args.bspline_iterations)
+bspline_optimizer.SetMaximumNumberOfCorrections(7)
+# Deliberately no scales estimator: see the module docstring.
+
+bspline_registration = itk.ImageRegistrationMethodv4[ImageType, ImageType].New()
+bspline_registration.SetFixedImage(fixed)
+bspline_registration.SetMovingImage(moving)
+bspline_registration.SetMetric(metric)
+bspline_registration.SetOptimizer(bspline_optimizer)
+bspline_registration.SetInitialTransform(bspline)
+# The rigid result is held in front of the B-spline rather than folded into it, which is how
+# ITKv4 composes a prior stage. A rigid transform can sit here; a B-spline cannot.
+bspline_registration.SetMovingInitialTransform(rigid_registration.GetTransform())
+# One level, unlike the rigid stage. A pyramid only helps a deformable stage if the control
+# grid is coarsened along with the image; held at its final spacing while the image shrinks by
+# four, the grid is fine relative to the blurred content and fits it, which the finer levels
+# then inherit. Coarsening it is what BSplineTransformParametersAdaptor does, and that class is
+# not wrapped for Python.
+pyramid(bspline_registration, [1], [0.0])
+bspline_registration.Update()
+print("bspline :", bspline_optimizer.GetStopConditionDescription())
+
+# --- 3. resample through both ---------------------------------------------------------------
+# CompositeTransform applies the last added first, so this is rigid(bspline(point)) -- the
+# same composition ITKv4 made internally during the second stage.
+composite = itk.CompositeTransform[itk.D, Dimension].New()
+composite.AddTransform(rigid_registration.GetTransform())
+composite.AddTransform(bspline)
+
+itk.imwrite(
+    itk.resample_image_filter(
+        moving, transform=composite, use_reference_image=True, reference_image=fixed,
+        default_pixel_value=-1024.0,
+    ),
+    args.out,
+)
+if args.save_transform:
+    itk.transformwrite([composite], args.save_transform)
+    print("wrote", args.save_transform)
+print("wrote", args.out)