Volume Rendering MRI data using PyOpenGL

  1'''
  2A rudimentary example of rendering basic MRI data 
  3using PyOpenGL and a 3D texture.
  4
  5Author: Stou Sandalski
  6Source: http://www.siafoo.net/snippet/310
  7License: New BSD
  8'''
  9
 10import numpy, wx
 11from numpy import arange, array, float32, int8
 12
 13from VolRenderSkel import *
 14
 15class RenderHead(VolumeRenderSkeleton):
 16       
 17    def __init__(self, parent):
 18        VolumeRenderSkeleton.__init__(self, parent)
 19
 20        self.data_scale = None
 21        self.iso_value = 0.0
 22        self.noise = True
 23
 24        # Shader names (if you have your own source)
 25        self.fragment_src_file = './shaders/head.f.c'
 26        self.vertex_src_file = './shaders/head.v.c'
 27        
 28        self.fragment_shader_src = '''
 29        // Core
 30        uniform sampler1D TransferFunction;
 31        uniform sampler3D VolumeData;
 32        uniform vec3 DataScale;
 33
 34        // Passed in from Vertex shader
 35        varying vec3 ecPosition;
 36        
 37        // Flags
 38        uniform bool EnableLighting;
 39        uniform int  NumberOfLights;
 40        uniform float IsoWeight;
 41
 42        // Lighting Constants (We should get these from the materials?)
 43        float k_a = 0.1; // Ambient
 44        float k_d = 0.7; // Diffuse
 45        float k_s = 0.2; // Specular
 46        float alpha = 1.1; // shininess
 47
 48        const float cellSize = 0.005;
 49
 50        vec3 getGradient(vec3 at)
 51        {
 52            return normalize(vec3(texture3D(VolumeData, at + vec3(cellSize, 0.0, 0.0)).w
 53                                    - texture3D(VolumeData, at - vec3(cellSize, 0.0, 0.0)).w,
 54                                  texture3D(VolumeData, at + vec3(0.0, cellSize, 0.0)).w
 55                                    - texture3D(VolumeData, at - vec3(0.0, cellSize, 0.0)).w,
 56                                  texture3D(VolumeData, at + vec3(0.0, 0.0, cellSize)).w
 57                                    - texture3D(VolumeData, at - vec3(0.0, 0.0, cellSize)).w
 58                                 ));
 59        }
 60        
 61        vec3 getPhongBlinnLighting(vec3 N, vec3 V)
 62        {
 63            vec3 I = vec3(0.0);
 64        
 65            for (int i = 0; i < 1; ++i)
 66            {
 67                vec3 L = normalize(gl_LightSource[i].position.xyz - V);
 68                vec3 R = reflect(L, N);
 69                vec3 H = (L+V) / length(L+V);
 70        
 71                I += k_a*gl_LightSource[i].ambient
 72                   + k_d*dot(L, N)*gl_LightSource[i].diffuse
 73                   + k_s*pow(abs(dot(N, H)), alpha)*gl_LightSource[i].specular;
 74        //           + k_s*pow(abs(dot(R, V)), alpha)*gl_LightSource[i].specular; // Phong
 75            }
 76        
 77            return I;
 78        }
 79        
 80        void main(void)
 81        {
 82            // This is a hack isn't it...
 83            if (gl_TexCoord[0].p < 0.0)
 84            {
 85                gl_FragColor = vec4(0.0, 0.0, 0.0, 0.0);
 86                return;
 87            }
 88            
 89            // Switch the order of the `stp` to like `spt`, etc. if the 
 90            // image is in a weird orientation... but you also need
 91            // to rearange the DataScale.xyz in the same way
 92            vec3 lookup = gl_TexCoord[0].stp;
 93
 94            float weight = texture3D(VolumeData, lookup).w;
 95
 96            if (weight >= IsoWeight + 0.02 || weight <= IsoWeight - 0.02)
 97            {
 98//                Using ISO Weight one could potentially generate
 99//                fake iso-surfaces which could serve as a preview mechanism
100//                for real surface extraction algorithms like MarchingCubes, etc.
101
102//                gl_FragColor = vec4(0.0, 0.0, 0.0, 0.0);
103//                return
104            }
105            
106            vec4 color = texture1D(TransferFunction, weight);
107
108            if (EnableLighting)
109            {
110                color.rgb *= getPhongBlinnLighting(getGradient(lookup), ecPosition);
111            }
112
113            gl_FragColor = color;
114        }
115        '''
116        
117        self.vertex_shader_src = '''
118        uniform vec3 DataScale;
119        varying vec3 ecPosition;
120        
121        void main(void)
122        {
123            gl_TexCoord[0] = vec4((gl_Vertex.xyz / DataScale) + DataScale - 1.0, 0.0);
124            ecPosition = vec3(gl_ModelViewMatrix * gl_Vertex);
125            gl_Position = ftransform();
126        }
127        '''
128
129    def SetupUniforms(self):
130        VolumeRenderSkeleton.SetupUniforms(self)
131        
132        # Init the texture units
133        glActiveTexture(GL_TEXTURE1)
134        glBindTexture(GL_TEXTURE_3D, self.vol_data)
135        # Setup VolumeData on Texture Unit 2 (we are using 1 for TransferFunction)
136        glUniform1i(glGetUniformLocation(self.program, "VolumeData"), 1)
137        glUniform1f(glGetUniformLocation(self.program, "IsoWeight"), 
138                    self.iso_value)
139        glUniform3f(glGetUniformLocation(self.program, "DataScale"), 
140                    self.data_scale[0], self.data_scale[1], self.data_scale[2])
141
142    def LoadVolumeData(self):
143        VolumeRenderSkeleton.LoadVolumeData(self)
144
145        if self.data_set.endswith('nii') or self.data_set.endswith('nii.gz'):
146            try:
147                from nifti import NiftiImage
148            except ImportError:
149                print "Apparently you don't have PyNIfTI installed, see http://www.siafoo.net/snippet/310 for instructions"
150                exit(1)
151                
152            nim = NiftiImage(canvas.data_set)
153            data = nim.data.reshape(nim.data.shape[2], nim.data.shape[1], 
154                                    nim.data.shape[0])
155        elif self.data_set.endswith('vol'):
156            print "Sorry I can't read '.vol' files yet, why don't you submit a patch?"
157            exit(1) 
158        elif self.data_set.endswith('hdr'):
159            # Use the header to figure out the shape of the data
160            # then load the raw data and reshape the array 
161            data = numpy.frombuffer(open(self.data_set.replace('.hdr', '.dat'), 'rb').read(), 
162                                    numpy.int8)\
163                        .reshape([int(x) for x in open(self.data_set).readline().split()])
164        else:
165            print "Sorry I can't read your file, why don't you submit a patch?"
166            exit(1) 
167            
168        self.data_set_size = array(data.shape, dtype=float32)
169        self.data_scale = self.data_set_size / self.data_set_size.max()
170        print 'Running with: ', self.data_set, self.data_set_size, self.data_scale, self.data_set_size.max()
171       
172        # Create Texture    
173        self.vol_data = glGenTextures(1)
174        glPixelStorei(GL_UNPACK_ALIGNMENT,1)
175        glBindTexture(GL_TEXTURE_3D, self.vol_data)
176        
177        glTexParameterf(GL_TEXTURE_3D, GL_TEXTURE_WRAP_S, GL_CLAMP)
178        glTexParameterf(GL_TEXTURE_3D, GL_TEXTURE_WRAP_T, GL_CLAMP)
179        glTexParameterf(GL_TEXTURE_3D, GL_TEXTURE_WRAP_R, GL_CLAMP)
180        
181        glTexParameterf(GL_TEXTURE_3D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
182        glTexParameterf(GL_TEXTURE_3D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
183    
184        glTexImage3D(GL_TEXTURE_3D, 0, GL_ALPHA, 
185                     self.data_set_size[0], self.data_set_size[1], self.data_set_size[2], 
186                     0, GL_ALPHA, GL_UNSIGNED_BYTE, data)
187        
188        self.texture_list.append(self.vol_data)
189        
190        return
191
192    def OnKeyDown(self, event):
193        VolumeRenderSkeleton.OnKeyDown(self, event)
194        
195        if event.GetKeyCode() == ord('n'):
196            self.noise = not self.noise
197            print 'Noise', self.noise 
198            self.Refresh()
199
200if __name__ == '__main__':
201    app = wx.App()
202    frame = wx.Frame(None, -1, 'Volume Data Visualizer', wx.DefaultPosition, wx.Size(600, 600))
203    canvas = RenderHead(frame)
204    
205    try:
206        canvas.data_set = sys.argv[1]
207    except:
208        canvas.data_set = 'avg152T1_RL_nifti.nii.gz' # NIfTI file
209        canvas.data_set = '../data/MRI_head.hdr'
210    
211    frame.Show()
212    app.MainLoop()