Add another loader function that is useful for chroma-cam.

Note that rendering is still broken by the new BVH format.

3 files changed, 63 insertions, 49 deletions

diff --git a/bin/chroma-cam b/bin/chroma-cam
index c75f28f..044cf31 100755
--- a/bin/chroma-cam
+++ b/bin/chroma-cam
@@ -6,9 +6,10 @@ if __name__ == '__main__':
     import sys
     import os
 
-    from chroma import view, build, EventViewer
+    from chroma import view, EventViewer
     from chroma import mesh_from_stl
     from chroma.tools import enable_debug_on_crash
+    import chroma.loader
     from chroma.log import logger, logging
     logger.setLevel(logging.INFO)
 
@@ -28,24 +29,10 @@ if __name__ == '__main__':
 
     size = [int(s) for s in options.resolution.split(',')]
 
-    if os.path.exists(args[0]) and args[0].lower().endswith(('.stl', '.bz2')):
-        obj = mesh_from_stl(args[0])
-    else:
-        module_name, function_name = args[0].rsplit('.', 1)
-        orig_sys_path = list(sys.path)
-        try:
-            sys.path.append('.')
-            module = __import__(module_name, fromlist=[function_name])
-            sys.path = orig_sys_path
-        except ImportError:
-            sys.path = orig_sys_path
-            raise
-
-        obj = getattr(module, function_name)
+    geometry = chroma.loader.load_geometry_from_string(args[0])
 
     if options.io_file is None:
-        view(obj, size)
+        view(geometry, size)
     else:
-        geometry = build(obj)
         viewer = EventViewer(geometry, options.io_file, size=size)
         viewer.start()
diff --git a/chroma/camera.py b/chroma/camera.py
index 09463d1..f23085f 100644
--- a/chroma/camera.py
+++ b/chroma/camera.py
@@ -19,30 +19,10 @@ from chroma.sample import uniform_sphere
 from chroma.project import from_film
 from chroma import make
 from chroma import gpu
-
+from chroma.loader import create_geometry_from_obj
 import pygame
 from pygame.locals import *
 
-def build(obj):
-    """Construct and build a geometry from `obj`."""
-    if callable(obj):
-        obj = obj()
-
-    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, color=0x33ffffff))
-    else:
-        raise TypeError('cannot build type %s' % type(obj))
-
-    geometry.build()
-
-    return geometry
-
 def bvh_mesh(geometry, layer):
     lower_bounds = geometry.lower_bounds[geometry.layers == layer]
     upper_bounds = geometry.upper_bounds[geometry.layers == layer]
@@ -84,7 +64,7 @@ class Camera(multiprocessing.Process):
         self.device_id = device_id
         self.size = size
 
-        self.unique_bvh_layers = np.unique(self.geometry.layers)
+        self.bvh_layer_count = len(self.geometry.bvh.layer_offsets)
         self.currentlayer = None
         self.bvh_layers = {}
 
@@ -437,7 +417,7 @@ class Camera(multiprocessing.Process):
             try:
                 gpu_geometry = self.bvh_layers[layer]
             except KeyError:
-                geometry = build(bvh_mesh(self.geometry, layer))
+                geometry = create_geometry_from_obj(bvh_mesh(self.geometry, layer))
                 gpu_geometry = gpu.GPUGeometry(geometry, print_usage=False)
                 self.bvh_layers[layer] = gpu_geometry
 
@@ -528,9 +508,9 @@ class Camera(multiprocessing.Process):
 
             elif event.key == K_PAGEDOWN:
                 if self.currentlayer is None:
-                    self.currentlayer = self.unique_bvh_layers[-1]
+                    self.currentlayer = self.bvh_layer_count - 1
                 else:
-                    if self.currentlayer > np.min(self.unique_bvh_layers):
+                    if self.currentlayer > 0:
                         self.currentlayer -= 1
                     else:
                         self.currentlayer = None
@@ -539,9 +519,9 @@ class Camera(multiprocessing.Process):
 
             elif event.key == K_PAGEUP:
                 if self.currentlayer is None:
-                    self.currentlayer = self.unique_bvh_layers[0]
+                    self.currentlayer = 0
                 else:
-                    if self.currentlayer < np.max(self.unique_bvh_layers):
+                    if self.currentlayer < self.bvh_layer_count:
                         self.currentlayer += 1
                     else:
                         self.currentlayer = None
@@ -703,7 +683,7 @@ class EventViewer(Camera):
         for pos in self.ev.photons_beg.pos[::100]:
             geometry.add_solid(marker, displacement=pos, rotation=make_rotation_matrix(np.random.uniform(0,2*np.pi), uniform_sphere()))
 
-        geometry.build(bits=11)
+        geometry = create_geometry_from_obj(geometry)
         gpu_geometry = gpu.GPUGeometry(geometry)
 
         self.gpu_geometries = [self.gpu_geometry, gpu_geometry]
@@ -819,6 +799,6 @@ class EventViewer(Camera):
         Camera.process_event(self, event)
 
 def view(obj, size=(800,600), **camera_kwargs):
-    geometry = build(obj)
+    geometry = create_geometry_from_obj(obj)
     camera = Camera(geometry, size, **camera_kwargs)
     camera.start()
diff --git a/chroma/loader.py b/chroma/loader.py
index fc98728..39842e1 100644
--- a/chroma/loader.py
+++ b/chroma/loader.py
@@ -5,7 +5,8 @@ import time
 from chroma.log import logger
 from chroma.cache import Cache
 from chroma.bvh import make_simple_bvh
-from chroma.geometry import Geometry, Solid
+from chroma.geometry import Geometry, Solid, Mesh
+from chroma.detector import Detector
 from chroma.stl import mesh_from_stl
 from chroma.gpu import create_cuda_context
 
@@ -112,7 +113,24 @@ def load_geometry_from_string(geometry_str,
             geometry = cache.load_geometry(geometry_id)
         # Cached geometries are flattened already
 
-    # Figure out the BVH situation
+
+    geometry.bvh = load_bvh(geometry, auto_build_bvh=auto_build_bvh,
+                            read_bvh_cache=read_bvh_cache,
+                            update_bvh_cache=update_bvh_cache,
+                            cache_dir=cache_dir,
+                            cuda_device=cuda_device)
+
+    return geometry
+
+def load_bvh(geometry,  bvh_name="default", 
+             auto_build_bvh=True, read_bvh_cache=True,
+             update_bvh_cache=True, cache_dir=None,
+             cuda_device=None):
+    if cache_dir is None:
+        cache = Cache()
+    else:
+        cache = Cache(cache_dir)
+
     mesh_hash = geometry.mesh.md5()
     bvh = None
     if read_bvh_cache and cache.exist_bvh(mesh_hash, bvh_name):
@@ -133,5 +151,34 @@ def load_geometry_from_string(geometry_str,
             logger.info('Saving BVH (%s:%s) to cache.' % (mesh_hash, bvh_name))
             cache.save_bvh(bvh, mesh_hash, bvh_name)
 
-    geometry.bvh = bvh
+    return bvh
+
+def create_geometry_from_obj(obj, bvh_name="default", 
+                             auto_build_bvh=True, read_bvh_cache=True,
+                             update_bvh_cache=True, cache_dir=None,
+                             cuda_device=None):
+    if callable(obj):
+        obj = obj()
+
+    if isinstance(obj, Detector):
+        geometry = obj
+    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, color=0x33ffffff))
+    else:
+        raise TypeError('cannot build type %s' % type(obj))
+
+    geometry.flatten()
+
+    geometry.bvh = load_bvh(geometry, auto_build_bvh=auto_build_bvh,
+                            read_bvh_cache=read_bvh_cache,
+                            update_bvh_cache=update_bvh_cache,
+                            cache_dir=cache_dir,
+                            cuda_device=cuda_device)
+
     return geometry