Oops, view.py is still in use.

1 files changed, 375 insertions, 0 deletions

diff --git a/view.py b/view.py
new file mode 100755
index 0000000..b10fbda
--- /dev/null
+++ b/view.py
@@ -0,0 +1,375 @@
+#!/usr/bin/env python
+import os
+import sys
+import time
+import numpy as np
+import inspect
+
+import pygame
+from pygame.locals import *
+
+import src
+from camera import get_rays
+from geometry import Mesh, Solid, Geometry
+from transform import rotate
+from optics import *
+from gpu import *
+
+#from pycuda import autoinit
+from pycuda.compiler import SourceModule
+from pycuda import gpuarray
+import pycuda.driver as cuda
+
+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 screenshot(screen, name='', dir='', index=0):
+    """Take a screenshot of `screen`."""
+    if name == '':
+        root, ext = 'screenshot', 'png'
+    else:
+        root, ext = name, 'png'
+
+    i = index
+
+    filename = os.path.join(dir,'.'.join([root + str(i).zfill(4), ext]))
+    while os.path.exists(filename):
+        filename = os.path.join(dir,'.'.join([root + str(i).zfill(4), ext]))
+        i += 1
+
+    pygame.image.save(screen, filename)
+    print 'image saved to %s' % filename
+
+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 box(lower_bound, upper_bound):
+    """
+    Return a mesh of the box defined by the opposing corners `lower_bound`
+    and `upper_bound`.
+    """
+    dx, dy, dz = upper_bound - lower_bound
+
+    vertices = np.array([lower_bound,
+                         lower_bound + (dx,0,0),
+                         lower_bound + (dx,dy,0),
+                         lower_bound + (0,dy,0),
+                         lower_bound + (0,0,dz),
+                         lower_bound + (dx,0,dz),
+                         lower_bound + (dx,dy,dz),
+                         lower_bound + (0,dy,dz)])
+
+    triangles = np.empty((12,3), dtype=np.int32)
+    # bottom
+    triangles[0] = (1,3,2)
+    triangles[1] = (1,4,3)
+    # top
+    triangles[2] = (5,7,6)
+    triangles[3] = (5,8,7)
+    # left
+    triangles[4] = (5,1,2)
+    triangles[5] = (5,2,6)
+    # right
+    triangles[6] = (3,4,8)
+    triangles[7] = (3,8,7)
+    # front
+    triangles[8] = (2,3,7)
+    triangles[9] = (2,7,6)
+    # back
+    triangles[10] = (1,5,8)
+    triangles[11] = (1,8,4)
+
+    triangles -= 1
+
+    return Mesh(vertices, triangles)
+
+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)
+
+    mesh = box(lower_bounds[0], upper_bounds[0])
+
+    for lower_bound, upper_bound in zip(lower_bounds, upper_bounds)[1:]:
+        mesh += box(lower_bound, upper_bound)
+
+    return mesh
+
+def view(viewable, size=(800,600), name='', bits=8, load_bvh=False):
+    """
+    Render `viewable` in a pygame window.
+
+    Movement:
+        - zoom: scroll the mouse wheel
+        - rotate: click and drag the mouse
+        - move: shift+click and drag the mouse
+    """
+
+    gpu = GPU()
+
+    geometry = build(viewable, bits)
+
+    if load_bvh:
+        print 'loading bvh...'
+        bvhg = []
+
+        bvhg.append(geometry)
+
+        for layer in sorted(np.unique(geometry.layers)):
+            print 'building layer %i' % layer
+            bvhg.append(build(bvh_mesh(geometry, layer), bits))
+
+    lower_bound, upper_bound = geometry.mesh.get_bounds()
+
+    scale = np.linalg.norm(upper_bound-lower_bound)
+
+    #from pycuda import autoinit
+
+    #print 'device %s' % autoinit.device.name()
+
+    module = gpu.module
+
+    #module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True)#, cache_dir=False)
+    #geometry.load(module)
+    gpu.load_geometry(geometry)
+    cuda_raytrace = module.get_function('ray_trace')
+    cuda_rotate = module.get_function('rotate')
+    cuda_translate = module.get_function('translate')
+
+    pygame.init()
+    width, height = size
+    screen = pygame.display.set_mode(size)
+    pygame.display.set_caption(name)
+
+    #camera = Camera(size)
+
+    diagonal = np.linalg.norm(upper_bound-lower_bound)
+
+    point = np.array([0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2])
+    axis1 = np.array([1,0,0], dtype=np.double)
+    axis2 = np.array([0,0,1], dtype=np.double)
+
+    #camera.position(point)
+
+    origins, directions = get_rays(point, size)
+
+    origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3))
+
+    directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3))
+
+    pixels_gpu = gpuarray.zeros(width*height, dtype=np.int32)
+
+    nblocks = 64
+
+    def update():
+        """Render the mesh and display to screen."""
+        t0 = time.time()
+        cuda_raytrace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, pixels_gpu, block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
+        cuda.Context.synchronize()
+        elapsed = time.time() - t0
+
+        print 'elapsed %f sec' % elapsed
+
+        pygame.surfarray.blit_array(screen, pixels_gpu.get().reshape(size))
+        pygame.display.flip()
+
+    update()
+
+    done = False
+    clicked = False
+    shift = False
+
+    current_layer = 0
+
+    while not done:
+        for event in pygame.event.get():
+            if event.type == MOUSEBUTTONDOWN:
+                if event.button == 4:
+                    v = scale*np.cross(axis1,axis2)/10.0
+
+                    cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
+
+                    point += v
+
+                    update()
+
+                if event.button == 5:
+                    v = -scale*np.cross(axis1,axis2)/10.0
+
+                    cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
+
+                    point += v
+
+                    update()
+
+                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
+
+                length = np.linalg.norm(movement)
+
+                mouse_direction = movement[0]*axis1 + movement[1]*axis2
+                mouse_direction /= np.linalg.norm(mouse_direction)
+
+                if shift:
+                    v = mouse_direction*scale*length/float(width)
+
+                    cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
+
+                    point += v
+
+                    update()
+                else:
+                    phi = np.float32(2*np.pi*length/float(width))
+                    n = rotate(mouse_direction, np.pi/2, \
+                                   -np.cross(axis1,axis2))
+
+                    cuda_rotate(np.int32(pixels_gpu.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
+
+                    cuda_rotate(np.int32(pixels_gpu.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
+
+                    point = rotate(point, phi, n)
+                    axis1 = rotate(axis1, phi, n)
+                    axis2 = rotate(axis2, phi, n)
+
+                    update()
+
+            if event.type == KEYDOWN:
+                if event.key == K_LSHIFT or event.key == K_RSHIFT:
+                    shift = 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
+
+                        gpu.load_geometry(geometry)
+                        #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
+
+                        gpu.load_geometry(geometry)
+                        #geometry.load(module, color=True)
+                        update()
+                    except IndexError:
+                        print 'no further layers to view'
+
+                if event.key == K_F12:
+                    screenshot(screen, name)
+
+            if event.type == KEYUP:
+                if event.key == K_LSHIFT or event.key == K_RSHIFT:
+                    shift = False
+
+            if event.type == pygame.QUIT:
+                done = True
+                break
+
+    pygame.display.quit()
+
+if __name__ == '__main__':
+    import optparse
+
+    from stl import mesh_from_stl
+
+    import solids
+    import detectors
+    import scenes
+
+    parser = optparse.OptionParser('%prog filename.stl')
+    parser.add_option('-b', '--bits', type='int', dest='bits',
+                      help='bits for z-ordering space axes', default=8)
+    parser.add_option('-l', action='store_true', dest='load_bvh',
+                      help='load bounding volumes', default=False)
+    parser.add_option('-r', '--resolution', dest='resolution',
+                      help='specify resolution', default='800,600')
+    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]):
+        head, tail = os.path.split(args[0])
+        root, ext = os.path.splitext(tail)
+
+        if ext.lower() in ('.stl', '.bz2'):
+            view(mesh_from_stl(args[0]), size, root, options.bits, options.load_bvh)
+    else:
+        members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(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:
+            view(buildable_lookup[args[0]], size, args[0], options.bits, options.load_bvh)
+        else:
+            sys.exit("couldn't find object %s" % args[0])