diff options
Diffstat (limited to 'camera.py')
| -rwxr-xr-x | camera.py | 99 |
1 files changed, 98 insertions, 1 deletions
@@ -5,11 +5,14 @@ from threading import Thread, Lock 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 chroma.gpu import GPU, CUDAFuncs from chroma.tools import timeit from chroma.transform import rotate +from chroma.optics import vacuum, lambertian_surface import pygame from pygame.locals import * @@ -22,6 +25,65 @@ from subprocess import call import shutil import tempfile +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): + """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__) + + if isinstance(obj, Geometry): + geometry = obj + elif isinstance(obj, Solid): + geometry = Geometry() + geometry.add_solid(obj) + elif isinstance(obj, Mesh): + geometry = Geometry() + geometry.add_solid(Solid(obj, vacuum, vacuum, surface=lambertian_surface, color=0x99ffffff)) + else: + raise Exception('cannot build type %s' % type(obj)) + + geometry.build(bits) + + return geometry + +def bvh_mesh(geometry, layer): + lower_bounds = geometry.lower_bounds[geometry.layers == layer] + upper_bounds = geometry.upper_bounds[geometry.layers == layer] + + if len(lower_bounds) == 0 or len(upper_bounds) == 0: + raise Exception('no nodes at layer %i' % layer) + + dx, dy, dz = upper_bounds[0] - lower_bounds[0] + center = np.mean([upper_bounds[0],lower_bounds[0]], axis=0) + + 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:]: + mesh += make.box(dx,dy,dz,center) + + return mesh + def encode_movie(dir): root, ext = 'movie', 'avi' for i in count(): @@ -59,6 +121,10 @@ class Camera(Thread): self.device_id = device_id self.size = size + self.unique_bvh_layers = np.unique(self.geometry.layers) + self.currentlayer = None + self.bvh_layers = {self.currentlayer : self.geometry} + try: import spnav as spnav_module self.spnav_module = spnav_module @@ -231,6 +297,16 @@ class Camera(Thread): self.screenshot(self.movie_dir, self.movie_index) self.movie_index += 1 + def loadlayer(self, layer): + try: + layergeometry = self.bvh_layers[layer] + except KeyError: + layergeometry = build(bvh_mesh(self.geometry, layer),8) + self.bvh_layers[layer] = layergeometry + + self.gpu.load_geometry(layergeometry) + self.update() + def process_event(self, event): if event.type == MOUSEBUTTONDOWN: if event.button == 4: @@ -310,6 +386,28 @@ class Camera(Thread): self.done = True return + elif event.key == K_PAGEDOWN: + if self.currentlayer is None: + self.currentlayer = self.unique_bvh_layers[-1] + else: + if self.currentlayer > np.min(self.unique_bvh_layers): + self.currentlayer -= 1 + else: + self.currentlayer = None + + self.loadlayer(self.currentlayer) + + elif event.key == K_PAGEUP: + if self.currentlayer is None: + self.currentlayer = self.unique_bvh_layers[0] + else: + if self.currentlayer < np.max(self.unique_bvh_layers): + self.currentlayer += 1 + else: + self.currentlayer = None + + self.loadlayer(self.currentlayer) + elif event.key == K_F12: self.screenshot() @@ -487,7 +585,6 @@ if __name__ == '__main__': import chroma.detectors import chroma.scenes from chroma.stl import mesh_from_stl - from chroma.view import build parser = optparse.OptionParser('%prog filename.stl') parser.add_option('-b', '--bits', type='int', dest='bits', |