Updated README and pretty much everything

README.md

@@ -1,25 +1,27 @@
-# FSRCNN-Tensorflow
-TensorFlow implementation of the Fast Super-Resolution Convolutional Neural Network (FSRCNN). This implements two models, FSRCNN which is more accurate and FSRCNN-s which is faster (approaches real-time performance). Based on this [project](http://mmlab.ie.cuhk.edu.hk/projects/FSRCNN.html).
+# FSRCNN-TensorFlow
+TensorFlow implementation of the Fast Super-Resolution Convolutional Neural Network (FSRCNN). This implements two models: FSRCNN which is more accurate and FSRCNN-s which is faster (approaches real-time performance). Based on this [project](http://mmlab.ie.cuhk.edu.hk/projects/FSRCNN.html).
 
 ## Prerequisites
+ * Python 2.7
  * TensorFlow
- * Scipy version > 0.18 ('mode' option from scipy.misc.imread function)
+ * Scipy version > 0.18
  * h5py
  * PIL
 
 ## Usage
-For training , `python main.py`
-Can specify epochs, learning rate,  `python main.py --epochs 10`
+For training: `python main.py`
 <br>
-For testing, `python main.py --is_train False`
+Can specify epochs, learning rate, data directory, etc:  `python main.py --epochs 10 --learning_rate 0.0001 --data_dir Train`
+<br>
+For testing: `python main.py --is_train False`
 
-To use FSCRNN-s over FSCRNN , `python main.py --fast True`
+To use FSCRNN-s instead of FSCRNN: `python main.py --fast True`
 
-Includes script expand_data.py which scales and rotates all the images in your training set to expand your dataset just like in the paper
-`python expand_data.py Train`
+Includes script `expand_data.py` which scales and rotates all the images in your training set to expand it: `python expand_data.py Train`
 
 ## Result
-After training 15,000 epochs, I got similar super-resolved image to reference paper. Training time takes 12 hours 16 minutes and 1.41 seconds. My desktop performance is Intel I7-6700 CPU, GTX970, and 16GB RAM. Result images are shown below.<br><br>
+
+<br><br>
 Original butterfly image:
 ![orig](https://github.com/tegg89/SRCNN-Tensorflow/blob/master/result/orig.png)<br>
 Bicubic interpolated image:

main.py

@@ -16,7 +16,7 @@ flags.DEFINE_integer("c_dim", 1, "Dimension of image color [1]")
 flags.DEFINE_integer("scale", 3, "The size of scale factor for preprocessing input image [3]")
 flags.DEFINE_integer("stride", 4, "The size of stride to apply to input image [4]")
 flags.DEFINE_string("checkpoint_dir", "checkpoint", "Name of checkpoint directory [checkpoint]")
-flags.DEFINE_string("test_dir", "result", "Name of test result directory [result]")
+flags.DEFINE_string("output_dir", "result", "Name of test output directory [result]")
 flags.DEFINE_string("data_dir", "FastTrain", "Name of data directory to train on [FastTrain]")
 flags.DEFINE_boolean("is_train", True, "True for training, False for testing [True]")
 flags.DEFINE_integer("threads", 1, "Number of processes to pre-process data with [1]")
@@ -29,10 +29,12 @@ pp = pprint.PrettyPrinter()
 def main(_):
   pp.pprint(flags.FLAGS.__flags)
 
+  if FLAGS.fast:
+    FLAGS.checkpoint_dir = 'fast_{}'.format(FLAGS.checkpoint_dir)
   if not os.path.exists(FLAGS.checkpoint_dir):
     os.makedirs(FLAGS.checkpoint_dir)
-  if not os.path.exists(FLAGS.sample_dir):
-    os.makedirs(FLAGS.sample_dir)
+  if not os.path.exists(FLAGS.output_dir):
+    os.makedirs(FLAGS.output_dir)
 
 
   with tf.Session() as sess:

model.py

@@ -24,12 +24,11 @@ class FSRCNN(object):
     self.sess = sess
     self.fast = config.fast
     self.is_train = config.is_train
-    self.save_image = config.save_image
     self.c_dim = config.c_dim
     self.is_grayscale = (self.c_dim == 1)
     self.epoch = config.epoch
-    self.stride = config.stride
     self.scale = config.scale
+    self.stride = config.stride
     self.batch_size = config.batch_size
     self.learning_rate = config.learning_rate
     self.momentum = config.momentum
@@ -39,13 +38,15 @@ class FSRCNN(object):
     # Different image/label sub-sizes for different scaling factors x2, x3, x4
     scale_factors = [[10, 20], [11, 21], [6, 24]]
     self.image_size, self.label_size = scale_factors[self.scale - 2]
+    if not self.is_train:
+      self.stride = [10, 7, 6][self.scale - 2]
 
     # Different model layer counts/filter sizes for FSRCNN vs FSRCNN-s (fast)
     model_params = [[56, 12, 4], [32, 5, 1]]
     self.model_params = model_params[self.fast]
     
     self.checkpoint_dir = config.checkpoint_dir
-    self.sample_dir = config.sample_dir
+    self.output_dir = config.output_dir
     self.data_dir = config.data_dir
     self.build_model()
 
@@ -139,6 +140,7 @@ class FSRCNN(object):
           print("Epoch: [%2d], step: [%2d], time: [%4.4f], loss: [%.8f]" \
             % ((ep+1), counter, time.time() - start_time, err))
 
+        # Save every 500 steps
         if counter % 500 == 0:
           self.save(self.checkpoint_dir, counter)
 
@@ -154,10 +156,11 @@ class FSRCNN(object):
     print("Start Average: [%.6f], End Average: [%.6f], Improved: [%.2f%%]" \
       % (start_average, end_average, 100 - (100*end_average/start_average)))
 
-    title = "Training complete - FSRCNN"
-    notification = "{}-{}-{} done training after {} epochs".format(self.image_size, self.label_size, self.stride, self.epoch);
-    notify_command = 'notify-send "{}" "{}"'.format(title, notification)
-    os.system(notify_command)
+    # Linux desktop notification when training has been completed
+    # title = "Training complete - FSRCNN"
+    # notification = "{}-{}-{} done training after {} epochs".format(self.image_size, self.label_size, self.stride, self.epoch);
+    # notify_command = 'notify-send "{}" "{}"'.format(title, notification)
+    # os.system(notify_command)
 
   
   def test(self):
@@ -173,7 +176,7 @@ class FSRCNN(object):
 
     result = merge(result, [nx, ny])
     result = result.squeeze()
-    image_path = os.path.join(os.getcwd(), self.sample_dir)
+    image_path = os.path.join(os.getcwd(), self.output_dir)
     image_path = os.path.join(image_path, "test_image.png")
 
     array_image_save(result * 255, image_path)
@@ -214,8 +217,6 @@ class FSRCNN(object):
   def save(self, checkpoint_dir, step):
     model_name = "FSRCNN.model"
     model_dir = "%s_%s" % ("fsrcnn", self.label_size)
-    if self.fast:
-      checkpoint_dir = 'fast_{}'.format(checkpoint_dir)
     checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
 
     if not os.path.exists(checkpoint_dir):
@@ -228,8 +229,6 @@ class FSRCNN(object):
   def load(self, checkpoint_dir):
     print(" [*] Reading checkpoints...")
     model_dir = "%s_%s" % ("fsrcnn", self.label_size)
-    if self.fast:
-      checkpoint_dir = 'fast_{}'.format(checkpoint_dir)
     checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
 
     ckpt = tf.train.get_checkpoint_state(checkpoint_dir)

utils.py

@@ -12,7 +12,6 @@ import struct
 import tensorflow as tf
 from PIL import Image  
 from scipy.misc import imread
-from scipy.ndimage import interpolation
 import numpy as np
 from multiprocessing import Pool, Lock, active_children
 
@@ -39,13 +38,9 @@ def preprocess(path, scale=3):
   Preprocess single image file 
     (1) Read original image as YCbCr format (and grayscale as default)
     (2) Normalize
-    (3) Downsampled by scale factor using bicubic interpolation
-
-  Args:
-    path: file path of desired file
-    input_: image downsampled (low-resolution)
-    label_: image with original resolution (high-resolution)
+    (3) Downsampled by scale factor
   """
+
   image = Image.open(path).convert('L')
   (width, height) = image.size
   label_ = np.array(list(image.getdata())).astype(np.float).reshape((height, width)) / 255
@@ -61,10 +56,6 @@ def preprocess(path, scale=3):
   (width, height) = scaled_image.size
   input_ = np.array(list(scaled_image.getdata())).astype(np.float).reshape((height, width))
 
-  if FLAGS.save_image and not FLAGS.is_train:
-    array_image_save(label_ * 255, "sample/test_image_original.bmp")
-    array_image_save(input_ * 255, "sample/test_image_downsampled.bmp")
-
   return input_, label_
 
 def prepare_data(sess, dataset):
@@ -88,7 +79,6 @@ def make_data(sess, checkpoint_dir, data, label):
   Make input data as h5 file format
   Depending on 'is_train' (flag value), savepath would be changed.
   """
-  checkpoint_dir = 'fast_{}'.format(checkpoint_dir)
   if FLAGS.is_train:
     savepath = os.path.join(os.getcwd(), '{}/train.h5'.format(checkpoint_dir))
   else:
@@ -132,8 +122,7 @@ def train_input_worker(args):
   image_data, config = args
   image_size, label_size, stride, scale, save_image = config
 
-  single_input_sequence = []
-  single_label_sequence = []
+  single_input_sequence, single_label_sequence = [], []
   padding = abs(image_size - label_size) / 2 # (21 - 11) / 2 = 5
   label_padding = label_size / scale # 21 / 3 = 7
 
@@ -186,8 +175,7 @@ def thread_train_setup(config):
 
   print("All worker processes done!")
 
-  sub_input_sequence = []
-  sub_label_sequence = []
+  sub_input_sequence, sub_label_sequence = [], []
 
   for image in range(len(results)):
     single_input_sequence, single_label_sequence = results[image]
@@ -210,8 +198,7 @@ def train_input_setup(config):
   # Load data path
   data = prepare_data(sess, dataset=config.data_dir)
 
-  sub_input_sequence = []
-  sub_label_sequence = []
+  sub_input_sequence, sub_label_sequence = [], []
   padding = abs(image_size - label_size) / 2 # (21 - 11) / 2 = 5
   label_padding = label_size / scale # 21 / 3 = 7
 
@@ -250,8 +237,7 @@ def test_input_setup(config):
   # Load data path
   data = prepare_data(sess, dataset="Test")
 
-  sub_input_sequence = []
-  sub_label_sequence = []
+  sub_input_sequence, sub_label_sequence = [], []
   padding = abs(image_size - label_size) / 2 # (21 - 11) / 2 = 5
   label_padding = label_size / scale # 21 / 3 = 7