Load NumPy Array into a VTK image

  1"""
  2vtkImageImportFromArray: a NumPy front-end to vtkImageImport
  3
  4Load a python array into a vtk image.
  5To use this class, you must have NumPy installed (http://numpy.scipy.org/)
  6
  7Methods:
  8
  9  GetOutput() -- connect to VTK image pipeline
 10  SetArray()  -- set the array to load in
 11  
 12Convert python 'Int' to VTK_UNSIGNED_SHORT:
 13(python doesn't support unsigned short, so this might be necessary)
 14
 15  SetConvertIntToUnsignedShort(yesno)
 16  ConvertIntToUnsignedShortOn()
 17  ConvertIntToUnsignedShortOff()
 18
 19Methods from vtkImageImport: 
 20(if you don't set these, sensible defaults will be used)
 21
 22  SetDataExtent()
 23  SetDataSpacing()
 24  SetDataOrigin()
 25"""
 26
 27import Numeric
 28from vtk import vtkImageImport
 29from vtk.util.vtkConstants import *
 30
 31class vtkImageImportFromArray:
 32    def __init__(self):
 33        self.__import = vtkImageImport()
 34        self.__ConvertIntToUnsignedShort = False
 35        self.__Array = None
 36
 37    # type dictionary: note that python doesn't support
 38    # unsigned integers properly!
 39    __typeDict = {'b':VTK_CHAR,             # int8
 40                  'B':VTK_UNSIGNED_CHAR,    # uint8
 41                  'h':VTK_SHORT,            # int16
 42                  'H':VTK_UNSIGNED_SHORT,   # uint16
 43                  'i':VTK_INT,              # int32
 44                  'I':VTK_UNSIGNED_INT,     # uint32
 45                  'l':VTK_LONG,             # int64
 46                  'L':VTK_UNSIGNED_LONG,    # uint64
 47                  'f':VTK_FLOAT,            # float32
 48                  'd':VTK_DOUBLE,           # float64
 49                  }
 50
 51    # convert 'Int32' to 'unsigned short'
 52    def SetConvertIntToUnsignedShort(self, yesno):
 53        self.__ConvertIntToUnsignedShort = yesno
 54
 55    def GetConvertIntToUnsignedShort(self):
 56        return self.__ConvertIntToUnsignedShort
 57    
 58    def ConvertIntToUnsignedShortOn(self):
 59        self.__ConvertIntToUnsignedShort = True
 60
 61    def ConvertIntToUnsignedShortOff(self):
 62        self.__ConvertIntToUnsignedShort = False
 63
 64    # get the output
 65    def GetOutputPort(self):
 66        return self.__import.GetOutputPort()
 67
 68    # get the output
 69    def GetOutput(self):
 70        return self.__import.GetOutput()
 71
 72    # import an array
 73    def SetArray(self, imArray):
 74        self.__Array = imArray
 75        imString = imArray.tostring()
 76        numComponents = 1
 77        dim = imArray.shape
 78
 79        if (len(dim) == 4):
 80            numComponents = dim[3]
 81            dim = (dim[0], dim[1], dim[2])
 82
 83
 84        typecode = imArray.dtype.char
 85       
 86        ar_type = self.__typeDict[typecode]
 87
 88        if (typecode == 'F' or typecode == 'D'):
 89            numComponents = numComponents * 2
 90
 91        if (self.__ConvertIntToUnsignedShort and typecode == 'i'):
 92            imString = imArray.astype(Numeric.Int16).tostring()
 93            ar_type = VTK_UNSIGNED_SHORT
 94        else:
 95            imString = imArray.tostring()
 96            
 97        self.__import.CopyImportVoidPointer(imString, len(imString))
 98        self.__import.SetDataScalarType(ar_type)
 99        self.__import.SetNumberOfScalarComponents(numComponents)
100        extent = self.__import.GetDataExtent()
101        self.__import.SetDataExtent(extent[0], extent[0] + dim[2] - 1,
102                                    extent[2], extent[2] + dim[1] - 1,
103                                    extent[4], extent[4] + dim[0] - 1)
104        self.__import.SetWholeExtent(extent[0], extent[0] + dim[2] - 1,
105                                     extent[2], extent[2] + dim[1] - 1,
106                                     extent[4], extent[4] + dim[0] - 1)
107
108    def GetArray(self):
109        return self.__Array
110        
111    # a whole bunch of methods copied from vtkImageImport
112    def SetDataExtent(self, extent):
113        self.__import.SetDataExtent(extent)
114
115    def GetDataExtent(self):
116        return self.__import.GetDataExtent()
117    
118    def SetDataSpacing(self, spacing):
119        self.__import.SetDataSpacing(spacing)
120
121    def GetDataSpacing(self):
122        return self.__import.GetDataSpacing()
123    
124    def SetDataOrigin(self, origin):
125        self.__import.SetDataOrigin(origin)
126
127    def GetDataOrigin(self):
128        return self.__import.GetDataOrigin()