Skip to content

Commit

Permalink
ENH: Add itk.vtk_image_from_image and itk.image_from_vtk_image
Browse files Browse the repository at this point in the history
For conversion to/from and itk.Image and a vtk.vtkImageData.
  • Loading branch information
thewtex committed Apr 29, 2020
1 parent f7500c5 commit f13e760
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 0 deletions.
25 changes: 25 additions & 0 deletions Wrapping/Generators/Python/Tests/extras.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,3 +390,28 @@ def custom_callback(name, progress):
except ImportError:
print('xarray not imported. Skipping xarray conversion tests')
pass

# vtk conversion
try:
import vtk
import numpy as np
print('Testing vtk conversion')

image = itk.image_from_array(np.random.rand(2,3,4))
vtk_image = itk.vtk_image_from_image(image)
image_round = itk.image_from_vtk_image(vtk_image)
assert(np.array_equal(itk.origin(image), itk.origin(image_round)))
assert(np.array_equal(itk.spacing(image), itk.spacing(image_round)))
assert(np.array_equal(itk.size(image), itk.size(image_round)))
assert(np.array_equal(itk.array_view_from_image(image), itk.array_view_from_image(image_round)))

image = itk.image_from_array(np.random.rand(5,4,2).astype(np.float32), is_vector=True)
vtk_image = itk.vtk_image_from_image(image)
image_round = itk.image_from_vtk_image(vtk_image)
assert(np.array_equal(itk.origin(image), itk.origin(image_round)))
assert(np.array_equal(itk.spacing(image), itk.spacing(image_round)))
assert(np.array_equal(itk.size(image), itk.size(image_round)))
assert(np.array_equal(itk.array_view_from_image(image), itk.array_view_from_image(image_round)))
except ImportError:
print('vtk not imported. Skipping vtk conversion tests')
pass
67 changes: 67 additions & 0 deletions Wrapping/Generators/Python/itkExtras.py
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,73 @@ def image_from_xarray(data_array):

return itk_image

def vtk_image_from_image(image):
"""Convert an itk.Image to a vtk.vtkImageData."""
import itk
import vtk
from vtk.util.numpy_support import numpy_to_vtk

array = itk.array_view_from_image(image)

vtk_image = vtk.vtkImageData()
data_array = numpy_to_vtk(array.reshape(-1))
data_array.SetNumberOfComponents(image.GetNumberOfComponentsPerPixel())
data_array.SetName('Scalars')
# Always set Scalars for (future?) multi-component volume rendering
vtk_image.GetPointData().SetScalars(data_array)
dim = image.GetImageDimension()
spacing = [1.0,] * 3
spacing[:dim] = image.GetSpacing()
vtk_image.SetSpacing(spacing)
origin = [0.0,] * 3
origin[:dim] = image.GetOrigin()
vtk_image.SetOrigin(origin)
dims = [1,] * 3
dims[:dim] = itk.size(image)
vtk_image.SetDimensions(dims)
# Todo: Add Direction with VTK 9
if image.GetImageDimension() == 3:
PixelType = itk.template(image)[1][0]
if PixelType == itk.Vector:
vtk_image.GetPointData().SetVectors(data_array)
elif PixelType == itk.CovariantVector:
vtk_image.GetPointData().SetVectors(data_array)
elif PixelType == itk.SymmetricSecondRankTensor:
vtk_image.GetPointData().SetTensors(data_array)
elif PixelType == itk.DiffusionTensor3D:
vtk_image.GetPointData().SetTensors(data_array)
return vtk_image

def image_from_vtk_image(vtk_image):
"""Convert a vtk.vtkImageData to an itk.Image."""
import itk
from vtk.util.numpy_support import vtk_to_numpy

point_data = vtk_image.GetPointData()
array = vtk_to_numpy(point_data.GetScalars())
array = array.reshape(-1)
is_vector = point_data.GetScalars().GetNumberOfComponents() != 1
dims = list(vtk_image.GetDimensions())
if is_vector and dims[-1] == 1:
# 2D
dims = dims[:2]
dims.reverse()
dims.append(point_data.GetScalars().GetNumberOfComponents())
else:
dims.reverse()
array.shape = tuple(dims)
image = itk.image_view_from_array(array, is_vector)

dim = image.GetImageDimension()
spacing = [1.0] * dim
spacing[:dim] = vtk_image.GetSpacing()[:dim]
image.SetSpacing(spacing)
origin = [0.0] * dim
origin[:dim] = vtk_image.GetOrigin()[:dim]
image.SetOrigin(origin)
# Todo: Add Direction with VTK 9
return image

# return an image
from itkTemplate import image, output

Expand Down

0 comments on commit f13e760

Please sign in to comment.