modified stylize, removed examples

RobGeada · Nov 20, 2017 · e1cacb5 · e1cacb5
1 parent b746181
commit e1cacb5
Show file tree

Hide file tree

Showing 16 changed files with 16 additions and 0 deletions.
diff --git a/.DS_Store b/.DS_Store
diff --git a/examples/1-content.jpg b/examples/1-content.jpg
diff --git a/examples/1-output.jpg b/examples/1-output.jpg
diff --git a/examples/1-style.jpg b/examples/1-style.jpg
diff --git a/examples/2-content.jpg b/examples/2-content.jpg
diff --git a/examples/2-output.jpg b/examples/2-output.jpg
diff --git a/examples/2-style1.jpg b/examples/2-style1.jpg
diff --git a/examples/2-style2.jpg b/examples/2-style2.jpg
diff --git a/examples/tweaks/cwe01.jpg b/examples/tweaks/cwe01.jpg
diff --git a/examples/tweaks/cwe10_default.jpg b/examples/tweaks/cwe10_default.jpg
diff --git a/examples/tweaks/swe02.jpg b/examples/tweaks/swe02.jpg
diff --git a/examples/tweaks/swe14_pavg.jpg b/examples/tweaks/swe14_pavg.jpg
diff --git a/examples/tweaks/swe14_pmax.jpg b/examples/tweaks/swe14_pmax.jpg
diff --git a/examples/tweaks/swe14_pmax_pcyuv.jpg b/examples/tweaks/swe14_pmax_pcyuv.jpg
diff --git a/examples/tweaks/swe20.jpg b/examples/tweaks/swe20.jpg
diff --git a/stylize.py b/stylize.py
@@ -5,6 +5,7 @@
 import tensorflow as tf
 import numpy as np
 
+import os
 from sys import stderr
 
 from PIL import Image
@@ -143,6 +144,18 @@ def print_progress():
             stderr.write('Optimization started...\n')
             if (print_iterations and print_iterations != 0):
                 print_progress()
+
+            pretrained = False
+            print os.listdir(str(os.getcwd()))
+            if "fns.ckpt" in os.listdir(str(os.getcwd())):
+                print("Evaluating pre-trained model")
+                pretrained = True
+                saver = tf.train.Saver()
+                ckpt = tf.train.get_checkpoint_state(os.getcwd())
+
+                saver.restore(sess, ckpt.model_checkpoint_path)
+                iterations = 1;
+
             for i in range(iterations):
                 stderr.write('Iteration %4d/%4d\n' % (i + 1, iterations))
                 train_step.run()
@@ -189,6 +202,9 @@ def print_progress():
 
                         # 5
                         img_out = np.array(Image.fromarray(combined_yuv, 'YCbCr').convert('RGB'))
+                    if not pretrained:
+                        saver = tf.train.Saver(write_version=tf.train.SaverDef.V1)
+                        res = saver.save(sess, os.getcwd()+"/fns.ckpt")
 
 
                     yield (