Save an image from a numpy array and an affine matrix

save_image_from_numpy_array_affine_matrix.py

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)