Merge speedup of build()

1 files changed, 21 insertions, 16 deletions

diff --git a/view.py b/view.py
index 730f2c9..952964f 100755
--- a/view.py
+++ b/view.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 import os
+import sys
 import numpy as np
 
 import pygame
@@ -16,9 +17,9 @@ from pycuda import autoinit
 from pycuda.compiler import SourceModule
 from pycuda import gpuarray
 
-def view(geometry, name=''):
+def view(viewable, name='', bits=8):
     """
-    Render `geometry` in a pygame window.
+    Render `viewable` in a pygame window.
 
     Movement:
         - zoom: scroll the mouse wheel
@@ -26,6 +27,20 @@ def view(geometry, name=''):
         - move: shift+click and drag the mouse
     """
 
+    if isinstance(viewable, Geometry):
+        geometry = viewable
+        geometry.build(bits)
+    elif isinstance(viewable, Solid):
+        geometry = Geometry()
+        geometry.add_solid(viewable)
+        geometry.build(bits)
+    elif isinstance(viewable, Mesh):
+        geometry = Geometry()
+        geometry.add_solid(Solid(viewable))
+        geometry.build(bits)
+    else:
+        sys.exit("can't display %s" % args[0])
+
     lower_bound = np.array([np.min(geometry.mesh[:][:,:,0]),
                             np.min(geometry.mesh[:][:,:,1]),
                             np.min(geometry.mesh[:][:,:,2])])
@@ -192,12 +207,12 @@ def view(geometry, name=''):
                 if event.key == K_LSHIFT or event.key == K_RSHIFT:
                     shift = False
 
+    pygame.display.quit()
+
 if __name__ == '__main__':
-    import os
-    import sys
     import optparse
 
-    from stl import *
+    from stl import mesh_from_stl
 
     import solids
     import detectors
@@ -224,16 +239,6 @@ if __name__ == '__main__':
         members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids))
 
         if args[0] in members:
-            if isinstance(members[args[0]], Geometry):
-                geometry = members[args[0]]
-                geometry.build(options.bits)
-                view(geometry, args[0])
-            elif isinstance(members[args[0]], Solid):
-                geometry = Geometry()
-                geometry.add_solid(members[args[0]])
-                geometry.build(options.bits)
-                view(geometry, args[0])
-            else:
-                sys.exit("can't display %s" % args[0])
+            view(members[args[0]], args[0], options.bits)
         else:
             sys.exit("couldn't find object %s" % args[0])