117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
"""
|
|
vtkImageExportToArray - a NumPy front-end to vtkImageExport
|
|
|
|
This class converts a VTK image to a numpy array. The output
|
|
array will always have 3 dimensions (or 4, if the image had
|
|
multiple scalar components).
|
|
|
|
To use this class, you must have numpy installed (http://numpy.scipy.org)
|
|
|
|
Methods
|
|
|
|
SetInputConnection(vtkAlgorithmOutput) -- connect to VTK image pipeline
|
|
SetInputData(vtkImageData) -- set an vtkImageData to export
|
|
GetArray() -- execute pipeline and return a numpy array
|
|
|
|
Methods from vtkImageExport
|
|
|
|
GetDataExtent()
|
|
GetDataSpacing()
|
|
GetDataOrigin()
|
|
"""
|
|
|
|
import numpy
|
|
import numpy.core.umath as umath
|
|
|
|
from vtkmodules.vtkIOImage import vtkImageExport
|
|
from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
|
|
from vtkmodules.vtkCommonCore import VTK_SIGNED_CHAR
|
|
from vtkmodules.vtkCommonCore import VTK_UNSIGNED_CHAR
|
|
from vtkmodules.vtkCommonCore import VTK_SHORT
|
|
from vtkmodules.vtkCommonCore import VTK_UNSIGNED_SHORT
|
|
from vtkmodules.vtkCommonCore import VTK_INT
|
|
from vtkmodules.vtkCommonCore import VTK_UNSIGNED_INT
|
|
from vtkmodules.vtkCommonCore import VTK_LONG
|
|
from vtkmodules.vtkCommonCore import VTK_UNSIGNED_LONG
|
|
from vtkmodules.vtkCommonCore import VTK_FLOAT
|
|
from vtkmodules.vtkCommonCore import VTK_DOUBLE
|
|
|
|
|
|
class vtkImageExportToArray:
|
|
def __init__(self):
|
|
self.__export = vtkImageExport()
|
|
self.__ConvertUnsignedShortToInt = False
|
|
|
|
# type dictionary
|
|
|
|
__typeDict = { VTK_SIGNED_CHAR:'b',
|
|
VTK_UNSIGNED_CHAR:'B',
|
|
VTK_SHORT:'h',
|
|
VTK_UNSIGNED_SHORT:'H',
|
|
VTK_INT:'i',
|
|
VTK_UNSIGNED_INT:'I',
|
|
VTK_FLOAT:'f',
|
|
VTK_DOUBLE:'d'}
|
|
|
|
__sizeDict = { VTK_SIGNED_CHAR:1,
|
|
VTK_UNSIGNED_CHAR:1,
|
|
VTK_SHORT:2,
|
|
VTK_UNSIGNED_SHORT:2,
|
|
VTK_INT:4,
|
|
VTK_UNSIGNED_INT:4,
|
|
VTK_FLOAT:4,
|
|
VTK_DOUBLE:8 }
|
|
|
|
# convert unsigned shorts to ints, to avoid sign problems
|
|
def SetConvertUnsignedShortToInt(self,yesno):
|
|
self.__ConvertUnsignedShortToInt = yesno
|
|
|
|
def GetConvertUnsignedShortToInt(self):
|
|
return self.__ConvertUnsignedShortToInt
|
|
|
|
def ConvertUnsignedShortToIntOn(self):
|
|
self.__ConvertUnsignedShortToInt = True
|
|
|
|
def ConvertUnsignedShortToIntOff(self):
|
|
self.__ConvertUnsignedShortToInt = False
|
|
|
|
# set the input
|
|
def SetInputConnection(self,input):
|
|
return self.__export.SetInputConnection(input)
|
|
|
|
def SetInputData(self,input):
|
|
return self.__export.SetInputData(input)
|
|
|
|
def GetInput(self):
|
|
return self.__export.GetInput()
|
|
|
|
def GetArray(self):
|
|
self.__export.Update()
|
|
input = self.__export.GetInput()
|
|
extent = input.GetExtent()
|
|
type = input.GetScalarType()
|
|
numComponents = input.GetNumberOfScalarComponents()
|
|
dim = (extent[5]-extent[4]+1,
|
|
extent[3]-extent[2]+1,
|
|
extent[1]-extent[0]+1)
|
|
if (numComponents > 1):
|
|
dim = dim + (numComponents,)
|
|
|
|
imArray = numpy.zeros(dim, self.__typeDict[type])
|
|
self.__export.Export(imArray)
|
|
|
|
# convert unsigned short to int to avoid sign issues
|
|
if (type == VTK_UNSIGNED_SHORT and self.__ConvertUnsignedShortToInt):
|
|
imArray = umath.bitwise_and(imArray.astype('i'),0xffff)
|
|
|
|
return imArray
|
|
|
|
def GetDataExtent(self):
|
|
return self.__export.GetDataExtent()
|
|
|
|
def GetDataSpacing(self):
|
|
return self.__export.GetDataSpacing()
|
|
|
|
def GetDataOrigin(self):
|
|
return self.__export.GetDataOrigin()
|