diff options
Diffstat (limited to 'camera.py')
| -rw-r--r-- | camera.py | 390 |
1 files changed, 373 insertions, 17 deletions
@@ -1,11 +1,25 @@ import numpy as np -from itertools import product +from itertools import product, count +from threading import Thread, Lock +import time +import os -class Camera(object): +from transform import rotate + +import pygame +from pygame.locals import * + +from pycuda import gpuarray +from pycuda.characterize import sizeof +import pycuda.driver as cuda + +def get_rays(position, size = (800, 600), film_size = (0.035, 0.024), focal_length=0.05): """ - Pinhole camera object. + Generate ray positions and directions from a pinhole camera facing the negative y direction. Args: + - position: tuple, + Position of the camera. - size: tuple, *optional* Pixel array shape. - film_size: tuple, *optional* @@ -13,21 +27,363 @@ class Camera(object): - focal_length: float, *optional* Focal length of camera. """ - def __init__(self, size = (800, 600), film_size = (0.035, 0.024), \ - focal_length=0.05): - x = np.linspace(-film_size[0]/2, film_size[0]/2, size[0]) - z = np.linspace(-film_size[1]/2, film_size[1]/2, size[1]) - self.grid = np.array(tuple(product(x,[0],z))) + x = np.linspace(-film_size[0]/2, film_size[0]/2, size[0]) + z = np.linspace(-film_size[1]/2, film_size[1]/2, size[1]) + + grid = np.array(tuple(product(x,[0],z))) + + grid += (0,focal_length,0) + focal_point = np.zeros(3) + + grid += position + focal_point += position + + return grid, focal_point-grid + +class Camera(Thread): + def __init__(self, geometry, module, context, lock, size=(800,600)): + Thread.__init__(self) + self.geometry = geometry + self.module = module + self.context = context + self.lock = lock + + self.size = size + self.width, self.height = size + + pygame.init() + self.screen = pygame.display.set_mode(size) + pygame.display.set_caption('') + + with self.lock: + self.context.push() + self.ray_trace_kernel = self.module.get_function('ray_trace') + self.rotate_kernel = self.module.get_function('rotate') + self.translate_kernel = self.module.get_function('translate') + self.build_rgb_lookup_kernel = self.module.get_function('build_rgb_lookup') + self.render_kernel = self.module.get_function('render') + self.init_rng_kernel = self.module.get_function('init_rng') + self.context.pop() + + lower_bound, upper_bound = self.geometry.mesh.get_bounds() + self.scale = np.linalg.norm(upper_bound-lower_bound) + self.source_position = [0,0,upper_bound[2]+1.0] + + self.nblocks = 64 + + with self.lock: + self.context.push() + print 'initializing random number states.' + rng_state_count = max(len(geometry.mesh.triangles), self.width*self.height) + self.rng_states_gpu = cuda.mem_alloc(rng_state_count*sizeof('curandStateXORWOW', '#include <curand_kernel.h>')) + self.init_rng_kernel(np.int32(rng_state_count), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(rng_state_count//self.nblocks+1,1)) + self.context.synchronize() + print 'done.' + + self.source_positions_gpu = gpuarray.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.float3) + self.source_positions_gpu.fill(gpuarray.vec.make_float3(*self.source_position)) + + source_directions = np.mean(self.geometry.mesh[:], axis=1) - self.source_position + self.source_directions_gpu = gpuarray.to_gpu(source_directions.astype(np.float32).view(gpuarray.vec.float3)) + + self.rgb_lookup1_gpu = gpuarray.zeros(self.source_positions_gpu.size, dtype=gpuarray.vec.float3) + self.rgb_lookup2_gpu = gpuarray.zeros(self.source_positions_gpu.size, dtype=gpuarray.vec.float3) + + self.max_steps = 10 + rgb_runs = 100 + + print 'building rgb lookup.' + self.build_rgb_lookup_kernel(np.int32(self.source_positions_gpu.size), self.rng_states_gpu, self.source_positions_gpu, self.source_directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(rgb_runs), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.source_positions_gpu.size//self.nblocks+1,1)) + self.context.synchronize() + print 'done.' + + rgb_lookup1 = self.rgb_lookup1_gpu.get().view(np.float32) + rgb_lookup1 /= rgb_runs + rgb_lookup1[rgb_lookup1 > 1.0] = 1.0 + self.rgb_lookup1_gpu.set(rgb_lookup1.view(gpuarray.vec.float3)) + + rgb_lookup2 = self.rgb_lookup2_gpu.get().view(np.float32) + rgb_lookup2 /= rgb_runs + rgb_lookup2[rgb_lookup2 > 1.0] = 1.0 + self.rgb_lookup2_gpu.set(rgb_lookup2.view(gpuarray.vec.float3)) + + self.nruns = 1 + + self.context.pop() + + self.point = np.array([0, self.scale*1.75, (lower_bound[2]+upper_bound[2])/2]) + self.axis1 = np.array([1,0,0], float) + self.axis2 = np.array([0,0,1], float) + + origins, directions = get_rays(self.point) + + with self.lock: + self.context.push() + self.origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3)) + self.directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3)) + self.pixels_gpu = gpuarray.zeros(self.width*self.height, dtype=np.int32) + self.context.pop() + + self.render = False + + self.render_run_count = 1 + + def screenshot(self, dir=''): + root, ext = 'screenshot', 'png' + + for i in count(): + filename = os.path.join(dir, '.'.join([root + str(i).zfill(4), ext])) + + if not os.path.exists(filename): + break + + pygame.image.save(self.screen, filename) + print 'image saved to %s' % filename + + def rotate(self, phi, n): + with self.lock: + self.context.push() + self.rotate_kernel(np.int32(self.pixels_gpu.size), self.origins_gpu, phi, gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + self.rotate_kernel(np.int32(self.pixels_gpu.size), self.directions_gpu, phi, gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + + self.point = rotate(self.point, phi, n) + self.axis1 = rotate(self.axis1, phi, n) + self.axis2 = rotate(self.axis2, phi, n) + self.context.pop() + + def translate(self, v): + with self.lock: + self.context.push() + self.translate_kernel(np.int32(self.pixels_gpu.size), self.origins_gpu, gpuarray.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks,1)) + + self.point += v + self.context.pop() + + def update(self): + with self.lock: + self.context.push() + t0 = time.time() + if self.render: + self.render_kernel(np.int32(self.pixels_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(self.render_run_count), self.pixels_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + else: + self.ray_trace_kernel(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + self.context.synchronize() + elapsed = time.time() - t0 + + print 'elapsed %f sec.' % elapsed + + pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size)) + pygame.display.flip() + self.context.pop() + + def run(self): + self.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 = self.scale*np.cross(self.axis1,self.axis2)/10.0 + + self.translate(v) + + self.update() + + if event.button == 5: + v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 + + self.translate(v) + + self.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]*self.axis1 + movement[1]*self.axis2 + mouse_direction /= np.linalg.norm(mouse_direction) + + if shift: + v = mouse_direction*self.scale*length/float(self.width) + + self.translate(v) + + self.update() + else: + phi = np.float32(2*np.pi*length/float(self.width)) + n = rotate(mouse_direction, np.pi/2, -np.cross(self.axis1,self.axis2)) + + self.rotate(phi, n) + + self.update() + + if event.type == KEYDOWN: + if event.key == K_a: + v = self.scale*self.axis1/10.0 + self.translate(v) + self.update() + + if event.key == K_d: + v = -self.scale*self.axis1/10.0 + self.translate(v) + self.update() + + if event.key == K_w: + v = self.scale*np.cross(self.axis1,self.axis2)/10.0 + self.translate(v) + self.update() + + if event.key == K_s: + v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 + self.translate(v) + self.update() + + if event.key == K_SPACE: + v = self.scale*self.axis2/10.0 + self.translate(v) + self.update() + + if event.key == K_LCTRL: + v = -self.scale*self.axis2/10.0 + self.translate(v) + self.update() + + 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 + + # 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 + + # geometry.load(module, color=True) + # update() + # except IndexError: + # print 'no further layers to view' + + if event.key == K_F12: + self.screenshot() + + if event.key == K_F5: + self.render = not self.render + self.update() + + if event.key == K_EQUALS: + self.render_run_count += 1 + self.update() + + if event.key == K_MINUS: + if self.render_run_count > 1: + self.render_run_count -= 1 + self.update() + + 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 + import inspect + + import solids + import detectors + import scenes + from stl import mesh_from_stl + import src + from view import build + + from pycuda.compiler import SourceModule + + 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]): + root, ext = os.path.splitext(os.path.split(args[0])[1]) + + if ext.lower() == '.stl': + obj = mesh_from_stl(args[0]) + 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: + obj = buildable_lookup[args[0]] + else: + raise Exception("can't find object %s" % args[0]) + + from pycuda.tools import make_default_context + + cuda.init() + context = make_default_context() + + module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False) + geometry = build(obj, options.bits) + geometry.load(module) - self.grid += (0,focal_length,0) - self.focal_point = np.zeros(3) + lock = Lock() - def position(self, position): - """Translate the camera to `position`.""" - self.grid += position - self.focal_point += position + camera = Camera(geometry, module, context, lock, size) - def get_rays(self): - """Return the position and direction for each pixel ray.""" - return self.grid, self.focal_point-self.grid + context.pop() + camera.start() |