Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save thewtex/9503448d2ad1dfacc2cbd620d95d3dac to your computer and use it in GitHub Desktop.
Save thewtex/9503448d2ad1dfacc2cbd620d95d3dac to your computer and use it in GitHub Desktop.
Save an image from a numpy array and an affine matrix
import itk
import numpy as np
def save(np_array: np.ndarray, filename: str, affine_matrix: np.ndarray):
# -1 for translation
Dimension = affine_matrix.shape[0] - 1
# Is this a multi-component image?
is_vector = np_array.ndim != Dimension
print('is_vector', is_vector)
itk_np_view = itk.image_view_from_array(np_array, is_vector=is_vector)
transform = itk.AffineTransform[itk.D, Dimension].New()
params = transform.GetParameters()
affine = affine_matrix[:Dimension, :Dimension]
for idx, val in enumerate(affine.flat):
params[idx] = val
for idx in range(Dimension):
params[Dimension*Dimension + idx] = affine_matrix[idx, Dimension]
transform.SetParameters(params)
# Depending on the direction of the transformation, we may need to use the
# inverse, instead.
# inverse = transform.GetInverseTransform()
origin = transform.TransformPoint([0.0,]*Dimension)
itk_np_view.SetOrigin(origin)
identity = np.eye(Dimension)
spacing = []
direction = np.eye(Dimension)
for dim in range(Dimension):
vector = identity[dim,:]
transformed = transform.TransformVector(vector)
spacing.append(transformed.Normalize())
direction[dim,:] = transformed
itk_np_view.SetSpacing(spacing)
itk_np_view.SetDirection(direction)
itk.imwrite(itk_np_view, filename)
def affine_transform_to_matrix(transform):
Dimension = itk.template(transform)[1][1]
params = np.asarray(transform.GetParameters())
affine = params.reshape((Dimension+1, Dimension))
affine_matrix = np.eye(Dimension+1)
affine_matrix[:Dimension, :Dimension] = affine[:Dimension, :Dimension]
affine_matrix[:Dimension, Dimension] = affine[Dimension, :]
print('affine_matrix:')
print(affine_matrix)
return affine_matrix
print('3D, scalar image with offset')
image = itk.imread('/home/matt/data/head.mha', itk.F)
np_array = np.asarray(image)
ParamType = itk.ctype('double')
Dimension = 3
transform = itk.AffineTransform[ParamType, Dimension].New()
transform.SetCenter(image.GetOrigin())
transform.Scale(image.GetSpacing())
transform.SetOffset(image.GetOrigin())
affine_matrix = affine_transform_to_matrix(transform)
save(np_array, '/tmp/head_apply_transform.nii.gz', affine_matrix)
print('\n3D, scalar image with origin')
image = itk.imread('/home/matt/data/head_origin.mha', itk.F)
np_array = np.asarray(image)
ParamType = itk.ctype('double')
Dimension = 3
transform = itk.AffineTransform[ParamType, Dimension].New()
transform.Scale(image.GetSpacing())
transform.SetOffset(image.GetOrigin())
affine_matrix = affine_transform_to_matrix(transform)
save(np_array, '/tmp/head_origin_apply_transform.nii.gz', affine_matrix)
print('\n3D, scalar image with direction')
image = itk.imread('/home/matt/data/head_direction.mha', itk.F)
np_array = np.asarray(image)
ParamType = itk.ctype('double')
Dimension = 3
transform = itk.AffineTransform[ParamType, Dimension].New()
transform.SetMatrix(image.GetDirection())
transform.Scale(image.GetSpacing(), True)
transform.SetOffset(image.GetOrigin())
affine_matrix = affine_transform_to_matrix(transform)
save(np_array, '/tmp/head_direction_apply_transform.nii.gz', affine_matrix)
print('\n3D, multi-channel image')
image = itk.imread('/home/matt/data/head_direction.mha', itk.F)
np_array = np.asarray(image)
np_array_multi = np.stack([np_array, (np_array * np.random.random(np_array.shape))], axis=np_array.ndim).astype(np.float32)
ParamType = itk.ctype('double')
Dimension = 3
transform = itk.AffineTransform[ParamType, Dimension].New()
transform.SetMatrix(image.GetDirection())
transform.Scale(image.GetSpacing(), True)
transform.SetOffset(image.GetOrigin())
affine_matrix = affine_transform_to_matrix(transform)
save(np_array_multi, '/tmp/head_multi_apply_transform.nii.gz', affine_matrix)
print('\n2D, scalar image')
image = itk.imread('/home/matt/data/cthead1.png', itk.F)
np_array = np.asarray(image)
ParamType = itk.ctype('double')
Dimension = 2
transform = itk.AffineTransform[ParamType, Dimension].New()
transform.SetCenter(image.GetOrigin())
transform.Scale(image.GetSpacing())
transform.SetOffset(image.GetOrigin())
affine_matrix = affine_transform_to_matrix(transform)
save(np_array, '/tmp/cthead1_apply_transform.nii.gz', affine_matrix)
print('\n2D, multi-channel image')
image = itk.imread('/home/matt/data/cthead1.png', itk.F)
np_array = np.asarray(image)
np_array_multi = np.stack([np_array, (np_array * np.random.random(np_array.shape))], axis=np_array.ndim).astype(np.float32)
ParamType = itk.ctype('double')
Dimension = 2
transform = itk.AffineTransform[ParamType, Dimension].New()
transform.SetCenter(image.GetOrigin())
transform.Scale(image.GetSpacing())
transform.SetOffset(image.GetOrigin())
affine_matrix = affine_transform_to_matrix(transform)
save(np_array_multi, '/tmp/cthead1_multi_apply_transform.nii.gz', affine_matrix)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment