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()