ImageRegistrationMethodExhaustive/ImageRegistrationMethodExhaustive.py

#!/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
#
#         http://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.
#
# =========================================================================
 
"""
This script demonstrates the use of the Exhaustive optimizer in the
ImageRegistrationMethod to estimate a good initial rotation position.
 
Because gradient descent base optimization can get stuck in local
minima, a good initial transform is critical for reasonable
results. Search a reasonable space on a grid with brute force may be a
reliable way to get a starting location for further optimization.
 
The initial translation and center of rotation for the transform is
initialized based on the first principle moments of the intensities of
the image. Then in either 2D or 3D a Euler transform is used to
exhaustively search a grid of the rotation space at a certain step
size. The resulting transform is a reasonable guess where to start
further registration.
"""
 
import SimpleITK as sitk
import sys
import os
from math import pi
 
 
def command_iteration(method):
    if method.GetOptimizerIteration() == 0:
        print("Scales: ", method.GetOptimizerScales())
    print(
        f"{method.GetOptimizerIteration():3} "
        + f"= {method.GetMetricValue():7.5f} "
        + f": {method.GetOptimizerPosition()}"
    )
 
 
if len(sys.argv) < 4:
    print(
        "Usage:",
        sys.argv[0],
        "<fixedImageFilter> <movingImageFile>",
        "<outputTransformFile>",
    )
    sys.exit(1)
 
fixed = sitk.ReadImage(sys.argv[1], sitk.sitkFloat32)
 
moving = sitk.ReadImage(sys.argv[2], sitk.sitkFloat32)
 
R = sitk.ImageRegistrationMethod()
 
R.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)
 
sample_per_axis = 12
if fixed.GetDimension() == 2:
    tx = sitk.Euler2DTransform()
    # Set the number of samples (radius) in each dimension, with a
    # default step size of 1.0
    R.SetOptimizerAsExhaustive([sample_per_axis // 2, 0, 0])
    # Utilize the scale to set the step size for each dimension
    R.SetOptimizerScales([2.0 * pi / sample_per_axis, 1.0, 1.0])
elif fixed.GetDimension() == 3:
    tx = sitk.Euler3DTransform()
    R.SetOptimizerAsExhaustive(
        [
            sample_per_axis // 2,
            sample_per_axis // 2,
            sample_per_axis // 4,
            0,
            0,
            0,
        ]
    )
    R.SetOptimizerScales(
        [
            2.0 * pi / sample_per_axis,
            2.0 * pi / sample_per_axis,
            2.0 * pi / sample_per_axis,
            1.0,
            1.0,
            1.0,
        ]
    )
 
# Initialize the transform with a translation and the center of
# rotation from the moments of intensity.
tx = sitk.CenteredTransformInitializer(fixed, moving, tx)
 
R.SetInitialTransform(tx)
 
R.SetInterpolator(sitk.sitkLinear)
 
R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R))
 
outTx = R.Execute(fixed, moving)
 
print("-------")
print(outTx)
print(f"Optimizer stop condition: {R.GetOptimizerStopConditionDescription()}")
print(f" Iteration: {R.GetOptimizerIteration()}")
print(f" Metric value: {R.GetMetricValue()}")
 
sitk.WriteTransform(outTx, sys.argv[3])
 
if "SITK_NOSHOW" not in os.environ:
    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(fixed)
    resampler.SetInterpolator(sitk.sitkLinear)
    resampler.SetDefaultPixelValue(1)
    resampler.SetTransform(outTx)
 
    out = resampler.Execute(moving)
 
    simg1 = sitk.Cast(sitk.RescaleIntensity(fixed), sitk.sitkUInt8)
    simg2 = sitk.Cast(sitk.RescaleIntensity(out), sitk.sitkUInt8)
    cimg = sitk.Compose(simg1, simg2, simg1 // 2.0 + simg2 // 2.0)
    sitk.Show(cimg, "ImageRegistrationExhaustive Composition")