added examples and fixed the io for checkpoints

43 files changed, 63 insertions, 26 deletions

diff --git a/__pycache__/hyperparameters.cpython-38.pyc b/__pycache__/hyperparameters.cpython-38.pyc
index 2ff6d7e9..25914e13 100644
--- a/__pycache__/hyperparameters.cpython-38.pyc
+++ b/__pycache__/hyperparameters.cpython-38.pyc
diff --git a/__pycache__/model.cpython-38.pyc b/__pycache__/model.cpython-38.pyc
new file mode 100644
index 00000000..a444f791
--- /dev/null
+++ b/__pycache__/model.cpython-38.pyc
diff --git a/checkpoint-images/img-epoch0.jpg b/checkpoint-images/img-epoch0.jpg
deleted file mode 100644
index 4836b9e8..00000000
--- a/checkpoint-images/img-epoch0.jpg
+++ /dev/null
diff --git a/checkpoint-images/img-epoch100.jpg b/checkpoint-images/img-epoch100.jpg
deleted file mode 100644
index dfcf28e2..00000000
--- a/checkpoint-images/img-epoch100.jpg
+++ /dev/null
diff --git a/checkpoint-images/img-epoch200.jpg b/checkpoint-images/img-epoch200.jpg
deleted file mode 100644
index 13fde882..00000000
--- a/checkpoint-images/img-epoch200.jpg
+++ /dev/null
diff --git a/checkpoint-images/img-epoch300.jpg b/checkpoint-images/img-epoch300.jpg
deleted file mode 100644
index 13fde882..00000000
--- a/checkpoint-images/img-epoch300.jpg
+++ /dev/null
diff --git a/checkpoint-images/img-epoch400.jpg b/checkpoint-images/img-epoch400.jpg
deleted file mode 100644
index 13fde882..00000000
--- a/checkpoint-images/img-epoch400.jpg
+++ /dev/null
diff --git a/checkpoint.npy b/checkpoint.npy
deleted file mode 100644
index 73b09998..00000000
--- a/checkpoint.npy
+++ /dev/null
diff --git a/hyperparameters.py b/hyperparameters.py
index 0ca70764..1d4c0e8d 100644
--- a/hyperparameters.py
+++ b/hyperparameters.py
@@ -18,12 +18,24 @@ optimizer is used. Refer to the default learning rate parameter
 """
 learning_rate = .002
 
+"""
+Beta_1 is the first hyperparameter for the Adam optimizer.
+"""
 beta_1 = .99
 
+"""
+epsilon for the Adam optimizer.
+"""
 epsilon = 1e-1
 
-momentum = 0.01
-
+"""
+A critical parameter for style transfer. The value for this will determine 
+how much the generated image is "influenced" by the CONTENT image.
+"""
 alpha = .05
 
+"""
+A critical parameter for style transfer. The value for this will determine 
+how much the generated image is "influenced" by the STYLE image.
+"""
 beta = 5
diff --git a/main.py b/main.py
index 527394cb..e60f14d1 100644
--- a/main.py
+++ b/main.py
@@ -31,54 +31,79 @@ def parse_args():
         '--savefile',
         required=True,
         help='Filename to save image')
-    parser.add_argument(
-        '--load',
-        required=False,
-        default='N',
-        help='Y if you want to load the most recent weights'
-    )
 
     return parser.parse_args()
 
 
+def save_tensor_as_image(tensor, path, img_name):
+    # make copy of tensor
+    copy = tf.identity(tensor)
+    copy = tf.squeeze(copy)
+    # convert tensor to back to uint8 image
+    copy = tf.image.convert_image_dtype(copy, tf.uint8)
+
+    # save image (make path if it doesn't exist)
+    if not os.path.exists(path):
+        os.makedirs(path)
+    imsave(path + img_name, copy)
+
+    # return copy if used
+    return copy
+
+
 def train(model: YourModel):
+    # do as many epochs from hyperparameters.py
     for i in range(hp.num_epochs):
-        if i % 50 == 0:
-            copy = tf.identity(model.x)
-            copy = tf.squeeze(copy)
-            copy = tf.image.convert_image_dtype(copy, tf.uint8)
-            imsave('save_images/epoch' + str(i) + '.jpg', copy)
-            np.save('checkpoint.npy', model.x)
+        # save a checkpoint every 100 epochs
+        if i % 100 == 0:
+            save_tensor_as_image(model.x, 'out/checkpoints/cp-{}/'.format(ARGS.savefile),
+                                 '{}_epoch.{}'.format(i, ARGS.savefile.split('.')[-1]))
+
+        # do the training step
         model.train_step(i)
 
 
 def main():
     """ Main function. """
+    # --------------------------------------------------------------------------------------------------------------
+    # PART 1 : parse the arguments #
+    # --------------------------------------------------------------------------------------------------------------
     if os.path.exists(ARGS.content):
         ARGS.content = os.path.abspath(ARGS.content)
     if os.path.exists(ARGS.style):
         ARGS.style = os.path.abspath(ARGS.style)
+    if os.path.exists('out/checkpoints/cp-{}/'.format(ARGS.savefile)):
+        # print an error to the console if the checkpoint directory already exists
+        print('Error: out/checkpoints/cp-{}/ already exists. Please choose a different name.'.format(ARGS.savefile))
+        return
     os.chdir(sys.path[0])
-    print('this is', ARGS.content)
 
+    # --------------------------------------------------------------------------------------------------------------
+    # PART 2 : read and process the style and content images #
+    # --------------------------------------------------------------------------------------------------------------
+    # 1) read content and style images
     content_image = imread(ARGS.content)
     style_image = imread(ARGS.style)
+    # 2) make the style image the same size as the content image
     style_image = transform.resize(style_image, content_image.shape, anti_aliasing=True)
 
+    # --------------------------------------------------------------------------------------------------------------
+    # PART 3 : make and train our model #
+    # --------------------------------------------------------------------------------------------------------------
+    # 1) initialize our model class
     my_model = YourModel(content_image=content_image, style_image=style_image)
-
-    if (ARGS.load == 'Y'):
-        checkpoint = np.load('checkpoint.npy')
-        image = tf.Variable(initial_value=checkpoint)
-
+    # 2) train the model calling the helper
     train(my_model)
 
-    # convert the tensor into an image
-    my_model.x = tf.squeeze(my_model.x)
-    final_image = tf.image.convert_image_dtype(my_model.x, tf.uint8)
-
-    imsave(ARGS.savefile, final_image)
+    # --------------------------------------------------------------------------------------------------------------
+    # PART 4 : save and show result from final epoch #
+    # --------------------------------------------------------------------------------------------------------------
+    # model.x is the most recent image created by the model
+    result_tensor = my_model.x
+    # save image in output folder
+    final_image = save_tensor_as_image(result_tensor, 'out/results/', ARGS.savefile)
 
+    # show the final result :)
     plt.imshow(final_image)
     plt.show()
 
diff --git a/model.py b/model.py
index d8ce4c88..3f1f83f0 100644
--- a/model.py
+++ b/model.py
@@ -76,7 +76,7 @@ class YourModel(tf.keras.Model):
         # style target is the second output of the vgg16 mode, the Gram Matrix of the style layers
         self.style_target = self.vgg16(img_to_np(style_image))[1]
 
-    # here for convention - here is the forward pass
+    # here for convention - this is the forward pass
     def call(self, x):
         # call only onto our created model
         x = self.vgg16(x * 255)
diff --git a/save_images/epoch0.jpg b/out/checkpoints/cp-example.jpg/0_epoch.jpg
index 81c239a8..81c239a8 100644
--- a/save_images/epoch0.jpg
+++ b/out/checkpoints/cp-example.jpg/0_epoch.jpg
diff --git a/save_images/epoch100.jpg b/out/checkpoints/cp-example.jpg/100_epoch.jpg
index ad5de5f5..ad5de5f5 100644
--- a/save_images/epoch100.jpg
+++ b/out/checkpoints/cp-example.jpg/100_epoch.jpg
diff --git a/save_images/epoch200.jpg b/out/checkpoints/cp-example.jpg/200_epoch.jpg
index 1096ea7f..1096ea7f 100644
--- a/save_images/epoch200.jpg
+++ b/out/checkpoints/cp-example.jpg/200_epoch.jpg
diff --git a/save_images/epoch300.jpg b/out/checkpoints/cp-example.jpg/300_epoch.jpg
index 225947f4..225947f4 100644
--- a/save_images/epoch300.jpg
+++ b/out/checkpoints/cp-example.jpg/300_epoch.jpg
diff --git a/save_images/epoch400.jpg b/out/checkpoints/cp-example.jpg/400_epoch.jpg
index 6e286ae0..6e286ae0 100644
--- a/save_images/epoch400.jpg
+++ b/out/checkpoints/cp-example.jpg/400_epoch.jpg
diff --git a/save_images/epoch500.jpg b/out/checkpoints/cp-example.jpg/500_epoch.jpg
index db82fe4a..db82fe4a 100644
--- a/save_images/epoch500.jpg
+++ b/out/checkpoints/cp-example.jpg/500_epoch.jpg
diff --git a/save_images/epoch600.jpg b/out/checkpoints/cp-example.jpg/600_epoch.jpg
index 6b3b999a..6b3b999a 100644
--- a/save_images/epoch600.jpg
+++ b/out/checkpoints/cp-example.jpg/600_epoch.jpg
diff --git a/save_images/epoch700.jpg b/out/checkpoints/cp-example.jpg/700_epoch.jpg
index fd8ec867..fd8ec867 100644
--- a/save_images/epoch700.jpg
+++ b/out/checkpoints/cp-example.jpg/700_epoch.jpg
diff --git a/save_images/epoch800.jpg b/out/checkpoints/cp-example.jpg/800_epoch.jpg
index 5f115a9c..5f115a9c 100644
--- a/save_images/epoch800.jpg
+++ b/out/checkpoints/cp-example.jpg/800_epoch.jpg
diff --git a/save_images/epoch900.jpg b/out/checkpoints/cp-example.jpg/900_epoch.jpg
index 57964612..57964612 100644
--- a/save_images/epoch900.jpg
+++ b/out/checkpoints/cp-example.jpg/900_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/0_epoch.jpg b/out/checkpoints/cp-example2.jpg/0_epoch.jpg
new file mode 100644
index 00000000..367d809c
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/0_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/100_epoch.jpg b/out/checkpoints/cp-example2.jpg/100_epoch.jpg
new file mode 100644
index 00000000..8c00d6f3
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/100_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/200_epoch.jpg b/out/checkpoints/cp-example2.jpg/200_epoch.jpg
new file mode 100644
index 00000000..b6f5b04a
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/200_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/300_epoch.jpg b/out/checkpoints/cp-example2.jpg/300_epoch.jpg
new file mode 100644
index 00000000..2d3393c2
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/300_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/400_epoch.jpg b/out/checkpoints/cp-example2.jpg/400_epoch.jpg
new file mode 100644
index 00000000..80c0ae8e
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/400_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/500_epoch.jpg b/out/checkpoints/cp-example2.jpg/500_epoch.jpg
new file mode 100644
index 00000000..d4e21a4a
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/500_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/600_epoch.jpg b/out/checkpoints/cp-example2.jpg/600_epoch.jpg
new file mode 100644
index 00000000..609f33df
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/600_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/700_epoch.jpg b/out/checkpoints/cp-example2.jpg/700_epoch.jpg
new file mode 100644
index 00000000..b4d777f4
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/700_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/800_epoch.jpg b/out/checkpoints/cp-example2.jpg/800_epoch.jpg
new file mode 100644
index 00000000..efd67fe7
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/800_epoch.jpg
diff --git a/out/checkpoints/cp-example2.jpg/900_epoch.jpg b/out/checkpoints/cp-example2.jpg/900_epoch.jpg
new file mode 100644
index 00000000..d274cea1
--- /dev/null
+++ b/out/checkpoints/cp-example2.jpg/900_epoch.jpg
diff --git a/out/results/example.jpg b/out/results/example.jpg
new file mode 100644
index 00000000..4da02115
--- /dev/null
+++ b/out/results/example.jpg
diff --git a/out/results/example2.jpg b/out/results/example2.jpg
new file mode 100644
index 00000000..50951e14
--- /dev/null
+++ b/out/results/example2.jpg
diff --git a/save_images/epoch150.jpg b/save_images/epoch150.jpg
deleted file mode 100644
index e18e565e..00000000
--- a/save_images/epoch150.jpg
+++ /dev/null
diff --git a/save_images/epoch250.jpg b/save_images/epoch250.jpg
deleted file mode 100644
index 6b0be244..00000000
--- a/save_images/epoch250.jpg
+++ /dev/null
diff --git a/save_images/epoch350.jpg b/save_images/epoch350.jpg
deleted file mode 100644
index eea0a167..00000000
--- a/save_images/epoch350.jpg
+++ /dev/null
diff --git a/save_images/epoch450.jpg b/save_images/epoch450.jpg
deleted file mode 100644
index 8ef2132c..00000000
--- a/save_images/epoch450.jpg
+++ /dev/null
diff --git a/save_images/epoch50.jpg b/save_images/epoch50.jpg
deleted file mode 100644
index da06f7e1..00000000
--- a/save_images/epoch50.jpg
+++ /dev/null
diff --git a/save_images/epoch550.jpg b/save_images/epoch550.jpg
deleted file mode 100644
index 0207735b..00000000
--- a/save_images/epoch550.jpg
+++ /dev/null
diff --git a/save_images/epoch650.jpg b/save_images/epoch650.jpg
deleted file mode 100644
index 29eab7a8..00000000
--- a/save_images/epoch650.jpg
+++ /dev/null
diff --git a/save_images/epoch750.jpg b/save_images/epoch750.jpg
deleted file mode 100644
index 7881cd98..00000000
--- a/save_images/epoch750.jpg
+++ /dev/null
diff --git a/save_images/epoch850.jpg b/save_images/epoch850.jpg
deleted file mode 100644
index ff1cab0d..00000000
--- a/save_images/epoch850.jpg
+++ /dev/null
diff --git a/save_images/epoch950.jpg b/save_images/epoch950.jpg
deleted file mode 100644
index 42deef1e..00000000
--- a/save_images/epoch950.jpg
+++ /dev/null