test_bspline_deform.py

import SimpleITK as sitk
import matplotlib.pyplot as plt
from neverfault.resample import resample_new
import numpy as np

def test(im):
    """ Use Simple ITK to random deform (elastic deform) an 2D image. Can be easily 
    used for 3D. There are two ways of elastic deforming an image. One is to define 
    a sparse control points, and use B-Spline to interpolate the deformation field. 
    Another is to directly define the deformation field. For data augmentation purpose,
    both methods work. This function use the first method.  """

    # Create the transformation (when working with images it is easier to use
    # the BSplineTransformInitializer function or its object oriented counterpart
    # BSplineTransformInitializerFilter).
    
    itk_im = sitk.GetImageFromArray(im)
    
    # (20, 20) define the number of control points. Large number will have more control points. 
    bspline = sitk.BSplineTransformInitializer(itk_im, (20,20), 3)

    # Random displacement of the control points.
    originalControlPointDisplacements = np.random.random(len(bspline.GetParameters())) * 10
    bspline.SetParameters(originalControlPointDisplacements)
    
    # use the original input image as the reference image for resampling. 
    itk_im_resampled = sitk.Resample(itk_im, itk_im,
                                    bspline, sitk.sitkLinear, 0)

    im_resampled = sitk.GetArrayFromImage(itk_im_resampled)

    fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (15, 10), sharex = True, sharey = True)

    ax[0].imshow(im, interpolation = 'nearest', cmap = 'bwr')
    ax[1].imshow(im_resampled, interpolation = 'nearest', cmap = 'bwr')

    plt.show()