updates to tensorflow code and processing additionsHEADmaster

1 files changed, 25 insertions, 3 deletions

diff --git a/Code/avg_runner.py b/Code/avg_runner.py
index 6809187..ed72b63 100644
--- a/Code/avg_runner.py
+++ b/Code/avg_runner.py
@@ -3,7 +3,11 @@ import getopt
 import sys
 import os
 
-from utils import get_train_batch, get_test_batch
+import pprint
+pp = pprint.PrettyPrinter(indent=2)
+
+from glob import glob
+from utils import get_train_batch, get_test_batch, get_all_clips
 import constants as c
 from g_model import GeneratorModel
 from d_model import DiscriminatorModel
@@ -26,8 +30,11 @@ class AVGRunner:
         self.num_steps = num_steps
         self.num_test_rec = num_test_rec
 
-        self.sess = tf.Session()
-        self.summary_writer = tf.train.SummaryWriter(c.SUMMARY_SAVE_DIR, graph=self.sess.graph)
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        self.sess = tf.Session(config = config)
+        #self.sess = tf.Session()
+        self.summary_writer = tf.summary.FileWriter(c.SUMMARY_SAVE_DIR, graph=self.sess.graph)
 
         if c.ADVERSARIAL:
             print 'Init discriminator...'
@@ -97,11 +104,20 @@ class AVGRunner:
         self.g_model.test_batch(
             batch, self.global_step, num_rec_out=self.num_test_rec)
 
+    def process(self):
+        """
+        Process a directory of images using the generator network.
+        """
+        batch = get_all_clips(c.PROCESS_DIR)
+        self.g_model.test_batch(
+            batch, self.global_step, num_rec_out=self.num_test_rec, process_only=True)
+
 
 def usage():
     print 'Options:'
     print '-l/--load_path=    <Relative/path/to/saved/model>'
     print '-t/--test_dir=     <Directory of test images>'
+    print '-p/--process_dir=  <Directory to process>'
     print '-r/--recursions=   <# recursive predictions to make on test>'
     print '-a/--adversarial=  <{t/f}> (Whether to use adversarial training. Default=True)'
     print '-n/--name=         <Subdirectory of ../Data/Save/*/ in which to save output of this run>'
@@ -123,6 +139,7 @@ def main():
 
     load_path = None
     test_only = False
+    process_only = False
     num_test_rec = 1  # number of recursive predictions to make on test
     num_steps = 1000001
     try:
@@ -140,6 +157,9 @@ def main():
             load_path = arg
         if opt in ('-t', '--test_dir'):
             c.set_test_dir(arg)
+        if opt in ('-p', '--process_dir'):
+            c.set_process_dir(arg)
+            process_only = True
         if opt in ('-r', '--recursions'):
             num_test_rec = int(arg)
         if opt in ('-a', '--adversarial'):
@@ -177,6 +197,8 @@ def main():
     runner = AVGRunner(num_steps, load_path, num_test_rec)
     if test_only:
         runner.test()
+    elif process_only:
+        runner.process()
     else:
         runner.train()