1 # GL interoperability example, by Peter Berrington.
   2 # Draws a rotating teapot, using cuda to invert the RGB value
   3 # each frame
   4 
   5 from OpenGL.GL import *
   6 from OpenGL.GLUT import *
   7 from OpenGL.GLU import *
   8 from OpenGL.GL.ARB.vertex_buffer_object import *
   9 from OpenGL.GL.ARB.pixel_buffer_object import *
  10 
  11 
  12 import numpy, sys, time
  13 import pycuda.driver as cuda_driver
  14 import pycuda.gl as cuda_gl
  15 from pycuda.compiler import SourceModule
  16 
  17 #this is all munged together from the CUDA SDK postprocessGL example.
  18 
  19 initial_size = 512,512
  20 current_size = initial_size
  21 animate = True
  22 enable_cuda = True
  23 window = None     # Number of the glut window.
  24 time_of_last_draw = 0.0
  25 time_of_last_titleupdate = 0.0
  26 frames_per_second = 0.0
  27 frame_counter = 0
  28 output_texture = None # pointer to offscreen render target
  29 (source_pbo, dest_pbo, cuda_module, invert,
  30  pycuda_source_pbo, pycuda_dest_pbo) = [None]*6
  31 heading,pitch,bank = [0.0]*3
  32 
  33 def create_PBOs(w,h):
  34     global source_pbo, dest_pbo, pycuda_source_pbo, pycuda_dest_pbo
  35     num_texels = w*h
  36     data = numpy.zeros((num_texels,4),numpy.uint8)
  37     source_pbo = glGenBuffers(1)
  38     glBindBuffer(GL_ARRAY_BUFFER, source_pbo)
  39     glBufferData(GL_ARRAY_BUFFER, data, GL_DYNAMIC_DRAW)
  40     glBindBuffer(GL_ARRAY_BUFFER, 0)
  41     pycuda_source_pbo = cuda_gl.BufferObject(long(source_pbo))
  42     dest_pbo = glGenBuffers(1)
  43     glBindBuffer(GL_ARRAY_BUFFER, dest_pbo)
  44     glBufferData(GL_ARRAY_BUFFER, data, GL_DYNAMIC_DRAW)
  45     glBindBuffer(GL_ARRAY_BUFFER, 0)
  46     pycuda_dest_pbo = cuda_gl.BufferObject(long(dest_pbo))
  47 
  48 def destroy_PBOs():
  49     global source_pbo, dest_pbo, pycuda_source_pbo, pycuda_dest_pbo
  50     for pbo in [source_pbo, dest_pbo]:
  51         glBindBuffer(GL_ARRAY_BUFFER, long(pbo))
  52         glDeleteBuffers(1, long(pbo));
  53         glBindBuffer(GL_ARRAY_BUFFER, 0)
  54     source_pbo,dest_pbo,pycuda_source_pbo,pycuda_dest_pbo = [None]*4
  55 
  56 def create_texture(w,h):
  57     global output_texture
  58     output_texture = glGenTextures(1)
  59     glBindTexture(GL_TEXTURE_2D, output_texture)
  60     # set basic parameters
  61     glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
  62     glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
  63     glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
  64     glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
  65     # buffer data
  66     glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA,
  67                  w, h, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
  68 
  69 def destroy_texture():
  70     global output_texture
  71     glDeleteTextures(output_texture);
  72     output_texture = None
  73 
  74 def init_gl():
  75     Width, Height = current_size
  76     glClearColor(0.1, 0.1, 0.5, 1.0)
  77     glDisable(GL_DEPTH_TEST)
  78     glViewport(0, 0, Width, Height)
  79     glMatrixMode(GL_PROJECTION);
  80     glLoadIdentity();
  81     gluPerspective(60.0, Width/float(Height), 0.1, 10.0)
  82     glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)
  83     glEnable(GL_LIGHT0)
  84     red   = ( 1.0, 0.1, 0.1, 1.0 )
  85     white = ( 1.0, 1.0, 1.0, 1.0 )
  86     glMaterialfv(GL_FRONT_AND_BACK, GL_DIFFUSE,  red  )
  87     glMaterialfv(GL_FRONT_AND_BACK, GL_SPECULAR, white)
  88     glMaterialf( GL_FRONT_AND_BACK, GL_SHININESS, 60.0)
  89 
  90 def resize(Width, Height):
  91     global current_size
  92     current_size = Width, Height
  93     glViewport(0, 0, Width, Height)        # Reset The Current Viewport And Perspective Transformation
  94     glMatrixMode(GL_PROJECTION)
  95     glLoadIdentity()
  96     gluPerspective(60.0, Width/float(Height), 0.1, 10.0)
  97 
  98 def do_tick():
  99     global time_of_last_titleupdate, frame_counter, frames_per_second
 100     if ((time.clock () * 1000.0) - time_of_last_titleupdate >= 1000.):
 101         frames_per_second = frame_counter                   # Save The FPS
 102         frame_counter = 0  # Reset The FPS Counter
 103         szTitle = "%d FPS" % (frames_per_second )
 104         glutSetWindowTitle ( szTitle )
 105         time_of_last_titleupdate = time.clock () * 1000.0
 106     frame_counter += 1
 107 
 108 # The function called whenever a key is pressed. Note the use of Python tuples to pass in: (key, x, y)
 109 def keyPressed(*args):
 110     global animate, enable_cuda
 111     # If escape is pressed, kill everything.
 112     if args[0] == '\033':
 113         print 'Closing..'
 114         destroy_PBOs()
 115         destroy_texture()
 116         exit()
 117     elif args[0] == 'a':
 118         print 'toggling animation'
 119         animate = not animate
 120     elif args[0] == 'e':
 121         print 'toggling cuda'
 122         enable_cuda = not enable_cuda
 123 
 124 def idle():
 125     global heading, pitch, bank
 126     if animate:
 127         heading += 0.2
 128         pitch   += 0.6
 129         bank    += 1.0
 130 
 131     glutPostRedisplay()
 132 
 133 def display():
 134     try:
 135         render_scene()
 136         if enable_cuda:
 137             process_image()
 138             display_image()
 139         glutSwapBuffers()
 140     except:
 141         from traceback import print_exc
 142         print_exc()
 143         from os import _exit
 144         _exit(0)
 145 
 146 def process(width, height):
 147     """ Use PyCuda """
 148     grid_dimensions   = (width//16,height//16)
 149 
 150     source_mapping = pycuda_source_pbo.map()
 151     dest_mapping   = pycuda_dest_pbo.map()
 152 
 153     invert.prepared_call(grid_dimensions, (16, 16, 1),
 154             source_mapping.device_ptr(),
 155             dest_mapping.device_ptr())
 156 
 157     cuda_driver.Context.synchronize()
 158 
 159     source_mapping.unmap()
 160     dest_mapping.unmap()
 161 
 162 def process_image():
 163     """ copy image and process using CUDA """
 164     global pycuda_source_pbo,source_pbo,current_size, dest_pbo
 165     image_width, image_height = current_size
 166     assert source_pbo is not None
 167 
 168     # tell cuda we are going to get into these buffers
 169     pycuda_source_pbo.unregister()
 170 
 171     # activate destination buffer
 172     glBindBufferARB(GL_PIXEL_PACK_BUFFER_ARB, long(source_pbo))
 173 
 174     # read data into pbo. note: use BGRA format for optimal performance
 175     glReadPixels(
 176              0,                  #start x
 177              0,                  #start y
 178              image_width,        #end   x
 179              image_height,       #end   y
 180              GL_BGRA,            #format
 181              GL_UNSIGNED_BYTE,   #output type
 182              ctypes.c_void_p(0))
 183 
 184     pycuda_source_pbo = cuda_gl.BufferObject(long(source_pbo))
 185 
 186     # run the Cuda kernel
 187     process(image_width, image_height)
 188     # blit convolved texture onto the screen
 189     # download texture from PBO
 190     glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, long(dest_pbo))
 191     glBindTexture(GL_TEXTURE_2D, output_texture)
 192 
 193     glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0,
 194                     image_width, image_height,
 195                     GL_BGRA, GL_UNSIGNED_BYTE, ctypes.c_void_p(0))
 196 
 197 def display_image():
 198     """ render a screen sized quad """
 199     glDisable(GL_DEPTH_TEST)
 200     glDisable(GL_LIGHTING)
 201     glEnable(GL_TEXTURE_2D)
 202     glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_REPLACE)
 203     glMatrixMode(GL_PROJECTION)
 204     glPushMatrix()
 205     glLoadIdentity()
 206     glOrtho(-1.0, 1.0, -1.0, 1.0, -1.0, 1.0)
 207     glMatrixMode( GL_MODELVIEW)
 208     glLoadIdentity()
 209     glViewport(0, 0, current_size[0], current_size[1])
 210     glBegin(GL_QUADS)
 211     glTexCoord2f(0.0, 0.0)
 212     glVertex3f(-1.0, -1.0, 0.5)
 213     glTexCoord2f(1.0, 0.0)
 214     glVertex3f(1.0, -1.0, 0.5)
 215     glTexCoord2f(1.0, 1.0)
 216     glVertex3f(1.0, 1.0, 0.5)
 217     glTexCoord2f(0.0, 1.0)
 218     glVertex3f(-1.0, 1.0, 0.5)
 219     glEnd()
 220     glMatrixMode(GL_PROJECTION)
 221     glPopMatrix()
 222     glDisable(GL_TEXTURE_2D)
 223     glBindBuffer(GL_PIXEL_PACK_BUFFER_ARB, 0)
 224     glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, 0)
 225 
 226 
 227 def render_scene():
 228     glClear (GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)# Clear Screen And Depth Buffer
 229     glMatrixMode(GL_MODELVIEW)
 230     glLoadIdentity ()      # Reset The Modelview Matrix
 231     glTranslatef(0.0, 0.0, -3.0);
 232     glRotatef(heading, 1.0, 0.0, 0.0)
 233     glRotatef(pitch  , 0.0, 1.0, 0.0)
 234     glRotatef(bank   , 0.0, 0.0, 1.0)
 235     glViewport(0, 0, current_size[0],current_size[1])
 236     glEnable(GL_LIGHTING)
 237     glEnable(GL_DEPTH_TEST)
 238     glDepthFunc(GL_LESS)
 239     glutSolidTeapot(1.0)
 240     do_tick()#just for fps display..
 241     return True
 242 
 243 def main():
 244     global window, cuda_module, cuda_gl, cuda_driver, invert
 245     glutInit(sys.argv)
 246     glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_ALPHA | GLUT_DEPTH)
 247     glutInitWindowSize(*initial_size)
 248     glutInitWindowPosition(0, 0)
 249     window = glutCreateWindow("PyCuda GL Interop Example")
 250     glutDisplayFunc(display)
 251     glutIdleFunc(idle)
 252     glutReshapeFunc(resize)
 253     glutKeyboardFunc(keyPressed)
 254     glutSpecialFunc(keyPressed)
 255     init_gl()
 256 
 257     # create texture for blitting to screen
 258     create_texture(*initial_size)
 259 
 260     #setup pycuda gl interop
 261     import pycuda.gl.autoinit
 262     import pycuda.gl
 263     cuda_gl = pycuda.gl
 264     cuda_driver = pycuda.driver
 265 
 266     cuda_module = SourceModule("""
 267     __global__ void invert(unsigned char *source, unsigned char *dest)
 268     {
 269       int block_num        = blockIdx.x + blockIdx.y * gridDim.x;
 270       int thread_num       = threadIdx.y * blockDim.x + threadIdx.x;
 271       int threads_in_block = blockDim.x * blockDim.y;
 272       //Since the image is RGBA we multiply the index 4.
 273       //We'll only use the first 3 (RGB) channels though
 274       int idx              = 4 * (threads_in_block * block_num + thread_num);
 275       dest[idx  ] = 255 - source[idx  ];
 276       dest[idx+1] = 255 - source[idx+1];
 277       dest[idx+2] = 255 - source[idx+2];
 278     }
 279     """)
 280     invert = cuda_module.get_function("invert")
 281     # The argument "PP" indicates that the invert function will take two PBOs as arguments
 282     invert.prepare("PP")
 283 
 284     # create source and destination pixel buffer objects for processing
 285     create_PBOs(*initial_size)
 286 
 287     glutMainLoop()
 288 
 289 # Print message to console, and kick off the main to get it rolling.
 290 if __name__ == "__main__":
 291     print "Hit ESC key to quit, 'a' to toggle animation, and 'e' to toggle cuda"
 292     main()