diff options
author | Anthony LaTorre <tlatorre9@gmail.com> | 2011-09-09 09:53:42 -0400 |
---|---|---|
committer | Anthony LaTorre <tlatorre9@gmail.com> | 2011-09-09 09:53:42 -0400 |
commit | 8f39536a7a5a15afeba52b1492f1d84fed95054f (patch) | |
tree | fdd78495d9fb3bfae47e182b386308a808420a1f /camera.py | |
parent | 25de37ef26b8356d8fa012c709cd68c1fb327227 (diff) | |
download | chroma-8f39536a7a5a15afeba52b1492f1d84fed95054f.tar.gz chroma-8f39536a7a5a15afeba52b1492f1d84fed95054f.tar.bz2 chroma-8f39536a7a5a15afeba52b1492f1d84fed95054f.zip |
chroma-cam script now used to view geometries.
Diffstat (limited to 'camera.py')
-rwxr-xr-x | camera.py | 333 |
1 files changed, 147 insertions, 186 deletions
@@ -1,59 +1,33 @@ #!/usr/bin/env python import numpy as np -from itertools import product, count -from threading import Thread, Lock +import itertools +import threading +import multiprocessing import os -import sys -from chroma.color import map_wavelength, map_to_color -from chroma.geometry import Mesh, Solid, Geometry -import chroma.make as make -import matplotlib.pyplot as plt +from subprocess import call +import shutil +import tempfile +import inspect -from chroma import gpu -from chroma.tools import timeit +import pycuda.driver as cuda +from pycuda import gpuarray as ga + +from chroma.color import map_to_color +from chroma.geometry import Mesh, Solid, Geometry from chroma.transform import rotate -from chroma.optics import vacuum, lambertian_surface -import chroma.project as project +from chroma.optics import vacuum +from chroma.project import from_film from chroma.fileio.root import RootReader +from chroma import make +from chroma import gpu import pygame from pygame.locals import * -from pycuda import gpuarray as ga -from pycuda.characterize import sizeof -import pycuda.driver as cuda - -from subprocess import call -import shutil -import tempfile -import inspect - -def buildable(identifier): - """ - Create a decorator which tags a function as buildable and assigns the - identifying string `identifier`. - - Example: - >>> @buildable('my_sphere') - >>> def build_my_sphere(): - >>> g = Geometry() - >>> g.add_solid(Solid(sphere(), vacuum, water)) - >>> return g - """ - def tag_as_buildable(func): - func.buildable = True - func.identifier = identifier - return func - return tag_as_buildable - -def build(obj, bits): +def build(obj): """Construct and build a geometry from `obj`.""" if inspect.isfunction(obj): - try: - if obj.buildable: - obj = obj() - except AttributeError: - raise Exception('function %s is not buildable.' % obj.__name__) + obj = obj() if isinstance(obj, Geometry): geometry = obj @@ -62,11 +36,11 @@ def build(obj, bits): geometry.add_solid(obj) elif isinstance(obj, Mesh): geometry = Geometry() - geometry.add_solid(Solid(obj, vacuum, vacuum, surface=lambertian_surface, color=0x99ffffff)) + geometry.add_solid(Solid(obj, vacuum, vacuum, color=0x99ffffff)) else: - raise Exception('cannot build type %s' % type(obj)) + raise TypeError('cannot build type %s' % type(obj)) - geometry.build(bits) + geometry.build() return geometry @@ -82,39 +56,43 @@ def bvh_mesh(geometry, layer): mesh = make.box(dx, dy, dz, center) - for center, dx, dy, dz in zip(np.mean([lower_bounds,upper_bounds],axis=0),*zip(*upper_bounds-lower_bounds))[1:]: + for center, dx, dy, dz in zip(np.mean([lower_bounds,upper_bounds],axis=0), + *zip(*upper_bounds-lower_bounds))[1:]: mesh += make.box(dx,dy,dz,center) return mesh def encode_movie(dir): root, ext = 'movie', 'avi' - for i in count(): - filename = '.'.join([root + str(i).zfill(5), ext]) + for i in itertools.count(): + path = '.'.join([root + str(i).zfill(5), ext]) - if not os.path.exists(filename): + if not os.path.exists(path): break - call(['mencoder', 'mf://' + dir + '/*.png', '-mf', 'fps=10', '-o', filename, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000']) + call(['mencoder', 'mf://' + dir + '/*.png', '-mf', 'fps=10', '-o', + path, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000']) shutil.rmtree(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, alpha_depth=10): - Thread.__init__(self) +class Camera(multiprocessing.Process): + def __init__(self, geometry, size=(800,600), device_id=None): + multiprocessing.Process.__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.display3d = False + self.green_magenta = False + self.max_alpha_depth = 10 + self.alpha_depth = 10 + try: import spnav as spnav_module self.spnav_module = spnav_module @@ -123,7 +101,8 @@ class Camera(Thread): self.spnav = False def init_gpu(self): - self.gpu_instance = gpu.GPU(self.device_id) + self.context = gpu.create_context(self.device_id) + 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')) @@ -134,10 +113,6 @@ class Camera(Thread): self.npixels = self.width*self.height - pygame.init() - 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() self.doom_mode = False @@ -156,55 +131,73 @@ class Camera(Thread): self.nblocks = 64 - self.point = np.array([0, -self.scale*1.0, (lower_bound[2]+upper_bound[2])/2]) + self.point = np.array([0, -self.scale*1.0, + (lower_bound[2]+upper_bound[2])/2]) + self.axis1 = np.array([0,0,1], float) self.axis2 = np.array([1,0,0], float) self.film_width = 0.035 - if self.enable3d: - self.point1 = self.point-(self.scale/60,0,0) - self.point2 = self.point+(self.scale/60,0,0) + pos, dir = from_film(self.point, axis1=self.axis1, axis2=self.axis2, + size=self.size, width=self.film_width) - self.viewing_angle = 0.0 + self.rays = gpu.GPURays(pos, dir, max_alpha_depth=self.max_alpha_depth) - pos1, dir1 = project.from_film(self.point1, size=self.size, width=self.film_width) - pos2, dir2 = project.from_film(self.point2, size=self.size, width=self.film_width) + self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32) - self.rays1 = gpu.GPURays(pos1, dir1) - self.rays2 = gpu.GPURays(pos2, dir2) + self.movie = False + self.movie_index = 0 + self.movie_dir = None + self.hybrid_render = False - scope_pos, scope_dir = project.from_film(self.point, size=np.array(self.size)/4.0, width=self.film_width/4.0) + def disable3d(self): + pos, dir = from_film(self.point, axis1=self.axis1, axis2=self.axis2, + size=self.size, width=self.film_width) - self.scope_rays = gpu.GPURays(scope_pos, scope_dir) + self.rays = gpu.GPURays(pos, dir, max_alpha_depth=self.max_alpha_depth) + + self.display3d = False - self.pixels1_gpu = ga.empty(self.width*self.height, dtype=np.int32) - self.pixels2_gpu = ga.empty(self.width*self.height, dtype=np.int32) + def enable3d(self): + self.point1 = self.point-(self.scale/60)*self.axis2 + self.point2 = self.point+(self.scale/60)*self.axis2 - self.distances_gpu = ga.empty(self.scope_rays.pos.size, dtype=np.float32) - else: - pos, dir = project.from_film(self.point, size=self.size, width=self.film_width) + self.viewing_angle = 0.0 - self.rays = gpu.GPURays(pos, dir) + pos1, dir1 = from_film(self.point1, axis1=self.axis1, axis2=self.axis2, + size=self.size, width=self.film_width) + pos2, dir2 = from_film(self.point2, axis1=self.axis1, axis2=self.axis2, + size=self.size, width=self.film_width) - 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.rays1 = gpu.GPURays(pos1, dir1, + max_alpha_depth=self.max_alpha_depth) + self.rays2 = gpu.GPURays(pos2, dir2, + max_alpha_depth=self.max_alpha_depth) - self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32) + scope_size = (self.size[0]//4, self.size[0]//4) - self.movie = False - self.movie_index = 0 - self.movie_dir = None - self.render = False + scope_pos, scope_dir = from_film(self.point, axis1=self.axis1, + axis2=self.axis2, size=scope_size, + width=self.film_width/4.0) + + self.scope_rays = gpu.GPURays(scope_pos, scope_dir) + + self.pixels1_gpu = ga.empty(self.width*self.height, dtype=np.int32) + self.pixels2_gpu = ga.empty(self.width*self.height, dtype=np.int32) + + self.distances_gpu = ga.empty(self.scope_rays.pos.size, + dtype=np.float32) + self.display3d = True - @timeit def initialize_render(self): self.rng_states_gpu = gpu.get_rng_states(self.npixels) - self.xyz_lookup1_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3) - self.xyz_lookup2_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3) + self.xyz_lookup1_gpu = ga.zeros(len(self.geometry.mesh.triangles), + dtype=ga.vec.float3) + self.xyz_lookup2_gpu = ga.zeros(len(self.geometry.mesh.triangles), + dtype=ga.vec.float3) - if self.enable3d: + if self.display3d: self.image1_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) self.image2_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) else: @@ -231,7 +224,7 @@ class Camera(Thread): self.nlookup_calls += 1 def clear_image(self): - if self.enable3d: + if self.display3d: self.image1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) self.image2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) else: @@ -245,7 +238,7 @@ class Camera(Thread): 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: + if self.display3d: self.update_image_from_rays(self.image1_gpu, self.rays1) self.update_image_from_rays(self.image2_gpu, self.rays2) else: @@ -254,7 +247,7 @@ class Camera(Thread): self.nimages += 1 def process_image(self): - if self.enable3d: + if self.display3d: 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: @@ -263,17 +256,17 @@ class Camera(Thread): def screenshot(self, dir='', start=0): root, ext = 'screenshot', 'png' - for i in count(start): - filename = os.path.join(dir, '.'.join([root + str(i).zfill(5), ext])) + for i in itertools.count(start): + path = os.path.join(dir, '.'.join([root + str(i).zfill(5), ext])) - if not os.path.exists(filename): + if not os.path.exists(path): break - pygame.image.save(self.screen, filename) - print 'image saved to %s' % filename + pygame.image.save(self.screen, path) + print 'image saved to %s' % path def rotate(self, phi, n): - if self.enable3d: + if self.display3d: self.rays1.rotate(phi, n) self.rays2.rotate(phi, n) self.scope_rays.rotate(phi, n) @@ -287,7 +280,7 @@ class Camera(Thread): self.axis1 = rotate(self.axis1, phi, n) self.axis2 = rotate(self.axis2, phi, n) - if self.render: + if self.hybrid_render: self.clear_image() self.update() @@ -296,7 +289,7 @@ class Camera(Thread): self.axis1 = rotate(self.axis1, phi, n) self.axis2 = rotate(self.axis2, phi, n) - if self.enable3d: + if self.display3d: self.rays1.rotate_around_point(phi, n, point) self.rays2.rotate_around_point(phi, n, point) self.scope_rays.rotate_around_point(phi, n, point) @@ -304,7 +297,7 @@ class Camera(Thread): self.rays.rotate_around_point(phi, n, point) if redraw: - if self.render: + if self.hybrid_render: self.clear_image() self.update() @@ -312,7 +305,7 @@ class Camera(Thread): def translate(self, v, redraw=True): self.point += v - if self.enable3d: + if self.display3d: self.rays1.translate(v) self.rays2.translate(v) self.scope_rays.translate(v) @@ -323,7 +316,7 @@ class Camera(Thread): self.rays.translate(v) if redraw: - if self.render: + if self.hybrid_render: self.clear_image() self.update() @@ -332,13 +325,13 @@ class Camera(Thread): if gpu_geometry is None: gpu_geometry = self.gpu_geometry - if self.render: + if self.hybrid_render: while self.nlookup_calls < 10: self.update_xyz_lookup(self.source_position) self.update_image() self.process_image() else: - if self.enable3d: + if self.display3d: self.rays1.render(gpu_geometry, self.pixels1_gpu, self.alpha_depth, keep_last_render) self.rays2.render(gpu_geometry, self.pixels2_gpu, @@ -348,7 +341,7 @@ class Camera(Thread): self.alpha_depth, keep_last_render) def update_viewing_angle(self): - if self.enable3d: + if self.display3d: 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() @@ -385,7 +378,7 @@ class Camera(Thread): self.viewing_angle = new_viewing_angle def update(self): - if self.enable3d: + if self.display3d: self.update_viewing_angle() n = len(self.gpu_geometries) @@ -395,7 +388,7 @@ class Camera(Thread): else: self.update_pixels(gpu_geometry, keep_last_render=True) - if self.enable3d: + if self.display3d: pixels1 = self.pixels1_gpu.get() pixels2 = self.pixels2_gpu.get() @@ -428,8 +421,8 @@ class Camera(Thread): try: gpu_geometry = self.bvh_layers[layer] except KeyError: - geometry = build(bvh_mesh(self.geometry, layer), 8) - gpu_geometry = gpu.GPUGeometry(geometry) + geometry = build(bvh_mesh(self.geometry, layer)) + gpu_geometry = gpu.GPUGeometry(geometry, print_usage=False) self.bvh_layers[layer] = gpu_geometry self.gpu_geometries = [self.gpu_geometry, gpu_geometry] @@ -470,7 +463,8 @@ class Camera(Thread): 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)) + 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) @@ -512,11 +506,12 @@ class Camera(Thread): return elif event.key == K_EQUALS: - self.alpha_depth += 1 + if self.alpha_depth < self.max_alpha_depth: + self.alpha_depth += 1 self.update() elif event.key == K_MINUS: - if self.alpha_depth > 0: + if self.alpha_depth > 1: self.alpha_depth -= 1 self.update() @@ -542,6 +537,16 @@ class Camera(Thread): self.loadlayer(self.currentlayer) + elif event.key == K_3: + if self.display3d: + self.disable3d() + else: + self.enable3d() + self.update() + + elif event.key == K_g: + self.green_magenta = not self.green_magenta + elif event.key == K_F12: self.screenshot() @@ -549,7 +554,7 @@ class Camera(Thread): if not hasattr(self, 'rng_states_gpu'): self.initialize_render() - self.render = not self.render + self.hybrid_render = not self.hybrid_render self.clear_image() self.update() @@ -598,9 +603,10 @@ class Camera(Thread): angle = length * 0.0001 * accelerate_factor axis /= length #print 'rotate:', angle, axis - self.rotate_around_point(angle, axis, self.point, redraw=False) + self.rotate_around_point(angle, axis, self.point, + redraw=False) - if self.render: + if self.hybrid_render: self.clear_image() self.update() @@ -611,7 +617,7 @@ class Camera(Thread): if not hasattr(self, 'rng_states_gpu'): self.initialize_render() - self.render = not self.render + self.hybrid_render = not self.hybrid_render self.clear_image() self.update() pygame.event.clear(pygame.SYSWMEVENT) @@ -621,6 +627,10 @@ class Camera(Thread): return def run(self): + self.window = pygame.display.set_mode(self.size) + self.screen = pygame.Surface(self.size, pygame.SRCALPHA) + pygame.display.set_caption('') + self.init_gpu() if self.spnav: @@ -629,7 +639,7 @@ class Camera(Thread): self.spnav_module.spnav_x11_open(wm_info['display'], wm_info['window']) pygame.event.set_allowed(pygame.SYSWMEVENT) - print 'Space Navigator support enabled.' + #print 'Space Navigator support enabled.' except: self.spnav = False @@ -641,11 +651,13 @@ class Camera(Thread): while not self.done: self.clock.tick(20) - if self.render and not self.clicked and not pygame.event.peek(KEYDOWN): + if self.hybrid_render and not self.clicked and \ + not pygame.event.peek(KEYDOWN): self.update() # Grab only last SYSWMEVENT (SpaceNav!) to avoid lagged controls - for event in pygame.event.get(pygame.SYSWMEVENT)[-1:] + pygame.event.get(): + for event in pygame.event.get(pygame.SYSWMEVENT)[-1:] + \ + pygame.event.get(): self.process_event(event) if self.movie: @@ -655,7 +667,7 @@ class Camera(Thread): if self.spnav: self.spnav_module.spnav_close() - del self.gpu_instance + self.context.pop() class EventViewer(Camera): def __init__(self, geometry, filename, **kwargs): @@ -674,8 +686,6 @@ class EventViewer(Camera): self.gpu_geometries = [self.gpu_geometry, gpu_geometry] - self.update() - def color_hit_pmts(self): self.gpu_geometry.reset_colors() @@ -686,7 +696,6 @@ class EventViewer(Camera): # Important: Compute range only with HIT channels solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max())) self.gpu_geometry.color_solids(hit, solid_colors) - self.update() def process_event(self, event): if event.type == KEYDOWN: @@ -697,7 +706,11 @@ class EventViewer(Camera): pass else: self.color_hit_pmts() - self.render_particle_track() + + if self.ev.photons_beg is not None: + self.render_particle_track() + + self.update() return elif event.key == K_PAGEDOWN: @@ -707,68 +720,16 @@ class EventViewer(Camera): pass else: self.color_hit_pmts() - self.render_particle_track() + + if self.ev.photons_beg is not None: + self.render_particle_track() + + self.update() return Camera.process_event(self, event) def view(obj, size=(800,600), **camera_kwargs): - geometry = build(obj, 8) + geometry = build(obj) camera = Camera(geometry, size, **camera_kwargs) camera.start() - camera.join() - -if __name__ == '__main__': - import optparse - import inspect - - import chroma.solids - import chroma.detectors - import chroma.scenes - from chroma.stl import mesh_from_stl - - parser = optparse.OptionParser('%prog filename.stl') - parser.add_option('-b', '--bits', type='int', dest='bits', - help='bits for z-ordering space axes', default=10) - parser.add_option('-r', '--resolution', dest='resolution', - help='specify resolution', default='1024,576') - parser.add_option('--3d', action='store_true', dest='enable3d', - help='enable 3d', default=False) - parser.add_option('--green', action='store_true', dest='green_magenta', - help='3d with green and magenta lenses', default=False) - parser.add_option('-i', dest='io_file', default=None) - 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() in ('.stl', '.bz2'): - obj = mesh_from_stl(args[0]) - - else: - members = dict(inspect.getmembers(chroma.detectors) + inspect.getmembers(chroma.solids) + inspect.getmembers(chroma.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]) - - geometry = build(obj, options.bits) - if options.io_file is not None: - camera = EventViewer(geometry, options.io_file, size=size, enable3d=options.enable3d, green_magenta=options.green_magenta) - else: - camera = Camera(geometry, size, enable3d=options.enable3d, green_magenta=options.green_magenta) - - camera.start() - camera.join() |