def set_metadata(image: sitk.Image, reference: sitk.Image):
assert image.GetSize() == reference.GetSize()
image.CopyInformation(reference)
def displacement(jacobian: sitk.Image,
*,
levels: int = 1,
pad: int = 0,
redistribute: bool = False,
mask: sitk.Image = None,
initial_guess: sitk.Image = None,
epsilon: float = 9.99e-4,
tolerance: float = 0.2,
it_max: Union[int, List[int]] = 50000,
alpha: float = 1.2,
beta: float = 0.5,
gamma: float = 0.1,
delta: float = 1e-3,
zeta: float = 10.0,
theta: float = 1e-6,
iota: float = 1e-9,
strict: bool = False,
eta: float = 0.1,
eta_max: float = 0.4,
algorithm: str = 'gradient',
gpu_id: int = -1) -> sitk.Image:
r""" Generate a displacement field that realises a given Jacobian.
Given a 3D scalar image encoding a Jacobian map, compute a
3D vector image encoding a vector field whose Jacobian map
matches the input up to a certain tolerance.
The three algorithms provided are:
* ``gradient``: a gradient descent method (default).
* ``greedy``: a greedy search method based on the method proposed in [1]_.
* ``matching``: a volume matching routine based on gradient descent,
published in [2]_ and [3]_. The implementation comes from
the `atrophysim tool`_.
The initial value of the step length in the gradient descent is
given by the parameter ``eta``. The ``gradient`` and ``matching``
algorithms use an `Armijo condition`_ to control the step length,
in the form
.. math::
E(d - \eta \nabla E(d)) - E(d) \le -\gamma \eta ||\nabla E(d)||^2
where :math:`d` is the displacement, :math:`E` the loss function,
:math:`\eta` the current step length, and :math:`\gamma \in (0, 1)` a
parameter of the condition. At each iteration the step length is increased
multiplying it by ``alpha``, and if the Armijo condition is not met after
the update, ``eta`` is decreased multiplying it by ``beta`` until the
truth of the inequality in the Armijo condition is restored. A maximum
value for ``eta`` can be set through the parameter ``eta_max``.
The ``gradient`` and ``matching`` algorithms have a regularisation term
that penalises values of the Jacobian below a certain threshold, given
by ``delta``. The importance of the regularisation term is controlled
by the parameter ``zeta`` (set to ``0`` to have no regularisation).
Termination is controlled by a condition on the improvement on the result
and one on the step length. If the percentual improvement of the cost
drops below the value given by ``theta``, the algorithm terminates.
Termination happens also if the step length becomes smaller than
``iota``.
.. _atrophysim tool: https://www.nitrc.org/projects/atrophysim
.. _Armijo condition: https://en.wikipedia.org/wiki/Wolfe_conditions
.. note::
The displacement is generally not accurate on image boundary voxels.
.. note::
The C verbose output is written to `stdout`. If you want to capture
it from within Python, the `wurlitzer package`_ might be helpful.
.. warning::
This function calls a C routine which cannot be interrupted from
the REPL.
.. _wurlitzer package: https://github.com/minrk/wurlitzer
References
----------
.. [1] van Eede, M. C., Scholz, J., Chakravarty, M. M., Henkelman, R. M., and Lerch, J. P.
"Mapping registration sensitivity in MR mouse brain images." Neuroimage 82 (2013), 226–236.
.. [2] Karaçali, B., and Davatzikos, C. "Estimating topology preserving and smooth displacement fields."
IEEE Transactions on Medical Imaging 23, 7 (2004), 868–880.
.. [3] Karaçali, B., and Davatzikos, C. "Simulation of tissue atrophy using a topology preserving
transformation model." IEEE transactions on medical imaging 25, 5 (2006), 649–652.
Parameters
----------
jacobian : sitk.Image
Input Jacobian.
levels : int
Number of levels in the multi-resolution pyramid; the size of
the image along each direction is halven at each level.
pad : int
Thickness of the zero-padding around the volume (0 for
the mask, 1.0 for the Jacobian) to be used during the
computation. The padding is removed before returning the result.
redistribute : bool
Redistribute the volume change inside the mask to the background.
mask : sitk.Image
Binary mask for the region of interest.
initial_guess : sitk.Image
Initial estimation of the solution. The default is a null
displacement field.
epsilon : float
A floating point value, representing the tolerance per
voxel on the Jacobian of the resulting field.
tolerance : float
Tolerance on Jacobian outside the mask.
it_max : Union[int, List[int]]
Maximum number of iterations allowed. If it is a list, its
length must match the number of levels in the multi-resolution
pyramid, and each value is used for a single level, with the
first element of the list representing the level with lowest
resolution. If it is a scalar, then the same number of
iterations is used for all pyramid levels.
alpha : float
Coefficient that controls the increase of the step length.
beta : float
Coefficient that controls the decrease of the step length.
gamma : float
Armijo-Goldstein parameter.
delta : float
Lower threshold for Jacobian regularisation.
zeta : float
Weight for the regularisation term.
theta : float
Terminate if the percentage improvement of the cost per
iteration drops below this value.
iota : float
Terminate if the step length drops below this value.
strict : bool
If True, reject iterations that not decrease the maximum
voxel error.
eta : float
Initial step length.
eta_max : float
Maximum step length allowed.
algorithm : str
Algorithm to generate the field, one of `greedy`, `gradient`, or `matching`.
gpu_id : int
Id of the CUDA device used to run the GPU implementation. If
equal to `-1`, the CPU implementation is used instead. Requires
a build of disptools with CUDA support enabled.
Returns
-------
sitk.Image
A displacement field whose Jacobian matches the input.
"""
# Filter parameters used locally and parameters propagated to the
# wrapped function
parameters = locals().copy()
used = [
'jacobian', 'levels', 'pad', 'redistribute', 'mask', 'initial_guess',
'it_max'
]
[parameters.pop(p) for p in used]
jacobian = sitk.Cast(jacobian, sitk_float_type)
if mask is None:
mask = np.ones(tuple(reversed(jacobian.GetSize())),
dtype=np_float_type)
mask = sitk.GetImageFromArray(mask)
mask.CopyInformation(jacobian)
else:
mask = sitk.Cast(mask, sitk_float_type)
size = jacobian.GetSize()
origin = jacobian.GetOrigin()
spacing = jacobian.GetSpacing()
direction = jacobian.GetDirection()
# Ensure consistency of the coordinate system
def make_consistent(img, name, interpolator):
if img is None:
return img
inconsitent = []
if img.GetSize() != size:
inconsitent.append('size')
if img.GetOrigin() != origin:
inconsitent.append('origin')
if img.GetSpacing() != spacing:
inconsitent.append('spacing')
if img.GetDirection() != direction:
inconsitent.append('direction')
if inconsitent != []:
inconsitent = ' and '.join(inconsitent)
warnings.warn("%s of '%s' " % (inconsitent, name) +
"inconsistent with the Jacobian, " +
"resampling to a common coordinate space")
if interpolator != sitk.sitkNearestNeighbor:
img = sitk.SmoothingRecursiveGaussian(img, 2.0)
return sitk.Resample(img, jacobian, sitk.Transform(), interpolator)
else:
return img
mask = make_consistent(mask, 'mask', sitk.sitkNearestNeighbor)
initial_guess = make_consistent(initial_guess, 'initial_guess',
sitk.sitkLinear)
# Add a voxel of zero-flux padding anyway since the algorithm
# will not compute the displacement field on boundary voxels
if pad > 0:
pad += 1
else:
pad = 1
pad = ((pad, pad, pad), (pad, pad, pad))
if redistribute:
jacobian = redistribute_volume_change(jacobian, mask)
mask = sitk.ConstantPad(mask, *pad, 0)
jacobian = sitk.ZeroFluxNeumannPad(jacobian, *pad)
# Create image pyramid
jacobian_pyramid = [jacobian]
mask_pyramid = [mask]
for i in range(1, levels):
new_size = tuple(
map(lambda x: x // 2, jacobian_pyramid[i - 1].GetSize()))
jacobian_pyramid.append(drawing.scale_image(jacobian, new_size))
mask_pyramid.append(
drawing.scale_image(mask, new_size, sitk.sitkNearestNeighbor))
# Set maximum number of iterations for each pyramid level
if not isinstance(it_max, list):
it_max = [it_max for i in range(levels)]
elif len(it_max) != levels:
raise ValueError('len(it_max) should equal the value of `levels`')
# Set initial guess
field = initial_guess
# Multi-resolution algorithm
for jacobian, mask, it in zip(jacobian_pyramid[::-1], mask_pyramid[::-1],
it_max):
size = jacobian.GetSize()
logging.info('Size %s' % str(size))
field = drawing.scale_image(field, size) if field is not None else None
field = _displacement(jacobian,
mask,
initial_guess=field,
it_max=it,
**parameters)
# Remove padding from the result
field = sitk.Crop(field, *pad)
field.SetOrigin(origin)
field.SetSpacing(spacing)
field.SetDirection(direction)
return field
