Commit febc3a96 authored by David Gobbi's avatar David Gobbi Committed by David Partyka
Browse files

ENH: Also convert vtkImageExportToArray from Numeric to numpy.

parent 74d7e025
"""
vtkImageExportToArray - a NumPy front-end to vtkImageExport
This class converts a VTK image to a Numeric Python array.
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 the LLNL Numeric Python distribution
(http://numpy.sf.net)
To use this class, you must have numpy installed (http://numpy.scipy.org)
Methods
SetInput(input) -- connect to VTK image pipeline
GetArray() -- execute pipeline and return a Numeric array
SetInputConnection(vtkAlgorithmOutput) -- connect to VTK image pipeline
SetInput(vtkImageData) -- set an vtkImageData to export
GetArray() -- execute pipeline and return a numpy array
Convert VTK_UNSIGNED_SHORT to python Int
(this might be necessary because Python doesn't support unsigned short,
the default is to cast unsigned short to signed short).
SetConvertUnsignedShortToInt(yesno)
ConvertUnsignedShortToIntOn()
ConvertUnsignedShortToIntOff()
From vtkImageExport
Methods from vtkImageExport
GetDataExtent()
GetDataSpacing()
GetDataOrigin()
"""
import Numeric
import umath
import numpy
from vtk import vtkImageExport
from vtk import VTK_SIGNED_CHAR
from vtk import VTK_UNSIGNED_CHAR
......@@ -40,39 +35,29 @@ from vtk import VTK_UNSIGNED_LONG
from vtk import VTK_FLOAT
from vtk import VTK_DOUBLE
_NEW_NUMERIC = 0
try:
val = float(Numeric.__version__)
except ValueError:
_NEW_NUMERIC = 0
else:
if val > 20.0:
_NEW_NUMERIC = 1
else:
_NEW_NUMERIC = 0
class vtkImageExportToArray:
def __init__(self):
self.__export = vtkImageExport()
self.__ConvertUnsignedShortToInt = 0
self.__ConvertUnsignedShortToInt = False
# type dictionary: note that python doesn't support
# unsigned integers!
# type dictionary
__typeDict = { VTK_SIGNED_CHAR:Numeric.Int8,
VTK_UNSIGNED_CHAR:Numeric.UnsignedInt8,
VTK_SHORT:Numeric.Int16,
VTK_UNSIGNED_SHORT:Numeric.Int16,
VTK_INT:Numeric.Int32,
VTK_FLOAT:Numeric.Float32,
VTK_DOUBLE:Numeric.Float64 }
__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 }
......@@ -84,12 +69,15 @@ class vtkImageExportToArray:
return self.__ConvertUnsignedShortToInt
def ConvertUnsignedShortToIntOn(self):
self.__ConvertUnsignedShortToInt = 1
self.__ConvertUnsignedShortToInt = True
def ConvertUnsignedShortToIntOff(self):
self.__ConvertUnsignedShortToInt = 0
self.__ConvertUnsignedShortToInt = False
# set the input
def SetInputConnection(self,input):
return self.__export.SetInputConnection(input)
def SetInput(self,input):
return self.__export.SetInput(input)
......@@ -107,24 +95,13 @@ class vtkImageExportToArray:
extent[1]-extent[0]+1)
if (numComponents > 1):
dim = dim + (numComponents,)
size = dim[0]*dim[1]*dim[2]*numComponents*self.__sizeDict[type]
if _NEW_NUMERIC:
imArray = Numeric.zeros((size,),Numeric.UnsignedInt8)
self.__export.Export(imArray)
else:
imString = Numeric.zeros((size,),
Numeric.UnsignedInt8).tostring()
self.__export.Export(imString)
imArray = Numeric.fromstring(imString,self.__typeDict[type])
# just to remind myself of the dangers of memory management
del imString
# reshape array appropriately.
imArray = Numeric.reshape(imArray, dim)
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(Numeric.Int32),0xffff)
imArray = umath.bitwise_and(imArray.astype('i'),0xffff)
return imArray
......
......@@ -2,19 +2,14 @@
vtkImageImportFromArray: a NumPy front-end to vtkImageImport
Load a python array into a vtk image.
To use this class,you must have NumPy installed (http://numpy.scipy.org/)
To use this class, you must have NumPy installed (http://numpy.scipy.org/)
Methods:
GetOutput() -- connect to VTK image pipeline
SetArray() -- set the array to load in
Convert python 'Int' to VTK_UNSIGNED_SHORT:
(python doesn't support unsigned short, so this might be necessary)
SetConvertIntToUnsignedShort(yesno)
ConvertIntToUnsignedShortOn()
ConvertIntToUnsignedShortOff()
SetArray() -- set the numpy array to load
Update() -- generate the output
GetOutput() -- get the image as vtkImageData
GetOutputPort() -- connect to VTK pipeline
Methods from vtkImageImport:
(if you don't set these, sensible defaults will be used)
......@@ -50,8 +45,6 @@ class vtkImageImportFromArray:
'H':VTK_UNSIGNED_SHORT, # uint16
'i':VTK_INT, # int32
'I':VTK_UNSIGNED_INT, # uint32
'l':VTK_LONG, # int64
'L':VTK_UNSIGNED_LONG, # uint64
'f':VTK_FLOAT, # float32
'd':VTK_DOUBLE, # float64
'F':VTK_FLOAT, # float32
......@@ -64,8 +57,6 @@ class vtkImageImportFromArray:
VTK_UNSIGNED_SHORT:2,
VTK_INT:4,
VTK_UNSIGNED_INT:4,
VTK_LONG:4,
VTK_UNSIGNED_LONG:4,
VTK_FLOAT:4,
VTK_DOUBLE:8 }
......@@ -96,7 +87,6 @@ class vtkImageImportFromArray:
# import an array
def SetArray(self,imArray):
self.__Array = imArray
imTmpArr = imArray.flat
numComponents = 1
dim = imArray.shape
if len(dim) == 0:
......@@ -119,13 +109,11 @@ class vtkImageImportFromArray:
complexComponents = 2
if (self.__ConvertIntToUnsignedShort and typecode == 'i'):
imTmpArr = imArray.astype('h').flat
imArray = imArray.astype('h')
ar_type = VTK_UNSIGNED_SHORT
else:
imTmpArr = imArray.flat
size = len(imTmpArr)*self.__sizeDict[ar_type]*complexComponents
self.__import.CopyImportVoidPointer(imTmpArr,size)
size = len(imArray.flat)*self.__sizeDict[ar_type]*complexComponents
self.__import.CopyImportVoidPointer(imArray, size)
self.__import.SetDataScalarType(ar_type)
self.__import.SetNumberOfScalarComponents(numComponents)
extent = self.__import.GetDataExtent()
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment