Compute the Deformation Gradient from an SimpleITK Displacement Field Image (shape (x, y, z, 3)) in numpy and SITK

deformation_gradient.py

import SimpleITK as sitk
import numpy as np

def construct_gradU_numpy(img: sitk.Image) -> np.array:
    """Constructs the deformation gradient using numpy (potentially slow)"""
    u = sitk.GetArrayViewFromImage(img)
    dx, dy, dz = img.GetSpacing()
    # Have to calculate the gradient row vise, as np.gradient can only use scalar values.

    # NOTE: The components of the vector are in x,y,z. However, the image itself is in z,y,x!
    # Thus, the ordering of the gradient is different.
    # --> For the tensor, we sort the vectors back into correct ordering, however, the image stays z,y,x.
    dx_dZ, dx_dY, dx_dX = np.gradient(u[..., 0], dz, dy, dx)
    dy_dZ, dy_dY, dy_dX = np.gradient(u[..., 1], dz, dy, dx)
    dz_dZ, dz_dY, dz_dX = np.gradient(u[..., 2], dz, dy, dx)

    gradU = np.stack([
        np.stack([dx_dX, dx_dY, dx_dZ], axis=-1),
        np.stack([dy_dX, dy_dY, dy_dZ], axis=-1),
        np.stack([dz_dX, dz_dY, dz_dZ], axis=-1),
    ], axis=3)
    return gradU


def construct_gradU_sitk(img: sitk.Image) -> np.array:
    """Construct the Deformation Gradient using SimpleITK"""
    return np.stack([
        sitk.GetArrayFromImage(sitk.Gradient(sitk.VectorIndexSelectionCast(img, i)))
        for i in range(3)
    ], axis=3)