Volume Rendering MRI data using VTK

  1'''
  2A quick and dirty VTK based viewer for volume data.
  3'''
  4
  5import numpy, sys, wx, vtk
  6
  7try:
  8    from transfer_function import TransferFunctionWidget
  9except ImportError:
 10    print 'You need to get the TransferFunction from: "http://www.siafoo.net/snippet/122"'
 11    exit(1)
 12
 13try:
 14    from wxVTKRenderWindowInteractor import wxVTKRenderWindowInteractor
 15except ImportError:
 16    print 'Warning: If you are getting flickering get wxVTKRenderWindowInteractor from: "http://www.siafoo.net/snippet/312"'
 17    from vtk.wx.wxVTKRenderWindowInteractor import wxVTKRenderWindowInteractor
 18
 19try:       
 20    from vtkImageImportFromArray import vtkImageImportFromArray
 21except ImportError:
 22    print 'You need to get the updated vtkImageImportFromArray from: "http://www.siafoo.net/snippet/313" '
 23    exit(1)
 24    
 25
 26class TransferGraph(wx.Dialog):
 27
 28    def __init__(self, parent, id=wx.ID_ANY, title="", pos=wx.DefaultPosition,
 29                 size=wx.DefaultSize, style=wx.DEFAULT_FRAME_STYLE):
 30
 31        wx.Dialog.__init__(self, parent, id, title, pos, size, style)
 32        self.mainPanel = wx.Panel(self, -1)
 33
 34        # Create some CustomCheckBoxes
 35        self.t_function = TransferFunctionWidget(self.mainPanel, -1, "", size=wx.Size(300, 150))
 36       
 37        # Layout the items with sizers
 38        mainSizer = wx.BoxSizer(wx.VERTICAL)
 39        mainSizer.Add(self.mainPanel, 1, wx.EXPAND)
 40        
 41        self.SetSizer(mainSizer)
 42        mainSizer.Layout()
 43        
 44    def get_function(self):
 45        return self.t_function
 46
 47class VolumeRenderHeadVTK(wxVTKRenderWindowInteractor):
 48    def __init__(self, parent, data_set):
 49        wxVTKRenderWindowInteractor.__init__(self, parent, -1, size=parent.GetSize())
 50        
 51        ren = vtk.vtkRenderer()
 52        self.GetRenderWindow().AddRenderer(ren)
 53
 54        img = self.LoadVolumeData(data_set)
 55
 56        pix_diag = 5.0
 57
 58#        volMapper = vtk.vtkVolumeTextureMapper3D()
 59
 60        volMapper = vtk.vtkVolumeRayCastMapper()
 61        compositeFunction = vtk.vtkVolumeRayCastCompositeFunction()
 62        compositeFunction.SetCompositeMethodToInterpolateFirst()
 63        volMapper.SetVolumeRayCastFunction(compositeFunction)
 64
 65        volMapper.SetSampleDistance(pix_diag / 5.0)
 66        volMapper.SetInputConnection(img.GetOutputPort())
 67
 68        # Transfer Functions
 69        self.opacity_tf = vtk.vtkPiecewiseFunction()
 70        self.color_tf = vtk.vtkColorTransferFunction()
 71
 72        # The property describes how the data will look
 73        self.volProperty = volProperty = vtk.vtkVolumeProperty()
 74        volProperty.SetColor(self.color_tf)
 75        volProperty.SetScalarOpacity(self.opacity_tf)
 76        volProperty.ShadeOn()
 77        volProperty.SetInterpolationTypeToLinear()
 78        volProperty.SetScalarOpacityUnitDistance(pix_diag)
 79
 80        vol = vtk.vtkVolume()
 81        vol.SetMapper(volMapper)
 82        vol.SetProperty(volProperty)
 83        
 84        ren.AddVolume(vol)
 85        
 86        boxWidget = vtk.vtkBoxWidget()
 87        boxWidget.SetInteractor(self)
 88        boxWidget.SetPlaceFactor(1.0)
 89        
 90        # The implicit function vtkPlanes is used in conjunction with the
 91        # volume ray cast mapper to limit which portion of the volume is
 92        # volume rendered.
 93        planes = vtk.vtkPlanes()
 94        def ClipVolumeRender(obj, event):
 95            obj.GetPlanes(planes)
 96            volMapper.SetClippingPlanes(planes)
 97        
 98        # Place the interactor initially. The output of the reader is used to
 99        # place the box widget.
100        boxWidget.SetInput(img.GetOutput())
101        boxWidget.PlaceWidget()
102        boxWidget.InsideOutOn()
103        boxWidget.AddObserver("InteractionEvent", ClipVolumeRender)
104
105        outline = vtk.vtkOutlineFilter()
106        outline.SetInputConnection(img.GetOutputPort())
107        outlineMapper = vtk.vtkPolyDataMapper()
108        outlineMapper.SetInputConnection(outline.GetOutputPort())
109        outlineActor = vtk.vtkActor()
110        outlineActor.SetMapper(outlineMapper)
111
112        outlineProperty = boxWidget.GetOutlineProperty()
113        outlineProperty.SetRepresentationToWireframe()
114        outlineProperty.SetAmbient(1.0)
115        outlineProperty.SetAmbientColor(1, 1, 1)
116        outlineProperty.SetLineWidth(3)
117        
118        selectedOutlineProperty = boxWidget.GetSelectedOutlineProperty()
119        selectedOutlineProperty.SetRepresentationToWireframe()
120        selectedOutlineProperty.SetAmbient(1.0)
121        selectedOutlineProperty.SetAmbientColor(1, 0, 0)
122        selectedOutlineProperty.SetLineWidth(3)
123
124        ren.AddActor(outlineActor)
125        
126        style = vtk.vtkInteractorStyleTrackballCamera()
127        self.SetInteractorStyle(style)
128
129        self.t_graph = TransferGraph(self)
130        self.SetFocus()
131        
132        # So we know when new values are added / changed on the tgraph
133        self.t_graph.Connect(-1, -1, wx.wxEVT_COMMAND_SLIDER_UPDATED, 
134                             self.OnTGraphUpdate)
135       
136        self.OnTGraphUpdate(None)
137        
138        self.lighting = True
139
140        # This is the transfer graph
141        self.t_graph.Show()
142
143    def OnTGraphUpdate(self, event):
144
145        self.color_tf.RemoveAllPoints()
146        self.opacity_tf.RemoveAllPoints()
147        
148        for p in self.t_graph.t_function.points:
149            rgba = p.get_rgba()
150            self.color_tf.AddRGBPoint(p.value, rgba[0]/255.0, 
151                                 rgba[1]/255.0, rgba[2]/255.0)
152
153            self.opacity_tf.AddPoint(p.value, rgba[3])
154
155        self.Refresh()
156
157    def LoadVolumeData(self, data_set):
158
159        if data_set.endswith('nii') or data_set.endswith('nii.gz'):
160            try:
161                from nifti import NiftiImage
162            except ImportError:
163                print "Apparently you don't have PyNIfTI installed, see http://www.siafoo.net/snippet/310 for instructions"
164                exit(1)
165            nim = NiftiImage(data_set)
166            img_data = nim.data
167        elif data_set.endswith('hdr'):
168            # Use the header to figure out the shape of the data
169            # then load the raw data and reshape the array 
170            shape = [int(x) for x in open(data_set).readline().split()]
171            img_data = numpy.frombuffer(open(data_set.replace('.hdr', '.dat'), 'rb').read(),
172                                        numpy.uint8)\
173                        .reshape((shape[2], shape[1], shape[0]))
174                        
175        img = vtkImageImportFromArray()
176        img.SetArray(img_data)
177        
178        return img
179
180    def OnKeyDown(self, event):
181        # Yes, I did just hack this... maybe one day when things work
182        # the way the docs claim... hacks won't be necessary.
183        
184        key_code = event.GetKeyCode()
185        
186        if key_code == ord('l'):
187            self.lighting = not self.lighting
188            print 'Lighting', self.lighting 
189            self.volProperty.SetShade(self.lighting)
190            self.Refresh()
191        else:
192            wxVTKRenderWindowInteractor.OnKeyDown(self, event)
193
194
195if __name__ == '__main__':
196    app = wx.App()
197    frame = wx.Frame(None, -1, 'Volume Rendering with VTK', wx.DefaultPosition, wx.Size(600, 600))
198    
199    try:
200        data_filename = sys.argv[1]
201    except:
202#        data_filename = 'avg152T1_RL_nifti.nii.gz' # NIfTI file
203        data_filename = '../data/MRI_head.hdr'
204        
205    canvas = VolumeRenderHeadVTK(frame, data_filename)
206
207    frame.Show()
208    app.MainLoop()