diff options
Diffstat (limited to 'render.py')
| -rwxr-xr-x | render.py | 468 |
1 files changed, 468 insertions, 0 deletions
diff --git a/render.py b/render.py new file mode 100755 index 0000000..2a4e0d5 --- /dev/null +++ b/render.py @@ -0,0 +1,468 @@ +#!/usr/bin/env python +import os +import sys +import numpy as np +import pickle +import tempfile + +import pygame +from pygame.locals import * + +import src +from camera import * +from geometry import * +from transform import * +from optics import * + +import time + +from pycuda import autoinit +from pycuda.compiler import SourceModule +from pycuda import gpuarray +from pycuda.characterize import sizeof + +def to_geometry(viewable, bits): + if isinstance(viewable, Geometry): + geometry = viewable + geometry.build(bits) + elif isinstance(viewable, Solid): + geometry = Geometry() + geometry.add_solid(viewable) + geometry.build(bits) + elif isinstance(viewable, Mesh): + geometry = Geometry() + geometry.add_solid(Solid(viewable, vacuum, vacuum, surface=lambertian_surface)) + geometry.build(bits) + else: + raise Exception("can't convert %s to a geometry" % viewable) + + return geometry + +def box(lower_bound, 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, name='', bits=8, load_bvh=False, camera_view=None, make_movie=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 + """ + + geometry = to_geometry(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(to_geometry(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])]) + + 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) + #geometry.load(module, color=True) + geometry.load(module, color=False) + #cuda_raytrace = module.get_function('ray_trace') + cuda_build_rgb_lookup = module.get_function('build_rgb_lookup') + cuda_render = module.get_function('render') + cuda_rotate = module.get_function('rotate') + cuda_translate = module.get_function('translate') + + pygame.init() + size = width, height = 800, 600 + + init_rng = module.get_function('init_rng') + rng_states_gpu = cuda.mem_alloc(len(geometry.mesh.triangles)*sizeof('curandStateXORWOW', '#include <curand_kernel.h>')) + init_rng(np.int32(len(geometry.mesh.triangles)), rng_states_gpu, np.int32(0), np.int32(0), block=(64,1,1), grid=(len(geometry.mesh.triangles)//64+1,1)) + + diagonal = np.linalg.norm(upper_bound-lower_bound) + #source_position = [0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2] + source_position = [0, 0, upper_bound[2]+5.0] + + source_positions_gpu = gpuarray.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.float3) + source_positions_gpu.fill(gpuarray.vec.make_float3(*source_position)) + + source_directions = np.mean(geometry.mesh[:], axis=1) - source_position + source_directions_float3 = np.empty(source_positions_gpu.size, dtype=gpuarray.vec.float3) + source_directions_float3['x'] = source_directions[:,0] + source_directions_float3['y'] = source_directions[:,1] + source_directions_float3['z'] = source_directions[:,2] + source_directions_gpu = cuda.to_device(source_directions_float3) + + if (len(geometry.mesh.triangles) < source_positions_gpu.size): + print width*height + print source_positions_gpu.size + raise Exception('not enough rng_states') + + rgb_lookup_gpu = gpuarray.zeros(source_positions_gpu.size, dtype=gpuarray.vec.float3) + + max_steps = 10 + rgb_runs = 1000 + + cuda_build_rgb_lookup(np.int32(source_positions_gpu.size), rng_states_gpu, source_positions_gpu, source_directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), rgb_lookup_gpu, np.int32(rgb_runs), np.int32(max_steps), block=(64,1,1), grid=(source_positions_gpu.size//64+1,1)) + + rgb_lookup = rgb_lookup_gpu.get() + rgb_lookup['x'] /= rgb_runs + rgb_lookup['y'] /= rgb_runs + rgb_lookup['z'] /= rgb_runs + rgb_lookup_gpu = cuda.to_device(rgb_lookup) + + + print rgb_lookup + + + + + + + + + + + + + if camera_view is None: + 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 = camera.get_rays() + + 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] + + 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 = cuda.to_device(origins_float3) + directions_gpu = cuda.to_device(directions_float3) + + pixels = np.empty(width*height, dtype=np.int32) + pixels_gpu = cuda.to_device(pixels) + + nruns = 5 + + + + screen = pygame.display.set_mode(size) + pygame.display.set_caption(name) + + + + + + + nblocks = 64 + + gpu_kwargs = {'block': (nblocks,1,1), 'grid':(pixels.size/nblocks+1,1)} + + if make_movie: + tempdir = tempfile.mkdtemp() + + def screenshot(dir=''): + if name == '': + root, ext = 'screenshot', 'png' + else: + root, ext = name, 'png' + + filename = os.path.join(dir,'.'.join([root, ext])) + + i = 1 + while os.path.exists(filename): + filename = os.path.join(dir,'.'.join([root + str(i), ext])) + i += 1 + + pygame.image.save(screen, filename) + print 'image saved to %s' % filename + + def render(): + """Render the mesh and display to screen.""" + t0 = time.time() + cuda_render(np.int32(pixels.size), rng_states_gpu, origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), rgb_lookup_gpu, np.int32(nruns), pixels_gpu, np.int32(max_steps), **gpu_kwargs) + 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.display.flip() + + if make_movie: + screenshot(tempdir) + + render() + + 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.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs) + + point += v + + render() + + 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) + + point += v + + render() + + 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.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs) + + point += v + + render() + 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.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs) + + point = rotate(point, phi, n) + axis1 = rotate(axis1, phi, n) + axis2 = rotate(axis2, phi, n) + + render() + + 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: + try: + if current_layer+1 >= len(bvhg): + raise IndexError + + geometry = bvhg[current_layer+1] + current_layer += 1 + + geometry.load(module, color=True) + render() + except IndexError: + print 'no further layers to view' + + if event.key == K_PAGEDOWN: + try: + if current_layer-1 < 0: + raise IndexError + + geometry = bvhg[current_layer-1] + current_layer -= 1 + + geometry.load(module, color=True) + render() + 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: + screenshot() + + if event.type == KEYUP: + if event.key == K_LSHIFT or event.key == K_RSHIFT: + shift = False + + pygame.display.quit() + + if make_movie: + if name == '': + root, ext = 'movie', 'tgz' + else: + root, ext = name + '_movie', 'tgz' + + filename = '.'.join([root, ext]) + + i=1 + while os.path.exists(filename): + filename = '.'.join([root + str(i), ext]) + i += 1 + + from subprocess import call + call(['tar', '-czvf', filename, tempdir]) + + os.removedirs(tempdir) + +if __name__ == '__main__': + import optparse + + from stl import mesh_from_stl + + import solids + import detectors + + 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('-m', '--movie', action='store_true', dest='movie', + help='create a movie', default=False) + 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) + + if ext.lower() == '.stl': + view(mesh_from_stl(args[0]), root, options.bits, options.load_bvh, options.camera_view, options.movie) + else: + import inspect + + members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids)) + + if args[0] in members: + view(members[args[0]], args[0], options.bits, options.load_bvh, options.camera_view, options.movie) + else: + sys.exit("couldn't find object %s" % args[0]) |