OpenGL screenshots using python

screenshotty.py

#! /usr/bin/env python
'''
	Simple example demonstrating how to take a screenshot
'''
import time
import sys
import os
import argparse

#OpenGL
import OpenGL 
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
from OpenGL.GLUT.special import *
from OpenGL.GL.shaders import *

from PIL import Image

#format loaders
sys.path.append("./formatloaders")
from fl_wavefront_obj import *

#Other
import screenshot_ini_loader


OpenGL.FORWARD_COMPATIBLE_ONLY = True
OpenGL.ERROR_ON_COPY = True 

# global object display list
g_dlObject = None


# global y-axis rotation (weeee, spinning)
g_Rotate = [0.0,0.0,0.0]
g_Translate = [0.0,0.0,-2.0]

# Screenshot variables
g_screenshot = 0
g_screens = None
g_take_screenshot = False
g_maintain_window = False

# A general OpenGL initialization function.  Sets all of the initial parameters. 
def InitGL(Width, Height):                # We call this right after our OpenGL window is created.
	glClearColor(0.0, 0.0, 0.0, 0.0)    # This Will Clear The Background Color To Black
	glClearDepth(1.0)                    # Enables Clearing Of The Depth Buffer
	glShadeModel(GL_SMOOTH)                # Enables Smooth Color Shading
	glEnable(GL_DEPTH_TEST)
	glEnable(GL_NORMALIZE)
	
	glDepthFunc(GL_LEQUAL)
	glHint(GL_PERSPECTIVE_CORRECTION_HINT,GL_NICEST)

	glMatrixMode(GL_PROJECTION)
	glLoadIdentity()                    # Reset The Projection Matrix
	
	InitLightModel()
	
	# Calculate The Aspect Ratio Of The Window
	gluPerspective(45.0, float(Width)/float(Height), 0.1, 100.0)

	glMatrixMode(GL_MODELVIEW)

# Init basic lighting scene, assumes data sits between -1,1 on all axes.
def InitLightModel():
	cAmbientLight = GLfloat_4(0.5,0.5,0.5,1.0)
	glLightfv(GL_LIGHT1,GL_AMBIENT,cAmbientLight)
	
	cDiffuseLight = GLfloat_4(0.5,0.5,0.5,1.0)
	glLightfv(GL_LIGHT1,GL_DIFFUSE,cDiffuseLight)
	
	vLightPos = GLfloat_4(1,3,4,1)
	glLightfv(GL_LIGHT1, GL_POSITION,vLightPos );
	glEnable(GL_LIGHT1)
	
	glEnable(GL_LIGHTING)
	
def InitScreenShots():
	global g_screens,g_take_screenshot
	g_screens = screenshot_ini_loader.load("screenshotty.ini")
	
	if g_screens == None:
		print "Could not initialize screenshot array due to errors in ini loader, fix and try again"
		#sys.exit(1)
	else:	
		g_take_screenshot = True
	
# Generate debug display list using 'teapot'
def dbgInitDisplayList():
	global g_dlObject
	
	g_dlObject = glGenLists(1)
	glNewList(g_dlObject,GL_COMPILE)
	glutSolidTeapot(1.0)
	glEndList()
	
	
# Generate display list for any obj
def InitDisplayList(filename):
	global g_dlObject,g_take_screenshot

	obj = WaveFront()
	obj.NormalizeVertices = True
	obj.load(filename)
	
	g_dlObject = glGenLists(1)
	glNewList(g_dlObject,GL_COMPILE)
	glBegin(GL_TRIANGLES)
	
	ct = 0
	progress = ""
	if obj.Colors != None:
		print "Object has texture, attempting render"
		
	for face in obj.Faces:
		if ct % (len(obj.Faces)/20) == 0:
			progress += "="
			space = ""
			for i in range(len(progress),20):
				space += " "
		
			sys.stdout.write("Initializing Normals  0[" + progress + ">" + space + "]100\r");
		
		ct+= 1
		
		
		v1 = [obj.Vertices[face[0]-1][0],obj.Vertices[face[0]-1][1],obj.Vertices[face[0]-1][2]]
		v2 = [obj.Vertices[face[1]-1][0],obj.Vertices[face[1]-1][1],obj.Vertices[face[1]-1][2]]
		v3 = [obj.Vertices[face[2]-1][0],obj.Vertices[face[2]-1][1],obj.Vertices[face[2]-1][2]]
		
		if obj.Colors != None:
			c1 = [obj.Colors[face[0]-1][0],obj.Colors[face[0]-1][1],obj.Colors[face[0]-1][2]]
			c2 = [obj.Colors[face[1]-1][0],obj.Colors[face[1]-1][1],obj.Colors[face[1]-1][2]]
			c3 = [obj.Colors[face[2]-1][0],obj.Colors[face[2]-1][1],obj.Colors[face[2]-1][2]]
		if len(obj.Normals) != 0:
			n = [obj.Normals[face[0]-1][0],obj.Normals[face[0]-1][1],obj.Normals[face[0]-1][2]]
		else:
			# calculate normals on the fly (slow as hell on any decently size models.)
			e1 = [v2[0] - v1[0],v2[1] - v1[1],v2[2] - v1[2]]
			e2 = [v3[0] - v1[0],v3[1] - v1[1],v3[2] - v1[2]]
			n = cross(e1,e2)
			
			length = sqrt(n[0]*n[0] + n[1] * n[1] + n[2] * n[2])
			if length > 0.0:
				n[0] = n[0] / length
				n[1] = n[1] / length
				n[2] = n[2] / length


		glNormal3f(n[0],n[1],n[2])
		if obj.Colors != None:
			glColor3f(c1[0],c1[1],c1[2])
		glVertex3f(v1[0],v1[1],v1[2])
		
		if obj.Colors != None:
			glColor3f(c2[0],c2[1],c2[2])
		glVertex3f(v2[0],v2[1],v2[2])
		
		if obj.Colors != None:
			glColor3f(c3[0],c3[1],c3[2])
		glVertex3f(v3[0],v3[1],v3[2])
		
	glEnd()
	glEndList()
	
	print "\nObj loaded"
	
# The function called when our window is resized (which shouldn't happen if you enable fullscreen, below)
def ReSizeGLScene(Width, Height):
    if Height == 0:                        # Prevent A Divide By Zero If The Window Is Too Small 
        Height = 1

	glViewport(0, 0, Width, Height)        # Reset The Current Viewport And Perspective Transformation
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(45.0, float(Width)/float(Height), 0.1, 100.0)
    glMatrixMode(GL_MODELVIEW)


#dump back buffer to image
def ScreenShot(filename):
	width = 800
	height = 600
	glReadBuffer(GL_FRONT)
	pixels = glReadPixels(0,0,width,height,GL_RGB,GL_UNSIGNED_BYTE)
	
	image = Image.fromstring("RGB", (width, height), pixels)
	image = image.transpose( Image.FLIP_TOP_BOTTOM)
	image.save(filename)


g_screenshot = 0
g_screenshot_name = ["screenx.jpg","screeny.jpg","screenz.jpg"]
g_screenshot_angle = [[90,1,0,0],[90,0,1,0],[90,0,0,1]]

# The main drawing function. 
def DrawGLScene():
	global g_Rotate,g_Translate,g_take_screenshot,g_screenshot
	glColorMaterial ( GL_FRONT, GL_AMBIENT )
	glEnable ( GL_COLOR_MATERIAL )

	glDrawBuffer(GL_BACK)
	glMatrixMode(GL_MODELVIEW)
	glLoadIdentity()

	# Clear The Screen And The Depth Buffer
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
	#glDisable(GL_LIGHTING)
	
	glTranslatef(g_Translate[0],g_Translate[1],g_Translate[2])
	#glMultMatrixf(g_Transform)
	
	
	if g_take_screenshot:
		print "Obj was loaded, getting first file:",g_screens[g_screenshot][0]
		glMultMatrixf(g_screens[g_screenshot][1])
		#anglev = g_screenshot_angle[g_screenshot]
		
		#glRotatef(90,0,0,1)
		# Draw list
		glCallList(g_dlObject)
		
		# Swap buffers
		glutSwapBuffers()
		glFlush()
		ScreenShot(g_screens[g_screenshot][0])
		g_screenshot += 1
		
		if g_screenshot >= len(g_screens):
			if not g_maintain_window:
				sys.exit(1)
			g_take_screenshot=False
	else:
		if g_maintain_window:
			glCallList(g_dlObject)
			
			# Swap buffers
			glutSwapBuffers()

# The function called whenever a key is pressed. Note the use of Python tuples to pass in: (key, x, y)  
def keyPressed(*args):
	# If escape is pressed, kill everything.
	if args[0] == '\x1b':
		sys.exit()

def windowClose(*args):
	print "Window closing"
	return
	
def main(args):
	global window
	global falloffValue
	global g_take_screenshot,g_maintain_window
	
	# Force hide window to fake "background" rendering
	if not args.unhide:
		glutHideWindow()
	else:
		g_maintain_window = True
		
	glutInit()

	# Select type of Display mode:   
	#  Double buffer 
	#  RGBA color
	# Alpha components supported 
	# Depth buffer
	glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)

	# get a 800 x 600 window 
	glutInitWindowSize(800, 600)

	# the window starts at the upper left corner of the screen 
	glutInitWindowPosition(0, 0)

	# Okay, like the C version we retain the window id to use when closing, but for those of you new
	# to Python (like myself), remember this assignment would make the variable local and not global
	# if it weren't for the global declaration at the start of main.
	window = glutCreateWindow("Display List Test")

	# Register the drawing function with glut, BUT in Python land, at least using PyOpenGL, we need to
	# set the function pointer and invoke a function to actually register the callback, otherwise it
	# would be very much like the C version of the code.    
	glutDisplayFunc(DrawGLScene)

	# Uncomment this line to get full screen.
	#glutFullScreen()

	# When we are doing nothing, redraw the scene.
	glutIdleFunc(DrawGLScene)

	# Register the function called when our window is resized.
	glutReshapeFunc(ReSizeGLScene)
	
	# Load screenshot array
	InitScreenShots()
	
	# Load display list w/ model
	InitDisplayList(args.filename)
		
	# Initialize our window. 
	InitGL(800, 600)
	
	
	
	# Start Event Processing Engine    
	glutMainLoop()


if __name__ == "__main__":
	parser = argparse.ArgumentParser(description=
		"Take screenshots* of a Wavefront .obj file\n"
		"\n"
		"*Screenshots must be defined in screenshots.ini.\n"
		"This file must reside at the same level as\n"
		"screenshotty.py (See file for examples)",
		formatter_class=argparse.RawTextHelpFormatter)
	parser.add_argument('filename', metavar='<obj file>', type=str,
                   help='Wavefront .obj file to load')
	parser.add_argument('-u','--unhide',action="store_true",help="Unhide background rendering window")
	args = parser.parse_args()

	print "Hit ESC key to quit."
	main(args)