diff options
Diffstat (limited to 'camera.py')
| -rw-r--r-- | camera.py | 175 |
1 files changed, 88 insertions, 87 deletions
@@ -2,8 +2,10 @@ import numpy as np from itertools import product, count from threading import Thread, Lock import time +import datetime import os import sys +from color import map_wavelength from transform import rotate @@ -14,6 +16,14 @@ from pycuda import gpuarray from pycuda.characterize import sizeof import pycuda.driver as cuda +def timeit(func): + def f(*args, **kwargs): + t0 = time.time() + func(*args, **kwargs) + elapsed = time.time() - t0 + print '%s elapsed in %s().' % (datetime.timedelta(seconds=elapsed), func.__name__) + return f + def get_rays(position, size = (800, 600), film_size = (0.035, 0.024), focal_length=0.05): """ Generate ray positions and directions from a pinhole camera facing the negative y direction. @@ -63,8 +73,9 @@ class Camera(Thread): 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.update_xyz_lookup_kernel = self.module.get_function('update_xyz_lookup') + self.update_xyz_image_kernel = self.module.get_function('update_xyz_image') + self.process_image_kernel = self.module.get_function('process_image') self.init_rng_kernel = self.module.get_function('init_rng') self.context.pop() @@ -77,7 +88,7 @@ class Camera(Thread): self.axis1 = np.array([1,0,0], float) self.axis2 = np.array([0,0,1], float) - origins, directions = get_rays(self.point) + origins, directions = get_rays(self.point, self.size) with self.lock: self.context.push() @@ -88,77 +99,73 @@ class Camera(Thread): self.render = False - def build_rgb_lookup(self, source_position=(0,0,0)): + @timeit + def initialize_render(self): with self.lock: self.context.push() - print '\ninitializing random number states.' self.rng_states_gpu = cuda.mem_alloc(self.width*self.height*sizeof('curandStateXORWOW', '#include <curand_kernel.h>')) self.init_rng_kernel(np.int32(self.width*self.height), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) + self.xyz_lookup1_gpu = gpuarray.zeros(len(self.geometry.mesh.triangles), dtype=gpuarray.vec.float3) + self.xyz_lookup2_gpu = gpuarray.zeros(len(self.geometry.mesh.triangles), dtype=gpuarray.vec.float3) + self.image_gpu = gpuarray.zeros(self.width*self.height, dtype=gpuarray.vec.float3) 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(*source_position)) + self.context.pop() - source_directions = np.mean(self.geometry.mesh[:], axis=1) - source_position - self.source_directions_gpu = gpuarray.to_gpu(source_directions.astype(np.float32).view(gpuarray.vec.float3)) + self.source_position = self.point - 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.nimages = 0 + self.nlookup_calls = 0 + self.max_steps = 10 - self.max_steps = 10 - rgb_runs = 100 - - print 'building rgb lookup.' - self.build_rgb_lookup_kernel(np.int32(self.source_positions_gpu.size), self.source_positions_gpu, self.source_directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.uint64(np.random.randint(np.iinfo(np.int64).max)), 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.' + def clear_xyz_lookup(self): + with self.lock: + self.context.push() + self.xyz_lookup1_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0)) + self.xyz_lookup2_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0)) + self.context.pop() - 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)) + self.nlookup_calls = 0 - 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)) + def update_xyz_lookup(self, source_position): + with self.lock: + self.context.push() + for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1): + self.update_xyz_lookup_kernel(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(685.0), gpuarray.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) - self.nimages = 0 - self.rgb_image = np.zeros((self.pixels_gpu.size, 3), float) - self.r_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32) - self.g_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32) - self.b_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32) + for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1): + self.update_xyz_lookup_kernel(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(545.0), gpuarray.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) + for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1): + self.update_xyz_lookup_kernel(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(445.0), gpuarray.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) self.context.pop() + self.nlookup_calls += 1 + def clear_image(self): - try: - self.rgb_image.fill(0.0) - self.nimages = 0 - except AttributeError: - pass + with self.lock: + self.context.push() + self.image_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0)) + self.context.pop() - def acquire_image(self): - if not hasattr(self, 'rng_states_gpu'): - self.build_rgb_lookup() + self.nimages = 0 + def update_image(self): with self.lock: self.context.push() - self.render_kernel(np.int32(self.pixels_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, self.r_gpu, self.g_gpu, self.b_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) - self.rgb_image[:,0] += self.r_gpu.get() - self.rgb_image[:,1] += self.g_gpu.get() - self.rgb_image[:,2] += self.b_gpu.get() - self.nimages += 1 - self.context.pop() + self.update_xyz_image_kernel(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(685.0), gpuarray.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) - def render_image(self): - scaled_rgb_image = ((self.rgb_image/self.nimages)*255).astype(np.int32) + self.update_xyz_image_kernel(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(545.0), gpuarray.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) - image = scaled_rgb_image[:,0] << 16 | scaled_rgb_image[:,1] << 8 | scaled_rgb_image[:,2] + self.update_xyz_image_kernel(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(445.0), gpuarray.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1)) + self.context.pop() - pygame.surfarray.blit_array(self.screen, image.reshape(self.size)) - pygame.display.flip() + self.nimages += 1 + + def process_image(self): + with self.lock: + self.context.push() + self.process_image_kernel(np.int32(self.width*self.height), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.width*self.height)//self.nblocks+1,1)) + self.context.pop() def screenshot(self, dir=''): root, ext = 'screenshot', 'png' @@ -183,6 +190,11 @@ class Camera(Thread): self.axis2 = rotate(self.axis2, phi, n) self.context.pop() + if self.render: + self.clear_image() + + self.update() + def translate(self, v): with self.lock: self.context.push() @@ -191,22 +203,25 @@ class Camera(Thread): self.point += v self.context.pop() + if self.render: + self.clear_image() + + self.update() + def update(self): if self.render: - self.acquire_image() - self.render_image() - return + while self.nlookup_calls < 10: + self.update_xyz_lookup(self.source_position) + self.update_image() + self.process_image() + else: + with self.lock: + self.context.push() + 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.pop() with self.lock: self.context.push() - t0 = time.time() - 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 '\relapsed %f sec.' % elapsed, - sys.stdout.flush() - pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size)) pygame.display.flip() self.context.pop() @@ -218,28 +233,23 @@ class Camera(Thread): clicked = False shift = False - current_layer = 0 + #current_layer = 0 while not done: self.clock.tick(20) if self.render and not clicked and not pygame.event.peek(KEYDOWN): - self.acquire_image() - self.render_image() + self.update() 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.clear_image() - self.update() if event.button == 5: v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) - self.clear_image() - self.update() if event.button == 1: clicked = True @@ -263,52 +273,41 @@ class Camera(Thread): if shift: v = mouse_direction*self.scale*length/float(self.width) self.translate(v) - self.clear_image() - 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.clear_image() - self.update() if event.type == KEYDOWN: if event.key == K_a: v = self.scale*self.axis1/10.0 self.translate(v) - self.clear_image() - self.update() if event.key == K_d: v = -self.scale*self.axis1/10.0 self.translate(v) - self.clear_image() - self.update() if event.key == K_w: v = self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) - self.clear_image() - self.update() if event.key == K_s: v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) - self.clear_image() - self.update() if event.key == K_SPACE: v = self.scale*self.axis2/10.0 self.translate(v) - self.clear_image() - self.update() if event.key == K_LCTRL: v = -self.scale*self.axis2/10.0 self.translate(v) - self.clear_image() - self.update() + + if event.key == K_p: + self.clear_xyz_lookup() + self.update_xyz_lookup(self.point) + self.source_position = self.point if event.key == K_LSHIFT or event.key == K_RSHIFT: shift = True @@ -347,6 +346,9 @@ class Camera(Thread): self.screenshot() if event.key == K_F5: + if not hasattr(self, 'rng_states_gpu'): + self.initialize_render() + self.render = not self.render self.update() @@ -358,7 +360,6 @@ class Camera(Thread): done = True break - print pygame.display.quit() if __name__ == '__main__': |