path: root/render.py

#!/usr/bin/env python
from view import *
from pycuda.characterize import sizeof
import tempfile

def render(viewable, size=(800,600), name='', bits=8, make_movie=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
    """

    geometry = build(viewable, bits)

    lower_bound, upper_bound = geometry.mesh.get_bounds()

    scale = np.linalg.norm(upper_bound-lower_bound)

    print 'device %s' % autoinit.device.name()

    module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)
    geometry.load(module, color=False)
    cuda_build_rgb_lookup = module.get_function('build_rgb_lookup')
    cuda_render = module.get_function('render')
    cuda_rotate = module.get_function('rotate')
    cuda_translate = module.get_function('translate')

    pygame.init()
    width, height = size

    init_rng = module.get_function('init_rng')

    rng_state_count = max(len(geometry.mesh.triangles), width*height)
    rng_states_gpu = cuda.mem_alloc(rng_state_count*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
    init_rng(np.int32(rng_state_count), rng_states_gpu, np.int32(0), np.int32(0), block=(64,1,1), grid=(rng_state_count//64+1,1))

    diagonal = np.linalg.norm(upper_bound-lower_bound)

    source_position = [0, 0, upper_bound[2]+1.0]

    source_positions_gpu = gpuarray.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.float3)
    source_positions_gpu.fill(gpuarray.vec.make_float3(*source_position))

    source_directions = np.mean(geometry.mesh[:], axis=1) - source_position
    source_directions_gpu = gpuarray.to_gpu(source_directions.astype(np.float32).view(gpuarray.vec.float3))

    rgb_lookup_gpu = gpuarray.zeros(source_positions_gpu.size, dtype=gpuarray.vec.float3)

    max_steps = 10
    rgb_runs = 100

    print 'building rgb lookup.'
    cuda_build_rgb_lookup(np.int32(source_positions_gpu.size), rng_states_gpu, source_positions_gpu, source_directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), rgb_lookup_gpu, np.int32(rgb_runs), np.int32(max_steps), block=(64,1,1), grid=(source_positions_gpu.size//64+1,1))
    cuda.Context.synchronize()
    print 'done.'

    rgb_lookup = rgb_lookup_gpu.get()
    rgb_lookup['x'] /= rgb_runs
    rgb_lookup['y'] /= rgb_runs
    rgb_lookup['z'] /= rgb_runs
    rgb_lookup['x'][rgb_lookup['x'] > 1.0] = 1.0
    rgb_lookup['y'][rgb_lookup['y'] > 1.0] = 1.0
    rgb_lookup['z'][rgb_lookup['z'] > 1.0] = 1.0
    rgb_lookup_gpu = cuda.to_device(rgb_lookup)

    camera = Camera(size)

    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 = camera.get_rays()

    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)

    nruns = 5

    screen = pygame.display.set_mode(size)
    pygame.display.set_caption(name)

    nblocks = 64

    if make_movie:
        image_index = 0
        tempdir = tempfile.mkdtemp()

    def update():
        """Render the mesh and display to screen."""

        global image_index

        t0 = time.time()
        cuda_render(np.int32(pixels_gpu.size), rng_states_gpu, origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), rgb_lookup_gpu, np.int32(nruns), pixels_gpu, np.int32(max_steps), 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()

    if make_movie:
        screenshot(screen, name, tempdir, image_index)
        image_index += 1

    done = False
    clicked = False
    shift = False

    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 make_movie:
                        screenshot(screen, name, tempdir, image_index)
                        image_index += 1

                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 make_movie:
                        screenshot(screen, name, tempdir, image_index)
                        image_index += 1

                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()

                    if make_movie:
                        screenshot(screen, name, tempdir, image_index)
                        image_index += 1

                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 make_movie:
                        screenshot(screen, name, tempdir, image_index)
                        image_index += 1

            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_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 make_movie:
        if name == '':
            root, ext = 'movie', 'avi'
        else:
            root, ext = name + '_movie', 'avi'

        filename = '.'.join([root, ext])

        i=1
        while os.path.exists(filename):
            filename = '.'.join([root + str(i), ext])
            i += 1

        from subprocess import call

        call(['mencoder', 'mf://' + tempdir + '/*.png', '-mf', 'fps=10', '-o', filename, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000'])

        import shutil

        shutil.rmtree(tempdir)

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('-r', '--resolution', dest='resolution',
                      help='specify resolution', default='800,600')
    parser.add_option('-m', '--movie', action='store_true', dest='movie',
                      help='create a movie', default=False)
    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() == '.stl':
            render(mesh_from_stl(args[0]), size, root, options.bits, options.movie)
    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:
            render(buildable_lookup[args[0]], size, args[0], options.bits, options.movie)
        else:
            sys.exit("couldn't find object %s" % args[0])