diff options
Diffstat (limited to 'camera.py')
| -rw-r--r-- | camera.py | 261 |
1 files changed, 166 insertions, 95 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,32 @@ from pycuda import gpuarray from pycuda.characterize import sizeof import pycuda.driver as cuda +from subprocess import call +import shutil +import tempfile + +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 encode_movie(dir): + root, ext = 'movie', 'avi' + for i in count(): + filename = '.'.join([root + str(i).zfill(5), ext]) + + if not os.path.exists(filename): + break + + call(['mencoder', 'mf://' + dir + '/*.png', '-mf', 'fps=10', '-o', filename, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000']) + + shutil.rmtree(dir) + + print 'movie saved to %s.' % filename + 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. @@ -62,9 +90,11 @@ class Camera(Thread): self.context.push() self.ray_trace_kernel = self.module.get_function('ray_trace') self.rotate_kernel = self.module.get_function('rotate') + self.rotate_around_point_kernel = self.module.get_function('rotate_around_point') 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 +107,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() @@ -86,85 +116,84 @@ class Camera(Thread): self.pixels_gpu = gpuarray.zeros(self.width*self.height, dtype=np.int32) self.context.pop() + self.movie = False + self.movie_index = 0 + self.movie_dir = None 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)) + + 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)) - 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(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() - image = scaled_rgb_image[:,0] << 16 | scaled_rgb_image[:,1] << 8 | scaled_rgb_image[:,2] + self.nimages += 1 - pygame.surfarray.blit_array(self.screen, image.reshape(self.size)) - pygame.display.flip() + 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=''): + def screenshot(self, dir='', start=0): root, ext = 'screenshot', 'png' - for i in count(): - filename = os.path.join(dir, '.'.join([root + str(i).zfill(4), ext])) + for i in count(start): + filename = os.path.join(dir, '.'.join([root + str(i).zfill(5), ext])) if not os.path.exists(filename): break @@ -175,14 +204,34 @@ class Camera(Thread): 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.rotate_kernel(np.int32(self.pixels_gpu.size), self.origins_gpu, np.float32(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, np.float32(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() + if self.render: + self.clear_image() + + self.update() + + def rotate_around_point(self, phi, n, point): + with self.lock: + self.context.push() + self.rotate_around_point_kernel(np.int32(self.origins_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), gpuarray.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) + self.rotate_kernel(np.int32(self.directions_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1)) + self.context.pop() + + self.axis1 = rotate(self.axis1, phi, n) + self.axis2 = rotate(self.axis2, phi, n) + + if self.render: + self.clear_image() + + self.update() + def translate(self, v): with self.lock: self.context.push() @@ -191,55 +240,58 @@ 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() + if self.movie: + self.screenshot(self.movie_dir, self.movie_index) + self.movie_index += 1 + def run(self): self.update() done = False clicked = False shift = False + ctrl = 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,56 +315,56 @@ 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 ctrl: + self.rotate_around_point(phi, n, self.point) + else: + self.rotate(phi, n) 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) + # if event.key == K_LCTRL: + # v = -self.scale*self.axis2/10.0 + # self.translate(v) + + if event.key == K_F6: + self.clear_xyz_lookup() self.clear_image() - self.update() + self.source_position = self.point + + if event.key == K_p: + for i in range(10): + self.update_xyz_lookup(self.point) + self.source_position = self.point if event.key == K_LSHIFT or event.key == K_RSHIFT: shift = True + if event.key == K_LCTRL or event.key == K_RCTRL: + ctrl = True + if event.key == K_ESCAPE: done = True break @@ -347,18 +399,37 @@ 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.clear_image() self.update() + if event.key == K_m: + if self.movie: + encode_movie(self.movie_dir) + self.movie_dir = None + self.movie = False + else: + self.movie_index = 0 + self.movie_dir = tempfile.mkdtemp() + self.movie = True + if event.type == KEYUP: if event.key == K_LSHIFT or event.key == K_RSHIFT: shift = False + if event.key == K_LCTRL or event.key == K_RCTRL: + ctrl = False + if event.type == pygame.QUIT: done = True break - print + if self.movie: + encode_movie(self.movie_dir) + pygame.display.quit() if __name__ == '__main__': |