diff options
Diffstat (limited to 'camera.py')
| -rwxr-xr-x | camera.py | 476 |
1 files changed, 216 insertions, 260 deletions
@@ -2,12 +2,12 @@ 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 color import map_wavelength, map_to_color +from gpu import GPU, CUDAFuncs +from tools import timeit from transform import rotate import pygame @@ -21,13 +21,9 @@ 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 +from histogram import * + +import matplotlib.pyplot as plt def encode_movie(dir): root, ext = 'movie', 'avi' @@ -72,24 +68,24 @@ def get_rays(position, size = (800, 600), film_size = (0.035, 0.024), focal_leng return grid, focal_point-grid class Camera(Thread): - def __init__(self, geometry, module, context, lock=None, size=(800,600), - spnav=False): + def __init__(self, geometry, size=(800,600), device_id=None, spnav=False): Thread.__init__(self) self.geometry = geometry - self.module = module - self.context = context + self.device_id = device_id + self.size = size + self.spnav = spnav if spnav: import spnav as spnav_module self.spnav_module = spnav_module - if lock is None: - self.lock = Lock() - else: - self.lock = lock + def init_gpu(self): + self.gpu = GPU(self.device_id) + self.gpu.load_geometry(geometry) + + self.kernels = CUDAFuncs(self.gpu.module, ['ray_trace', 'rotate', 'rotate_around_point', 'translate', 'update_xyz_lookup', 'update_xyz_image', 'process_image', 'init_rng']) - self.size = size self.width, self.height = size pygame.init() @@ -97,18 +93,6 @@ class Camera(Thread): pygame.display.set_caption('') self.clock = pygame.time.Clock() - 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.rotate_around_point_kernel = self.module.get_function('rotate_around_point') - self.translate_kernel = self.module.get_function('translate') - 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() - lower_bound, upper_bound = self.geometry.mesh.get_bounds() self.scale = np.linalg.norm(upper_bound-lower_bound) @@ -120,12 +104,9 @@ class Camera(Thread): origins, directions = get_rays(self.point, self.size) - 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.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.movie = False self.movie_index = 0 @@ -134,16 +115,12 @@ class Camera(Thread): @timeit def initialize_render(self): - with self.lock: - self.context.push() - print 'beginning initialize_render()' - 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() - self.context.pop() + self.rng_states_gpu = cuda.mem_alloc(self.width*self.height*sizeof('curandStateXORWOW', '#include <curand_kernel.h>')) + self.kernels.init_rng(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() self.source_position = self.point @@ -152,54 +129,39 @@ class Camera(Thread): self.max_steps = 10 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() + 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.nlookup_calls = 0 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)) + for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1): + self.kernels.update_xyz_lookup(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)) - 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.kernels.update_xyz_lookup(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() + for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1): + self.kernels.update_xyz_lookup(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.nlookup_calls += 1 def clear_image(self): - with self.lock: - self.context.push() - self.image_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0)) - self.context.pop() + self.image_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0)) self.nimages = 0 def update_image(self): - with self.lock: - self.context.push() - 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.kernels.update_xyz_image(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)) + self.kernels.update_xyz_image(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)) - 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() + self.kernels.update_xyz_image(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.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() + self.kernels.process_image(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)) def screenshot(self, dir='', start=0): root, ext = 'screenshot', 'png' @@ -214,15 +176,12 @@ class Camera(Thread): 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, 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.kernels.rotate(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.kernels.rotate(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() + self.point = rotate(self.point, phi, n) + self.axis1 = rotate(self.axis1, phi, n) + self.axis2 = rotate(self.axis2, phi, n) if self.render: self.clear_image() @@ -230,11 +189,8 @@ class Camera(Thread): self.update() def rotate_around_point(self, phi, n, point, redraw=True): - 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.kernels.rotate_around_point(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.kernels.rotate(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.axis1 = rotate(self.axis1, phi, n) self.axis2 = rotate(self.axis2, phi, n) @@ -246,12 +202,9 @@ class Camera(Thread): self.update() def translate(self, v, redraw=True): - 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.kernels.translate(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() + self.point += v if redraw: if self.render: @@ -266,38 +219,135 @@ class Camera(Thread): 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() + self.kernels.ray_trace(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)) - with self.lock: - self.context.push() - pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size)) - pygame.display.flip() - self.context.pop() + pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size)) + pygame.display.flip() if self.movie: self.screenshot(self.movie_dir, self.movie_index) self.movie_index += 1 + def process_event(self, event): + if event.type == MOUSEBUTTONDOWN: + if event.button == 4: + v = self.scale*np.cross(self.axis1,self.axis2)/10.0 + self.translate(v) + + elif event.button == 5: + v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 + self.translate(v) + + elif event.button == 1: + mouse_position = pygame.mouse.get_rel() + self.clicked = True + + elif event.type == MOUSEBUTTONUP: + if event.button == 1: + self.clicked = False + + elif event.type == MOUSEMOTION and self.clicked: + movement = np.array(pygame.mouse.get_rel()) + + if (movement == 0).all(): + return + + length = np.linalg.norm(movement) + + mouse_direction = movement[0]*self.axis1 + movement[1]*self.axis2 + mouse_direction /= np.linalg.norm(mouse_direction) + + if pygame.key.get_mods() & (KMOD_LSHIFT | KMOD_RSHIFT): + v = mouse_direction*self.scale*length/float(self.width) + self.translate(v) + else: + phi = np.float32(2*np.pi*length/float(self.width)) + n = rotate(mouse_direction, np.pi/2, -np.cross(self.axis1,self.axis2)) + + if pygame.key.get_mods() & KMOD_LCTRL: + self.rotate_around_point(phi, n, self.point) + else: + self.rotate(phi, n) + + elif event.type == KEYDOWN: + if event.key == K_a: + v = self.scale*self.axis1/10.0 + self.translate(v) + + elif event.key == K_d: + v = -self.scale*self.axis1/10.0 + self.translate(v) + + elif event.key == K_w: + v = self.scale*np.cross(self.axis1,self.axis2)/10.0 + self.translate(v) + + elif event.key == K_s: + v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 + self.translate(v) + + elif event.key == K_SPACE: + v = self.scale*self.axis2/10.0 + self.translate(v) + + elif event.key == K_F6: + self.clear_xyz_lookup() + self.clear_image() + self.source_position = self.point + + elif event.key == K_p: + for i in range(100): + self.update_xyz_lookup(self.point) + self.source_position = self.point + + elif event.key == K_ESCAPE: + self.done = True + return + + elif event.key == K_F12: + self.screenshot() + + elif event.key == K_F5: + if not hasattr(self, 'rng_states_gpu'): + self.initialize_render() + + self.render = not self.render + self.clear_image() + self.update() + + elif 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 + + elif event.type == pygame.QUIT: + self.done = True + return + def run(self): + self.init_gpu() + if self.spnav: self.spnav_module.spnav_open() self.update() - done = False - clicked = False - shift = False - ctrl = False + self.done = False + self.clicked = False + #shift = False + #ctrl = False #current_layer = 0 - while not done: + while not self.done: self.clock.tick(20) - if self.render and not clicked and not pygame.event.peek(KEYDOWN): + if self.render and not self.clicked and not pygame.event.peek(KEYDOWN): self.update() # Space Navigator controls @@ -353,156 +403,70 @@ class Camera(Thread): 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) - - if event.button == 5: - v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 - self.translate(v) - - 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 + self.process_event(event) - 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) - else: - phi = np.float32(2*np.pi*length/float(self.width)) - n = rotate(mouse_direction, np.pi/2, -np.cross(self.axis1,self.axis2)) - - 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) - - if event.key == K_d: - v = -self.scale*self.axis1/10.0 - self.translate(v) - - if event.key == K_w: - v = self.scale*np.cross(self.axis1,self.axis2)/10.0 - self.translate(v) - - if event.key == K_s: - v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 - self.translate(v) - - if event.key == K_SPACE: - 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.source_position = self.point - - if event.key == K_p: - for i in range(100): - 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 - - #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 + if self.movie: + encode_movie(self.movie_dir) - # geometry.load(module, color=True) - # update() - # except IndexError: - # print 'no further layers to view' + pygame.display.quit() + if self.spnav: + self.spnav_module.spnav_close() - if event.key == K_F12: - self.screenshot() + del self.gpu - if event.key == K_F5: - if not hasattr(self, 'rng_states_gpu'): - self.initialize_render() +class EventViewer(Camera): + def __init__(self, geometry, filename, **kwargs): + Camera.__init__(self, geometry, **kwargs) - self.render = not self.render - self.clear_image() - self.update() + import ROOT - 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 + self.f = ROOT.TFile(filename) + self.T = self.f.Get('T') + self.T.GetEntry(0) - if event.type == KEYUP: - if event.key == K_LSHIFT or event.key == K_RSHIFT: - shift = False + @timeit + def color_hit_pmts(self): + self.gpu.reset_colors() - if event.key == K_LCTRL or event.key == K_RCTRL: - ctrl = False + solid_ids = np.empty(len(self.T.ev.channel), np.uint32) + t = np.empty(len(self.T.ev.channel), np.float32) + q = np.empty(len(self.T.ev.channel), np.float32) - if event.type == pygame.QUIT: - done = True - break + for i, channel in enumerate(self.T.ev.channel): + solid_ids[i] = channel.channel_id + t[i] = channel.t + q[i] = channel.q - if self.movie: - encode_movie(self.movie_dir) + self.gpu.color_solids(solid_ids, map_to_color(t)) + self.update() - pygame.display.quit() - if self.spnav: - self.spnav_module.spnav_close() + plt.clf() + plt.hist(t*1e9, 100) + plt.xlabel('Time (ns)') + plt.ylabel('Entries') + plt.draw() + + def process_event(self, event): + if event.type == KEYDOWN: + if event.key == K_PAGEUP: + entry = self.T.GetReadEntry() + if entry < self.T.GetEntries() - 1: + self.T.GetEntry(entry+1) + self.color_hit_pmts() + return + + if event.key == K_PAGEDOWN: + entry = self.T.GetReadEntry() + if entry > 0: + self.T.GetEntry(entry-1) + self.color_hit_pmts() + return + + Camera.process_event(self, event) + def run(self): + plt.ion() + Camera.run(self) if __name__ == '__main__': import optparse @@ -512,11 +476,8 @@ if __name__ == '__main__': 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=10) @@ -524,6 +485,7 @@ if __name__ == '__main__': # help='load bounding volumes', default=False) parser.add_option('-r', '--resolution', dest='resolution', help='specify resolution', default='800,600') + parser.add_option('-i', dest='io_file', default=None) parser.add_option('--spnav', action='store_true', dest='spnav', help='activate Space Navigator support', default=False) options, args = parser.parse_args() @@ -538,6 +500,7 @@ if __name__ == '__main__': if ext.lower() in ('.stl', '.bz2'): obj = mesh_from_stl(args[0]) + else: members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(scenes)) @@ -552,16 +515,9 @@ if __name__ == '__main__': else: raise Exception("can't find object %s" % args[0]) - from pycuda.tools import make_default_context - from gpu import * - geometry = build(obj, options.bits) - - lock = Lock() - - gpu = GPU() - gpu.load_geometry(geometry) - gpu.context.pop() - camera = Camera(geometry, gpu.module, gpu.context, lock, size=size, - spnav=options.spnav) + if options.io_file is not None: + camera = EventViewer(geometry, options.io_file, size=size, spnav=options.spnav) + else: + camera = Camera(geometry, size, spnav=options.spnav) camera.start() |