Modernize the runner; add preview support; etc.

main.py

@@ -11,16 +11,14 @@
 
 import os
 import sys
+import argparse
+import multiprocessing
 from subprocess import Popen
-
-from pprint import pprint
 from ctypes import *
 
 import numpy as np
-np.set_printoptions(precision=5, edgeitems=20)
+import Image
 import scipy
-
-import pyglet
 import pycuda.autoinit
 
 import cuburn._pyflam3_hacks
@@ -28,65 +26,198 @@ from fr0stlib import pyflam3
 from cuburn.render import *
 from cuburn.code.mwc import MWCTest
 
-# Required on my system; CUDA doesn't yet work with GCC 4.5
-os.environ['PATH'] = ('/usr/x86_64-pc-linux-gnu/gcc-bin/4.4.5:'
-                     + os.environ['PATH'])
+np.set_printoptions(precision=5, edgeitems=20)
+
+def fmt_time(time):
+    # Stupid precision modifier garbage
+    frac = np.round((time - np.floor(time)) * 1000)
+    if frac:
+        return '%05d.%03d' % (np.floor(time), frac)
+    return '%05d' % time
+
+def save(args, time, raw):
+    name = os.path.join(args.dir, '%s_%s' % (args.name, fmt_time(time)))
+    noalpha = raw[:,:,:3]
+    img = scipy.misc.toimage(noalpha, cmin=0, cmax=1)
+    img.save(name+'.png')
+
+    if args.jpg is not None:
+        img.save(name+'.jpg', quality=args.jpg)
+    print 'saved', name
+
+def error(msg):
+    print "Error:", msg
+    sys.exit(1)
 
 def main(args):
-    if '-t' in args:
+    if args.test:
         MWCTest.test_mwc()
         return
 
-    with open(args[1]) as fp:
-        genome_ptr, ngenomes = pyflam3.Genome.from_string(fp.read())
-        genomes = cast(genome_ptr, POINTER(pyflam3.Genome*ngenomes)).contents
-    anim = Animation(genomes)
-    if '-g' in args:
-        anim.compile(keep=True,
-                 cmp_options=('-use_fast_math', '-maxrregcount', '32', '-G'))
+    genome_ptr, ngenomes = pyflam3.Genome.from_string(args.flame.read())
+    genomes = cast(genome_ptr, POINTER(pyflam3.Genome*ngenomes)).contents
+
+    if not np.all([g.width == genomes[0].width for g in genomes]):
+        error("Inconsistent width. Force with --width.")
+    if not np.all([g.height == genomes[0].height for g in genomes]):
+        error("Inconsistent height. Force with --height.")
+
+    if args.qs:
+        for g in genomes:
+            g.sample_density *= args.qs
+            g.ntemporal_samples = max(1, int(g.ntemporal_samples * args.qs))
+
+    if args.width:
+        if not args.scale:
+            args.scale = float(args.width) / genomes[0].width
+        for g in genomes:
+            g.width = args.width
+    if args.height:
+        for g in genomes:
+            g.height = args.height
+    if args.scale:
+        for g in genomes:
+            g.pixels_per_unit *= args.scale
+
+    if args.skip and not args.tempscale:
+        args.tempscale = float(args.skip)
+    if args.tempscale:
+        for g in genomes:
+            g.temporal_filter_width *= args.tempscale
+
+    if not args.name:
+        if args.flame.name == '<stdin>':
+            args.name = genomes[0].flame_name
+        else:
+            args.name = os.path.splitext(os.path.basename(args.flame.name))[0]
+
+    cp_times = [cp.time for cp in genomes]
+    if args.renumber:
+        for t, g in enumerate(genomes):
+            g.time = t - args.renumber
+        cp_times = [cp.time for cp in genomes]
+    elif np.any(np.diff(cp_times) <= 0):
+        error("Genome times are non-monotonic. Try using --renumber.")
+
+    if args.skip:
+        if args.start is None:
+            args.start = cp_times[0]
+        if args.end is None:
+            args.end = cp_times[-1]
+        times = np.arange(args.start, args.end, args.skip)
     else:
-        anim.compile(keep='-k' in args)
+        times = [t for t in cp_times
+                 if (args.start is None or t >= args.start)
+                 and (args.end is None or t < args.end)]
+
+    anim = Animation(genomes)
+    if args.debug:
+        anim.cmp_options.append('-G')
+    anim.keep = args.keep or args.debug
+    anim.compile()
     anim.load()
-    for n, out in enumerate(anim.render_frames()):
-        noalpha = np.delete(out, 3, axis=2)
-        name = 'rendered_%05d' % n
-        scipy.misc.toimage(noalpha, cmin=0, cmax=1).save(name+'.png')
 
-        if '-j' in args:
-            # Convert using imagemagick, to set custom quality
-            Popen(['convert', name+'.png', '-quality', '90', name+'.jpg'])
-        print 'saved', name, np.min(noalpha), np.max(noalpha)
-    return
+    if args.gfx:
+        import pyglet
+        window = pyglet.window.Window(anim.features.width, anim.features.height)
+        image = pyglet.image.CheckerImagePattern().create_image(
+                anim.features.width, anim.features.height)
+        label = pyglet.text.Label('Rendering first frame', x=5, y=5,
+                font_size=24, bold=True)
 
-    #if '-g' not in args:
-    #    return
+        @window.event
+        def on_draw():
+            print 'redrawing'
+            window.clear()
+            image.texture.blit(0, 0)
+            label.draw()
 
-    window = pyglet.window.Window(anim.features.width, anim.features.height)
-    imgbuf = (np.minimum(accum * 255, 255)).astype(np.uint8)
-    image = pyglet.image.ImageData(anim.features.width, anim.features.height,
-                                   'RGBA', imgbuf.tostring(),
-                                   -anim.features.width * 4)
-    tex = image.texture
+        @window.event
+        def on_key_press(sym, mod):
+            if sym == pyglet.window.key.Q:
+                pyglet.app.exit()
 
-    #pal = (anim.ctx.ptx.instances[PaletteLookup].pal * 255.).astype(np.uint8)
-    #image2 = pyglet.image.ImageData(256, 16, 'RGBA', pal.tostring())
+        @window.event
+        def on_mouse_motion(x, y, dx, dy):
+            pass
 
-    @window.event
-    def on_draw():
-        window.clear()
-        tex.blit(0, 0)
-        #image2.blit(0, 0)
+        frames = anim.render_frames(times, block=False)
+        def poll(dt):
+            out = next(frames, False)
+            if out is False:
+                label.text = "Done. ('q' to quit)"
+                pyglet.clock.unschedule(poll)
+            elif out is not None:
+                time, buf = out
+                save(args, time, buf)
+                imgbuf = np.uint8(buf.flatten() * 255)
+                image.set_data('RGBA', -anim.features.width*4, imgbuf.tostring())
+                label.text = '%s %4g' % (args.name, time)
 
-    @window.event
-    def on_key_press(sym, mod):
-        if sym == pyglet.window.key.Q:
-            pyglet.app.exit()
+        pyglet.clock.set_fps_limit(30)
+        pyglet.clock.schedule_interval(poll, 1/30.)
+        pyglet.app.run()
 
-    pyglet.app.run()
+    else:
+        for time, out in anim.render_frames(times):
+            save(args, time, out)
 
 if __name__ == "__main__":
-    if len(sys.argv) < 2 or not os.path.isfile(sys.argv[1]):
-        print "Last argument must be a path to a genome file"
-        sys.exit(1)
-    main(sys.argv)
+    parser = argparse.ArgumentParser(description='Render fractal flames.')
+
+    parser.add_argument('flame', metavar='FILE', type=file,
+        help="Path to genome file ('-' for stdin)")
+    parser.add_argument('-g', action='store_true', dest='gfx',
+        help="Show output in OpenGL window")
+    parser.add_argument('-j', metavar='QUALITY', nargs='?',
+        action='store', type=int, dest='jpg', const=90,
+        help="Write .jpg in addition to .png (default quality 90)")
+    parser.add_argument('-n', metavar='NAME', type=str, dest='name',
+        help="Prefix to use when saving files (default is basename of input)")
+    parser.add_argument('-o', metavar='DIR', type=str, dest='dir',
+        help="Output directory", default='.')
+
+    time = parser.add_argument_group('Sequence options', description="""
+        Control which frames are rendered from a genome sequence. If '-k' is
+        not given, '-s' and '-e' act as limits, and any control point with a
+        time in bounds is rendered at its central time. If '-k' is given,
+        a list of times to render is given according to the semantics of
+        Python's range operator, as in range(start, end, skip).
+
+        If no options are given, all control points are rendered.""")
+    time.add_argument('-s', dest='start', metavar='TIME', type=float,
+        help="Start time of image sequence (inclusive, clamped)")
+    time.add_argument('-e', dest='end', metavar='TIME', type=float,
+        help="End time of image sequence (exclusive, clamped)")
+    time.add_argument('-k', dest='skip', metavar='TIME', type=float,
+        help="Skip time between frames in image sequence. Auto-sets "
+             "--tempscale, use '--tempscale 1' to override.")
+    time.add_argument('--renumber', metavar="START", type=float,
+        dest='renumber', nargs='?', const=0,
+        help="Renumber frame times monotonically, with the first frame at time "
+             "START. Skip is 1 and default start is 0.")
+
+    genome = parser.add_argument_group('Genome options')
+    genome.add_argument('--qs', type=float,
+        help="Scale quality and number of temporal samples by this factor")
+    genome.add_argument('--scale', type=float,
+        help="Scale pixels per unit (camera zoom) by this factor")
+    genome.add_argument('--width', type=int,
+        help="Use this width. Auto-sets scale, use '--scale 1' to override.")
+    genome.add_argument('--height', type=int,
+        help="Use this height (does *not* auto-set scale)")
+    genome.add_argument('--tempscale', type=float,
+        help="Scale temporal filter width by this factor")
+
+    debug = parser.add_argument_group('Debug options')
+    debug.add_argument('--test', action='store_true', dest='test',
+        help='Run some internal tests')
+    debug.add_argument('--keep', action='store_true', dest='keep',
+        help='Keep compilation directory (disables kernel caching)')
+    debug.add_argument('--debug', action='store_true', dest='debug',
+        help='Compile kernel with debugging enabled (implies --keep)')
+
+    args = parser.parse_args()
+
+    main(args)