qsyao/cuda_spatial_deform

minimal working example for elastic_deform

vbeliveau opened this issue · 7 comments

This package is exactly what I need to speed up the data augmentation in my code! Unfortunately, I am unable to get the elastic deformation to work. Here is a tentative minimal working example for the FLAIR image included with the code..

import nibabel as nib
import time
import numpy as np
from cuda_backend.py_api import Handle

img = nib.load('data/FLAIR.nii.gz')
data = img.get_fdata().astype('float32')

start = time.time()
cuda_handle = Handle(data.shape, RGB=False, mode='constant')
cuda_handle.elastic(sigma=1., alpha=20., mode='constant')
cuda_handle.end_flag()
data_aug = cuda_handle.augment(data)
print('%f' % (time.time() - start))

nib.save(nib.Nifti1Image(np.array(data_aug[0]), img.affine, img.header),
         'Flair.deformed.nii.gz')

And here are the input/output:

example1
example2

Any idea what's is going on?

qsyao commented

Can u compare the result with cpu-backend(like simpleITK, https://www.kaggle.com/bguberfain/elastic-transform-for-data-augmentation, etc). The sigma of the elastic deform is between 12-15 in most case.

Sure... Here's an example comparing your code to a CPU equivalent deformation.

import nibabel as nib
import time
import numpy as np
from cuda_backend.py_api import Handle
from scipy.ndimage.interpolation import map_coordinates
from scipy.ndimage.filters import gaussian_filter

# Function to distort image
def elastic_transform(image, alpha, sigma, random_state=None, order=3, mode='constant'):
    """Elastic deformation of images as described in [Simard2003]_ (with modifications).
    .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
         Convolutional Neural Networks applied to Visual Document Analysis", in
         Proc. of the International Conference on Document Analysis and
         Recognition, 2003.

     Based on https://gist.github.com/erniejunior/601cdf56d2b424757de5
    """
    if random_state is None:
        random_state = np.random.RandomState(None)

    shape = image.shape
    dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha
    dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha
    dz = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha

    x, y, z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]), np.arange(shape[2]))
    indices = np.reshape(y+dy, (-1, 1)), np.reshape(x+dx, (-1, 1)), np.reshape(z+dz, (-1, 1))

    return map_coordinates(image, indices, order=order, mode=mode).reshape(shape)

# Load data and set deform parameters
img = nib.load('data/FLAIR.nii.gz')
data = img.get_fdata().astype('float32')
sigma = 100
alpha = 5

# CUDA deform
start = time.time()
cuda_handle = Handle(data.shape, RGB=False, mode='constant')
# cuda_handle = Cuda_Spatial_Deform(data.shape, mode="constant")
cuda_handle.elastic(sigma=sigma, alpha=alpha, mode='constant')
cuda_handle.end_flag()
data_aug = cuda_handle.augment(data, order=1)
print('CUDA deform: %f' % (time.time() - start))
nib.save(nib.Nifti1Image(np.array(data_aug[0]), img.affine, img.header),
         'Flair.cuda_deform.nii.gz')

# CPU deform
start = time.time()
data_aug2 = elastic_transform(data, sigma, alpha)
print('CPU deform: %f' % (time.time() - start))
nib.save(nib.Nifti1Image(data_aug2, img.affine, img.header),
         'Flair.cpu_deform.nii.gz')

The corresponding CUDA deformation:

flair_cuda_deform

and the CPU deformation:

flair_cpu_deform

qsyao commented

Thank you very much! I am so stupid to make such a big bug hear.
I have fixed it now. Here is my results.
By the way, it seems that the order sigma and alpha is wrong in your code :).
捕获

Thanks for the fix. Unfortunately, I think there is something else fishy going on.. Your code works when using SimpleITK, but not when using another library to open the nifti file. See this for example:

import nibabel as nib
import SimpleITK as sitk
import numpy as np
from cuda_backend.py_api import Handle

sigma = 5.
alpha = 200.

# CUDA deform SimpleITK
data_pth = 'data/FLAIR.nii.gz'
sitk_image = sitk.ReadImage(data_pth)
array_image = sitk.GetArrayFromImage(sitk_image).copy()
print('SimpleITK image shape:')
print(array_image.shape)
cuda_handle = Handle(array_image.shape, mode='mirror')
cuda_handle.elastic(sigma=sigma, alpha=alpha, mode='constant')
cuda_handle.end_flag()
output = cuda_handle.augment(array_image, order=1)
volOut=sitk.GetImageFromArray(output[0])
sitk.WriteImage( volOut,"Flair.cuda_sitk.nii.gz", True)

# CUDA deform nibabel
img = nib.load('data/FLAIR.nii.gz')
data = img.get_fdata().astype('float32')
print('Nibabel image shape:')
print(data.shape)
cuda_handle = Handle(data.shape, mode='mirror')
cuda_handle.elastic(sigma=sigma, alpha=alpha, mode='constant')
cuda_handle.end_flag()
data_aug = cuda_handle.augment(data, order=1)
nib.save(nib.Nifti1Image(data_aug[0], img.affine, img.header),
         'Flair.cuda_nibabel.nii.gz')

The only difference between sitk and nibabel is that sitk is "channel first". For 2D images, sitk will output [c x y], but for 3D images with no channel it will output [z y x]. In theory, the ordering of the data shouldn't change anything if its also saved using the same standard, but somehow something scrambles the results in your code.

qsyao commented

Ops if you want to process 2D image with [c,y.x], you should do the augment c times ....... The coodinates are determinated once you define your flow.

I think you misunderstand. In [c,y,x], c refers to the channels, i.e. for RGB that would be the 3 color channel. My point with the example above was that your code is sensitive to the order of the dimensions. It works fine when the data is loaded as [z,y,x], as done by sitk, but something is scrambled when the data is loaded as [x,y,z]. This is an unexpected behavior and indicate that something weird is going on in your code.

qsyao commented

Thank you!
The numpy array python frontend deliverd to the cuda backend must be continuous and have the correct order of dimensions in the real memory.
So if you want to deliver the array with shape [z,y,x] when it is loaded as [x,y,z] by transpose, you should use array.transpose().copy().
I have not used nibabel before , so I do not know if the array loaded by nibable is continunous in memory.