init

2026-02-12 23:22:11 +08:00
parent 7b09eb3d89
commit 89660bba4e
5988 changed files with 2517516 additions and 0 deletions
@@ -0,0 +1,116 @@
+"""
+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()