Add an option to save animation images

gte/__main__.py

@@ -7,7 +7,7 @@ import time
 import datetime
 import numpy as np
 from random import randint
-from multiprocessing import Pool
+from multiprocessing import Pool, Process, Queue
 
 logging.basicConfig(level=logging.INFO,
                     format="%(created)d %(asctime)s:%(module)s(%(lineno)d):%(levelname)s:%(message)s")
@@ -103,7 +103,8 @@ parser_gen.add_argument('--gray', '-G', action='store_const',
                         default=False, const=True)
 parser_gen.add_argument('--no-diagonals', '-d', action='store_const',
                         default=False, const=True)
-
+parser_gen.add_argument('--threads', '-t', type=int, metavar='N',
+                        default=2, help="Animation writting threads")
 
 args = parser.parse_args()
 
@@ -221,6 +222,24 @@ if 'pool_size' in args:
     exit(0)
 else:
     # Generate
+    if args.output_dir is not None:
+        #from PIL import Image
+        import scipy
+        def save_anim(data_q):
+            while True:
+                wval, fname = data_q.get()
+                if wval is None:
+                    return
+                #im = Image.fromarray(wval, 'RGB')
+                #im.save(fname)
+                scipy.misc.imsave(fname, wval)
+                logger.info('Written %s / %d' % (fname, args.steps//args.anim_div))
+
+        data_q = Queue()
+        wps = [Process(target=save_anim, args=(data_q,))
+               for _ in range(args.threads)]
+        for wp in wps:
+            wp.start()
     w = World(args.world_height, args.world_width, gray=args.gray)
     prog = rpnlib.RpnExpr.from_string(args.prog)
     turmits = [LivingTurmit(world=w, prog=prog, diag=not args.no_diagonals)
@@ -238,8 +257,20 @@ else:
         for turmit in turmits:
             turmit()
         if args.output_dir is not None and step % args.anim_div == 0:
-            #w.save('/tmp/anim/GTE_ANIM_%d.png' % (step/10))
-            w.save('/tmp/anim/GTE_ANIM_%d.png' % (step/args.anim_div))
+            fname = 'GTE_ANIM_%d.png' % (step/args.anim_div)
+            fname = os.path.join(args.output_dir, fname)
+            data_q.put((w._val, fname))
+        #if args.output_dir is not None and step % args.anim_div == 0:
+        #    #w.save('/tmp/anim/GTE_ANIM_%d.png' % (step/10))
+        #    fname = 'GTE_ANIM_%d?png' % (step/args.anim_div)
+        #    w.save(os.path.join(args.output_dir, fname))
+    
+    if args.output_dir is not None:
+        for _ in wps:
+            data_q.put((None, None))
+        for wp in wps:
+            wp.join()
+        data_q.close()
     stop = time.time()
 
     score_dir = sum([t.dirvar() for t in turmits]) / len(turmits)