diff options
| -rwxr-xr-x | view.py | 375 |
1 files changed, 375 insertions, 0 deletions
@@ -0,0 +1,375 @@ +#!/usr/bin/env python +import os +import sys +import time +import numpy as np +import inspect + +import pygame +from pygame.locals import * + +import src +from camera import get_rays +from geometry import Mesh, Solid, Geometry +from transform import rotate +from optics import * +from gpu import * + +#from pycuda import autoinit +from pycuda.compiler import SourceModule +from pycuda import gpuarray +import pycuda.driver as cuda + +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(obj) + elif isinstance(obj, Mesh): + geometry = Geometry() + geometry.add_solid(Solid(obj, vacuum, vacuum, surface=lambertian_surface, color=0x99ffffff)) + else: + 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, + lower_bound + (dx,0,0), + lower_bound + (dx,dy,0), + lower_bound + (0,dy,0), + lower_bound + (0,0,dz), + lower_bound + (dx,0,dz), + lower_bound + (dx,dy,dz), + lower_bound + (0,dy,dz)]) + + triangles = np.empty((12,3), dtype=np.int32) + # bottom + triangles[0] = (1,3,2) + triangles[1] = (1,4,3) + # top + triangles[2] = (5,7,6) + triangles[3] = (5,8,7) + # left + triangles[4] = (5,1,2) + triangles[5] = (5,2,6) + # right + triangles[6] = (3,4,8) + triangles[7] = (3,8,7) + # front + triangles[8] = (2,3,7) + triangles[9] = (2,7,6) + # back + triangles[10] = (1,5,8) + triangles[11] = (1,8,4) + + triangles -= 1 + + return Mesh(vertices, triangles) + +def bvh_mesh(geometry, layer): + lower_bounds = geometry.lower_bounds[geometry.layers == layer] + upper_bounds = geometry.upper_bounds[geometry.layers == layer] + + if len(lower_bounds) == 0 or len(upper_bounds) == 0: + raise Exception('no nodes at layer %i' % layer) + + mesh = box(lower_bounds[0], upper_bounds[0]) + + for lower_bound, upper_bound in zip(lower_bounds, upper_bounds)[1:]: + mesh += box(lower_bound, upper_bound) + + return mesh + +def view(viewable, size=(800,600), name='', bits=8, load_bvh=False): + """ + Render `viewable` in a pygame window. + + Movement: + - zoom: scroll the mouse wheel + - rotate: click and drag the mouse + - move: shift+click and drag the mouse + """ + + gpu = GPU() + + geometry = build(viewable, bits) + + if load_bvh: + print 'loading bvh...' + bvhg = [] + + bvhg.append(geometry) + + for layer in sorted(np.unique(geometry.layers)): + print 'building layer %i' % layer + bvhg.append(build(bvh_mesh(geometry, layer), bits)) + + lower_bound, upper_bound = geometry.mesh.get_bounds() + + scale = np.linalg.norm(upper_bound-lower_bound) + + #from pycuda import autoinit + + #print 'device %s' % autoinit.device.name() + + module = gpu.module + + #module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True)#, cache_dir=False) + #geometry.load(module) + gpu.load_geometry(geometry) + cuda_raytrace = module.get_function('ray_trace') + cuda_rotate = module.get_function('rotate') + cuda_translate = module.get_function('translate') + + pygame.init() + width, height = size + screen = pygame.display.set_mode(size) + pygame.display.set_caption(name) + + #camera = Camera(size) + + diagonal = np.linalg.norm(upper_bound-lower_bound) + + 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) + + #camera.position(point) + + origins, directions = get_rays(point, size) + + origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3)) + + directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3)) + + pixels_gpu = gpuarray.zeros(width*height, dtype=np.int32) + + nblocks = 64 + + def update(): + """Render the mesh and display to screen.""" + t0 = time.time() + cuda_raytrace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, 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 + + pygame.surfarray.blit_array(screen, pixels_gpu.get().reshape(size)) + pygame.display.flip() + + update() + + done = False + clicked = False + shift = False + + current_layer = 0 + + while not done: + for event in pygame.event.get(): + if event.type == MOUSEBUTTONDOWN: + if event.button == 4: + v = scale*np.cross(axis1,axis2)/10.0 + + 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 + + update() + + if event.button == 5: + v = -scale*np.cross(axis1,axis2)/10.0 + + 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 + + update() + + if event.button == 1: + clicked = True + mouse_position = pygame.mouse.get_rel() + + if event.type == MOUSEBUTTONUP: + if event.button == 1: + clicked = False + + if event.type == MOUSEMOTION and clicked: + movement = np.array(pygame.mouse.get_rel()) + + if (movement == 0).all(): + continue + + length = np.linalg.norm(movement) + + mouse_direction = movement[0]*axis1 + movement[1]*axis2 + mouse_direction /= np.linalg.norm(mouse_direction) + + if shift: + v = mouse_direction*scale*length/float(width) + + 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 + + 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_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_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) + + update() + + if event.type == KEYDOWN: + if event.key == K_LSHIFT or event.key == K_RSHIFT: + shift = True + + if event.key == K_ESCAPE: + done = True + break + + if event.key == K_PAGEUP and load_bvh: + try: + if current_layer+1 >= len(bvhg): + raise IndexError + + geometry = bvhg[current_layer+1] + current_layer += 1 + + gpu.load_geometry(geometry) + #geometry.load(module, color=True) + update() + except IndexError: + print 'no further layers to view' + + if event.key == K_PAGEDOWN and load_bvh: + try: + if current_layer-1 < 0: + raise IndexError + + geometry = bvhg[current_layer-1] + current_layer -= 1 + + gpu.load_geometry(geometry) + #geometry.load(module, color=True) + update() + except IndexError: + print 'no further layers to view' + + if event.key == K_F12: + 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__': + import optparse + + from stl import mesh_from_stl + + 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('-r', '--resolution', dest='resolution', + help='specify resolution', default='800,600') + options, args = parser.parse_args() + + if len(args) < 1: + sys.exit(parser.format_help()) + + 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() in ('.stl', '.bz2'): + view(mesh_from_stl(args[0]), size, root, options.bits, options.load_bvh) + else: + members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(scenes)) + + 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 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]) |