diff options
Diffstat (limited to 'view.py')
| -rwxr-xr-x | view.py | 248 |
1 files changed, 120 insertions, 128 deletions
@@ -1,42 +1,87 @@ #!/usr/bin/env python import os import sys +import time import numpy as np -import pickle +import inspect import pygame from pygame.locals import * import src -from camera import * -from geometry import * -from transform import * -#from bvh_view import * - -import time +from camera import Camera +from geometry import Mesh, Solid, Geometry +from transform import rotate from pycuda import autoinit from pycuda.compiler import SourceModule from pycuda import gpuarray +import pycuda.driver as cuda -def to_geometry(viewable, bits): - if isinstance(viewable, Geometry): - geometry = viewable - geometry.build(bits) - elif isinstance(viewable, Solid): +def buildable(identifier): + """ + Create a decorator which tags a function as buildable and assigns the + identifying string `identifier`. + + Example: + >>> @buildable('my_sphere') + >>> def build_my_sphere(): + >>> g = Geometry() + >>> g.add_solid(Solid(sphere(), vacuum, water)) + >>> return g + """ + def tag_as_buildable(func): + func.buildable = True + func.identifier = identifier + return func + return tag_as_buildable + +def screenshot(screen, name='', dir='', index=0): + """Take a screenshot of `screen`.""" + if name == '': + root, ext = 'screenshot', 'png' + else: + root, ext = name, 'png' + + i = index + + filename = os.path.join(dir,'.'.join([root + str(i).zfill(4), ext])) + while os.path.exists(filename): + filename = os.path.join(dir,'.'.join([root + str(i).zfill(4), ext])) + i += 1 + + pygame.image.save(screen, filename) + print 'image saved to %s' % filename + +def build(obj, bits): + """Construct and build a geometry from `obj`.""" + if inspect.isfunction(obj): + try: + if obj.buildable: + obj = obj() + except AttributeError: + raise Exception('function %s is not buildable.' % obj.__name__) + + if isinstance(obj, Geometry): + geometry = obj + elif isinstance(obj, Solid): geometry = Geometry() - geometry.add_solid(viewable) - geometry.build(bits) - elif isinstance(viewable, Mesh): + geometry.add_solid(obj) + elif isinstance(obj, Mesh): geometry = Geometry() - geometry.add_solid(Solid(viewable)) - geometry.build(bits) + geometry.add_solid(Solid(obj)) else: - raise Exception("can't convert %s to a geometry" % viewable) + raise Exception('cannot build type %s' % type(obj)) + + geometry.build(bits) return geometry def box(lower_bound, upper_bound): + """ + Return a mesh of the box defined by the opposing corners `lower_bound` + and `upper_bound`. + """ dx, dy, dz = upper_bound - lower_bound vertices = np.array([lower_bound, @@ -86,7 +131,7 @@ def bvh_mesh(geometry, layer): return mesh -def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): +def view(viewable, size=(800,600), name='', bits=8, load_bvh=False): """ Render `viewable` in a pygame window. @@ -96,7 +141,7 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): - move: shift+click and drag the mouse """ - geometry = to_geometry(viewable, bits) + geometry = build(viewable, bits) if load_bvh: print 'loading bvh...' @@ -106,87 +151,60 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): for layer in sorted(np.unique(geometry.layers)): print 'building layer %i' % layer - bvhg.append(to_geometry(bvh_mesh(geometry, layer), bits)) + bvhg.append(build(bvh_mesh(geometry, layer), bits)) - lower_bound = np.array([np.min(geometry.mesh[:][:,:,0]), - np.min(geometry.mesh[:][:,:,1]), - np.min(geometry.mesh[:][:,:,2])]) - - upper_bound = np.array([np.max(geometry.mesh[:][:,:,0]), - np.max(geometry.mesh[:][:,:,1]), - np.max(geometry.mesh[:][:,:,2])]) + lower_bound, upper_bound = geometry.mesh.get_bounds() scale = np.linalg.norm(upper_bound-lower_bound) print 'device %s' % autoinit.device.name() - module = SourceModule(src.kernel, options=['-I' + src.dir], - no_extern_c=True, cache_dir=False) + module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False) geometry.load(module, color=True) cuda_raytrace = module.get_function('ray_trace') cuda_rotate = module.get_function('rotate') cuda_translate = module.get_function('translate') pygame.init() - size = width, height = 800, 600 + width, height = size screen = pygame.display.set_mode(size) pygame.display.set_caption(name) - if camera_view is None: - camera = Camera(size) - - diagonal = np.linalg.norm(upper_bound-lower_bound) + camera = Camera(size) - point = np.array([0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2]) - axis1 = np.array([1,0,0], dtype=np.double) - axis2 = np.array([0,0,1], dtype=np.double) + diagonal = np.linalg.norm(upper_bound-lower_bound) - camera.position(point) + point = np.array([0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2]) + axis1 = np.array([1,0,0], dtype=np.double) + axis2 = np.array([0,0,1], dtype=np.double) - origins, directions = camera.get_rays() + camera.position(point) - origins_float3 = np.empty(origins.shape[0], dtype=gpuarray.vec.float3) - origins_float3['x'] = origins[:,0] - origins_float3['y'] = origins[:,1] - origins_float3['z'] = origins[:,2] + origins, directions = camera.get_rays() - directions_float3 = np.empty(directions.shape[0], dtype=gpuarray.vec.float3) - directions_float3['x'] = directions[:,0] - directions_float3['y'] = directions[:,1] - directions_float3['z'] = directions[:,2] - else: - f = open(camera_view, 'rb') - origins_float3 = pickle.load(f) - directions_float3 = pickle.load(f) - point = pickle.load(f) - axis1 = pickle.load(f) - axis2 = pickle.load(f) - f.close() + origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3)) - origins_gpu = cuda.to_device(origins_float3) - directions_gpu = cuda.to_device(directions_float3) + directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3)) - pixels = np.empty(width*height, dtype=np.int32) - pixels_gpu = cuda.to_device(pixels) + pixels_gpu = gpuarray.zeros(width*height, dtype=np.int32) nblocks = 64 - gpu_kwargs = {'block': (nblocks,1,1), 'grid':(pixels.size/nblocks+1,1)} + gpu_kwargs = {'block': (nblocks,1,1), 'grid':(pixels_gpu.size/nblocks+1,1)} - def render(): + def update(): """Render the mesh and display to screen.""" t0 = time.time() - cuda_raytrace(np.int32(pixels.size), origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), pixels_gpu, **gpu_kwargs) + cuda_raytrace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), pixels_gpu, block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) cuda.Context.synchronize() elapsed = time.time() - t0 #print 'elapsed %f sec' % elapsed - cuda.memcpy_dtoh(pixels, pixels_gpu) - pygame.surfarray.blit_array(screen, pixels.reshape(size)) + pygame.surfarray.blit_array(screen, pixels_gpu.get().reshape(size)) pygame.display.flip() - render() + update() done = False clicked = False @@ -200,20 +218,20 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): if event.button == 4: v = scale*np.cross(axis1,axis2)/10.0 - cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs) + cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) point += v - render() + update() if event.button == 5: v = -scale*np.cross(axis1,axis2)/10.0 - cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs) + cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) point += v - render() + update() if event.button == 1: clicked = True @@ -237,25 +255,25 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): if shift: v = mouse_direction*scale*length/float(width) - cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs) + cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) point += v - render() + update() else: phi = np.float32(2*np.pi*length/float(width)) n = rotate(mouse_direction, np.pi/2, \ -np.cross(axis1,axis2)) - cuda_rotate(np.int32(pixels.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs) + cuda_rotate(np.int32(pixels_gpu.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - cuda_rotate(np.int32(pixels.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs) + cuda_rotate(np.int32(pixels_gpu.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) point = rotate(point, phi, n) axis1 = rotate(axis1, phi, n) axis2 = rotate(axis2, phi, n) - render() + update() if event.type == KEYDOWN: if event.key == K_LSHIFT or event.key == K_RSHIFT: @@ -265,7 +283,7 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): done = True break - if event.key == K_PAGEUP: + if event.key == K_PAGEUP and load_bvh: try: if current_layer+1 >= len(bvhg): raise IndexError @@ -274,11 +292,11 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): current_layer += 1 geometry.load(module, color=True) - render() + update() except IndexError: print 'no further layers to view' - if event.key == K_PAGEDOWN: + if event.key == K_PAGEDOWN and load_bvh: try: if current_layer-1 < 0: raise IndexError @@ -287,55 +305,21 @@ def view(viewable, name='', bits=8, load_bvh=False, camera_view=None): current_layer -= 1 geometry.load(module, color=True) - render() + update() except IndexError: print 'no further layers to view' - if event.key == K_F11: - if name == '': - root, ext = 'camera_view', 'pkl' - else: - root, ext = name + '_camera_view', 'pkl' - - filename = '.'.join([root, ext]) - - i = 1 - while os.path.exists(filename): - filename = '.'.join([root + str(i), ext]) - i += 1 - - f = open(filename, 'wb') - cuda.memcpy_dtoh(origins_float3, origins_gpu) - cuda.memcpy_dtoh(directions_float3, directions_gpu) - pickle.dump(origins_float3, f) - pickle.dump(directions_float3, f) - pickle.dump(point, f) - pickle.dump(axis1, f) - pickle.dump(axis2, f) - f.close() - - print 'camera view saved to %s' % filename - if event.key == K_F12: - if name == '': - root, ext = 'screenshot', 'png' - else: - root, ext = name, 'png' - - filename = '.'.join([root, ext]) - - i = 1 - while os.path.exists(filename): - filename = '.'.join([root + str(i), ext]) - i += 1 - - pygame.image.save(screen, filename) - print 'image saved to %s' % filename + screenshot(screen, name) if event.type == KEYUP: if event.key == K_LSHIFT or event.key == K_RSHIFT: shift = False + if event.type == pygame.QUIT: + done = True + break + pygame.display.quit() if __name__ == '__main__': @@ -345,30 +329,38 @@ if __name__ == '__main__': import solids import detectors + import scenes parser = optparse.OptionParser('%prog filename.stl') parser.add_option('-b', '--bits', type='int', dest='bits', help='bits for z-ordering space axes', default=8) parser.add_option('-l', action='store_true', dest='load_bvh', help='load bounding volumes', default=False) - parser.add_option('-v', '--camera_view', dest='camera_view', - help='load a camera view', default=None) + parser.add_option('-r', '--resolution', dest='resolution', + help='specify resolution', default='800,600') options, args = parser.parse_args() if len(args) < 1: sys.exit(parser.format_help()) - head, tail = os.path.split(args[0]) - root, ext = os.path.splitext(tail) + size = [int(s) for s in options.resolution.split(',')] + + if os.path.exists(args[0]): + head, tail = os.path.split(args[0]) + root, ext = os.path.splitext(tail) - if ext.lower() == '.stl': - view(mesh_from_stl(args[0]), root, options.bits, options.load_bvh, options.camera_view) + if ext.lower() == '.stl': + view(mesh_from_stl(args[0]), size, root, options.bits, options.load_bvh) else: - import inspect + members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(scenes)) - members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids)) + buildable_lookup = {} + for member in members.values(): + if inspect.isfunction(member) and \ + hasattr(member, 'buildable') and member.buildable == True: + buildable_lookup[member.identifier] = member - if args[0] in members: - view(members[args[0]], args[0], options.bits, options.load_bvh, options.camera_view) + if args[0] in buildable_lookup: + view(buildable_lookup[args[0]], size, args[0], options.bits, options.load_bvh) else: sys.exit("couldn't find object %s" % args[0]) |