diff options
Diffstat (limited to 'camera.py')
| -rwxr-xr-x | camera.py | 129 |
1 files changed, 92 insertions, 37 deletions
@@ -26,6 +26,7 @@ import pycuda.driver as cuda from subprocess import call import shutil import tempfile +import inspect def buildable(identifier): """ @@ -101,17 +102,18 @@ def encode_movie(dir): print 'movie saved to %s.' % filename class Camera(Thread): - def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False): + def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False, alpha_depth=10): Thread.__init__(self) self.geometry = geometry self.device_id = device_id self.size = size self.enable3d = enable3d self.green_magenta = green_magenta + self.alpha_depth = alpha_depth self.unique_bvh_layers = np.unique(self.geometry.layers) self.currentlayer = None - self.bvh_layers = {self.currentlayer : self.geometry} + self.bvh_layers = {} try: import spnav as spnav_module @@ -122,15 +124,19 @@ class Camera(Thread): def init_gpu(self): self.gpu_instance = gpu.GPU(self.device_id) - self.gpu_geometry = gpu.GPUGeometry(self.gpu_instance, geometry) - self.gpu_funcs = gpu.GPUFuncs(self.gpu_instance.module) + self.gpu_geometry = gpu.GPUGeometry(self.geometry) + self.gpu_funcs = gpu.GPUFuncs(gpu.get_cu_module('mesh.h')) + self.hybrid_funcs = gpu.GPUFuncs(gpu.get_cu_module('hybrid_render.cu')) - self.width, self.height = size + self.gpu_geometries = [self.gpu_geometry] + + self.width, self.height = self.size self.npixels = self.width*self.height pygame.init() - self.screen = pygame.display.set_mode(size) + self.window = pygame.display.set_mode(self.size) + self.screen = pygame.Surface(self.size, pygame.SRCALPHA) pygame.display.set_caption('') self.clock = pygame.time.Clock() @@ -181,9 +187,12 @@ class Camera(Thread): self.rays = gpu.GPURays(pos, dir) + self.distance_array = ga.empty(self.alpha_depth*self.rays.pos.size, dtype=np.float32) + self.index_array = ga.empty(self.alpha_depth*self.rays.pos.size, dtype=np.uint32) + self.n_array = ga.zeros(self.rays.pos.size, dtype=np.uint32) + self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32) - self.alpha = True self.movie = False self.movie_index = 0 self.movie_dir = None @@ -217,7 +226,7 @@ class Camera(Thread): for wavelength, rgb_tuple in \ zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]): for i in range(self.xyz_lookup1_gpu.size//(self.npixels)+1): - self.gpu_funcs.update_xyz_lookup(np.int32(self.npixels), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.npixels), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.npixels//self.nblocks+1,1)) + self.hybrid_funcs.update_xyz_lookup(np.int32(self.npixels), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.npixels), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), self.gpu_geometry.gpudata, block=(self.nblocks,1,1), grid=(self.npixels//self.nblocks+1,1)) self.nlookup_calls += 1 @@ -233,7 +242,7 @@ class Camera(Thread): def update_image_from_rays(self, image, rays): for wavelength, rgb_tuple in \ zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]): - self.gpu_funcs.update_xyz_image(np.int32(rays.pos.size), self.rng_states_gpu, rays.pos, rays.dir, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, image, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(rays.pos.size//self.nblocks+1,1)) + self.hybrid_funcs.update_xyz_image(np.int32(rays.pos.size), self.rng_states_gpu, rays.pos, rays.dir, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, image, np.int32(self.nlookup_calls), np.int32(self.max_steps), self.gpu_geometry.gpudata, block=(self.nblocks,1,1), grid=(rays.pos.size//self.nblocks+1,1)) def update_image(self): if self.enable3d: @@ -246,10 +255,10 @@ class Camera(Thread): def process_image(self): if self.enable3d: - self.gpu_funcs.process_image(np.int32(self.pixels1_gpu.size), self.image1_gpu, self.pixels1_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels1_gpu.size)//self.nblocks+1,1)) - self.gpu_funcs.process_image(np.int32(self.pixels2_gpu.size), self.image2_gpu, self.pixels2_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels2_gpu.size)//self.nblocks+1,1)) + self.hybrid_funcs.process_image(np.int32(self.pixels1_gpu.size), self.image1_gpu, self.pixels1_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels1_gpu.size)//self.nblocks+1,1)) + self.hybrid_funcs.process_image(np.int32(self.pixels2_gpu.size), self.image2_gpu, self.pixels2_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels2_gpu.size)//self.nblocks+1,1)) else: - self.gpu_funcs.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1)) + self.hybrid_funcs.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1)) def screenshot(self, dir='', start=0): root, ext = 'screenshot', 'png' @@ -319,29 +328,28 @@ class Camera(Thread): self.update() - def update_pixels(self): + def update_pixels(self, gpu_geometry=None, keep_last_render=False): + if gpu_geometry is None: + gpu_geometry = self.gpu_geometry + if self.render: while self.nlookup_calls < 10: self.update_xyz_lookup(self.source_position) self.update_image() self.process_image() else: - if self.alpha: - if self.enable3d: - self.gpu_funcs.ray_trace_alpha(np.int32(self.rays1.pos.size), self.rays1.pos, self.rays1.dir, self.pixels1_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1)) - self.gpu_funcs.ray_trace_alpha(np.int32(self.rays2.pos.size), self.rays2.pos, self.rays2.dir, self.pixels2_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1)) - else: - self.gpu_funcs.ray_trace_alpha(np.int32(self.rays.pos.size), self.rays.pos, self.rays.dir, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.rays.pos.size//self.nblocks+1,1)) + if self.enable3d: + self.rays1.render(gpu_geometry, self.pixels1_gpu, + self.alpha_depth, keep_last_render) + self.rays2.render(gpu_geometry, self.pixels2_gpu, + self.alpha_depth, keep_last_render) else: - if self.enable3d: - self.gpu_funcs.ray_trace(np.int32(self.rays1.pos.size), self.rays1.pos, self.rays1.dir, self.pixels1_gpu, block=(self.nblocks,1,1), grid=(self.rays1.pos.size//self.nblocks+1,1)) - self.gpu_funcs.ray_trace(np.int32(self.rays2.pos.size), self.rays2.pos, self.rays2.dir, self.pixels2_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1)) - else: - self.gpu_funcs.ray_trace(np.int32(self.rays.pos.size), self.rays.pos, self.rays.dir, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.rays.pos.size//self.nblocks+1,1)) + self.rays.render(gpu_geometry, self.pixels_gpu, + self.alpha_depth, keep_last_render) - def update(self): + def update_viewing_angle(self): if self.enable3d: - self.gpu_funcs.distance_to_mesh(np.int32(self.scope_rays.pos.size), self.scope_rays.pos, self.scope_rays.dir, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_rays.pos.size//self.nblocks,1)) + self.gpu_funcs.distance_to_mesh(np.int32(self.scope_rays.pos.size), self.scope_rays.pos, self.scope_rays.dir, self.gpu_geometry.gpudata, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_rays.pos.size//self.nblocks,1)) baseline = ga.min(self.distances_gpu).get().item() @@ -376,7 +384,16 @@ class Camera(Thread): self.viewing_angle = new_viewing_angle - self.update_pixels() + def update(self): + if self.enable3d: + self.update_viewing_angle() + + n = len(self.gpu_geometries) + for i, gpu_geometry in enumerate(self.gpu_geometries): + if i == 0: + self.update_pixels(gpu_geometry) + else: + self.update_pixels(gpu_geometry, keep_last_render=True) if self.enable3d: pixels1 = self.pixels1_gpu.get() @@ -386,12 +403,18 @@ class Camera(Thread): pixels = (pixels1 & 0x00ff00) | (pixels2 & 0xff00ff) else: pixels = (pixels1 & 0xff0000) | (pixels2 & 0x00ffff) + + alpha = ((0xff & (pixels1 >> 24)) + (0xff & (pixels2 >> 24)))/2 + + pixels |= (alpha << 24) else: pixels = self.pixels_gpu.get() pygame.surfarray.blit_array(self.screen, pixels.reshape(self.size)) if self.doom_mode: self.screen.blit(self.doom_hud, self.doom_rect) + self.window.fill(0) + self.window.blit(self.screen, (0,0)) pygame.display.flip() if self.movie: @@ -399,12 +422,17 @@ class Camera(Thread): self.movie_index += 1 def loadlayer(self, layer): - try: - layergeometry = self.bvh_layers[layer] - layergeometry.activate() - except KeyError: - layergeometry = build(bvh_mesh(self.geometry, layer),8) - self.bvh_layers[layer] = gpu.GPUGeometry(self.gpu, layergeometry) + if layer is None: + self.gpu_geometries = [self.gpu_geometry] + else: + try: + gpu_geometry = self.bvh_layers[layer] + except KeyError: + geometry = build(bvh_mesh(self.geometry, layer), 8) + gpu_geometry = gpu.GPUGeometry(geometry) + self.bvh_layers[layer] = gpu_geometry + + self.gpu_geometries = [self.gpu_geometry, gpu_geometry] self.update() @@ -483,6 +511,15 @@ class Camera(Thread): self.done = True return + elif event.key == K_EQUALS: + self.alpha_depth += 1 + self.update() + + elif event.key == K_MINUS: + if self.alpha_depth > 0: + self.alpha_depth -= 1 + self.update() + elif event.key == K_PAGEDOWN: if self.currentlayer is None: self.currentlayer = self.unique_bvh_layers[-1] @@ -516,10 +553,6 @@ class Camera(Thread): self.clear_image() self.update() - elif event.key == K_F8: - self.alpha = not self.alpha - self.update() - elif event.key == K_m: if self.movie: encode_movie(self.movie_dir) @@ -629,6 +662,20 @@ class EventViewer(Camera): Camera.__init__(self, geometry, **kwargs) self.rr = RootReader(filename) + def render_particle_track(self): + marker = Solid(make.cube(0.1), vacuum, vacuum) + + geometry = Geometry() + for pos in self.ev.photons_beg.pos[::100]: + geometry.add_solid(marker, displacement=pos) + + geometry.build(bits=11) + gpu_geometry = gpu.GPUGeometry(geometry) + + self.gpu_geometries = [self.gpu_geometry, gpu_geometry] + + self.update() + def color_hit_pmts(self): self.gpu_geometry.reset_colors() @@ -650,6 +697,7 @@ class EventViewer(Camera): pass else: self.color_hit_pmts() + self.render_particle_track() return elif event.key == K_PAGEDOWN: @@ -659,10 +707,17 @@ class EventViewer(Camera): pass else: self.color_hit_pmts() + self.render_particle_track() return Camera.process_event(self, event) +def view(obj, size=(800,600), **camera_kwargs): + geometry = build(obj, 8) + camera = Camera(geometry, size, **camera_kwargs) + camera.start() + camera.join() + if __name__ == '__main__': import optparse import inspect |