Refactor test.py

src/data/dataset.py

@@ -1,7 +1,3 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import os
 
 import tensorflow as tf
@@ -38,13 +34,13 @@ def load_train_data(self):
  # Throwing away the remainder allows the pipeline to report a fixed sized
  # batch size, aiding in model definition downstream.
  train_datasetA = train_datasetA.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
- batch_size=self.opt.batch_size,
- num_parallel_calls=self.opt.num_threads,
- drop_remainder=True))
+  batch_size=self.opt.batch_size,
+  num_parallel_calls=self.opt.num_threads,
+  drop_remainder=True))
  train_datasetB = train_datasetB.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
- batch_size=self.opt.batch_size,
- num_parallel_calls=self.opt.num_threads,
- drop_remainder=True))
+  batch_size=self.opt.batch_size,
+  num_parallel_calls=self.opt.num_threads,
+  drop_remainder=True))
  # Queue up a number of batches on CPU side:
  train_datasetA = train_datasetA.prefetch(buffer_size=self.opt.num_threads)
  train_datasetB = train_datasetB.prefetch(buffer_size=self.opt.num_threads)
@@ -61,13 +57,13 @@ def load_test_data(self):
  test_datasetA = tf.data.Dataset.list_files(self.testA_path + os.sep + '*.jpg', shuffle=False)
  test_datasetB = tf.data.Dataset.list_files(self.testB_path + os.sep + '*.jpg', shuffle=False)
  test_datasetA = test_datasetA.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
- batch_size=1,
- num_parallel_calls=self.opt.num_threads,
- drop_remainder=False))
+  batch_size=1,
+  num_parallel_calls=self.opt.num_threads,
+  drop_remainder=False))
  test_datasetB = test_datasetB.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
- batch_size=1,
- num_parallel_calls=self.opt.num_threads,
- drop_remainder=False))
+  batch_size=1,
+  num_parallel_calls=self.opt.num_threads,
+  drop_remainder=False))
  test_datasetA = test_datasetA.prefetch(buffer_size=self.opt.num_threads)
  test_datasetB = test_datasetB.prefetch(buffer_size=self.opt.num_threads)
  if self.opt.gpu_id != -1:
@@ -90,17 +86,22 @@ def load_image(self, image_file):
  image = (image - 0.5) * 2
  return image
 
- def save_images(self, image_to_save, save_dir, image_index):
- save_file = os.path.join(save_dir,'test' + str(image_index) + '.jpg')
- # Reshape to get rid of batch size dimension in the tensor.
- image = tf.reshape(image_to_save, shape=[self.opt.img_size, self.opt.img_size, 3])
- # Scale from [-1, 1] to [0, 1).
- image = (image * 0.5) + 0.5
- # Convert to uint8 (range [0, 255]), saturate to avoid possible under/overflow.
- image = tf.image.convert_image_dtype(image, dtype=tf.uint8, saturate=True)
- # JPEG encode image into string Tensor.
- image_string = tf.image.encode_jpeg(image, format='rgb', quality=95)
- tf.write_file(filename=save_file, contents=image_string)
+ def save_images(self, test_images, index):
+ image_paths = []
+ image_paths.append(os.path.join(opt.results_dir, 'generatedA', 'test' + str(index) + '_real.jpg'))
+ image_paths.append(os.path.join(opt.results_dir, 'generatedA', 'test' + str(index) + '_fake.jpg'))
+ image_paths.append(os.path.join(opt.results_dir, 'generatedB', 'test' + str(index) + '_real.jpg'))
+ image_paths.append(os.path.join(opt.results_dir, 'generatedB', 'test' + str(index) + '_fake.jpg'))
+ for index in range(len(test_images)):
+ # Reshape to get rid of batch size dimension in the tensor.
+ image = tf.reshape(test_images[index], shape=[self.opt.img_size, self.opt.img_size, 3])
+ # Scale from [-1, 1] to [0, 1).
+ image = (image * 0.5) + 0.5
+ # Convert to uint8 (range [0, 255]), saturate to avoid possible under/overflow.
+ image = tf.image.convert_image_dtype(image, dtype=tf.uint8, saturate=True)
+ # JPEG encode image into string Tensor.
+ image_string = tf.image.encode_jpeg(image, format='rgb', quality=95)
+ tf.write_file(filename=image_paths[index], contents=image_string)
 
  def get_batches_per_epoch(self, opt):
  # floor(Avg dataset size / batch_size)

src/models/cyclegan.py

@@ -1,7 +1,3 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import os
 
 import tensorflow as tf
@@ -30,7 +26,8 @@ def __init__(self, opt):
  self.discB_buffer = ImageHistoryBuffer(opt)
  # Restore latest checkpoint:
  self.initialize_checkpoint()
- self.restore_checkpoint()
+ if not opt.training or opt.load_checkpoint:
+ self.restore_checkpoint()
 
  def initialize_checkpoint(self):
  if self.opt.training:
@@ -53,13 +50,11 @@ def restore_checkpoint(self):
  if self.opt.load_checkpoint and latest_checkpoint is not None:
  # Use assert_existing_objects_matched() instead of asset_consumed() here because
  # optimizers aren't initialized fully until first gradient update.
- # This will throw an exception if checkpoint does not restore the model weights.
+ # This will throw an exception if the checkpoint does not restore the model weights.
  self.checkpoint.restore(latest_checkpoint).assert_existing_objects_matched()
  print("Checkpoint restored from ", latest_checkpoint)
- # Uncomment below to print full list of checkpoint metadata.
- #print(tf.contrib.checkpoint.object_metadata(latest_checkpoint))
  else:
- print("No checkpoint found, initializing model.")
+ print("Failed to restore checkpoint, initializing model.")
 
  def set_input(self, input):
  # Get next batches:
@@ -106,8 +101,8 @@ def backward_G(self):
  genA2B_loss = generator_loss(self.discB(self.fakeB))
  genB2A_loss = generator_loss(self.discA(self.fakeA))
 
- cyc_lossA = cycle_consistency_loss(self.dataA, self.reconstructedA) * opt.cyc_lambda
- cyc_lossB = cycle_consistency_loss(self.dataB, self.reconstructedB) * opt.cyc_lambda
+ cyc_lossA = cycle_loss(self.dataA, self.reconstructedA) * opt.cyc_lambda
+ cyc_lossB = cycle_loss(self.dataB, self.reconstructedB) * opt.cyc_lambda
 
  gen_loss = genA2B_loss + genB2A_loss + cyc_lossA + cyc_lossB + id_lossA + id_lossB
  return gen_loss
@@ -126,8 +121,8 @@ def optimize_parameters(self):
  gen_variables = [self.genA2B.variables, self.genB2A.variables]
  gen_gradients = genTape.gradient(gen_loss, gen_variables)
  self.gen_opt.apply_gradients(list(zip(gen_gradients[0], gen_variables[0])) \
- + list(zip(gen_gradients[1], gen_variables[1])),
- global_step=self.global_step)
+  + list(zip(gen_gradients[1], gen_variables[1])),
+  global_step=self.global_step)
 
  for net in (self.discA, self.discB):
  for layer in net.layers:
@@ -142,15 +137,20 @@ def optimize_parameters(self):
  discA_gradients = discTape.gradient(discA_loss, self.discA.variables)
  discB_gradients = discTape.gradient(discB_loss, self.discB.variables)
  self.disc_opt.apply_gradients(zip(discA_gradients, self.discA.variables),
-  global_step=self.global_step)
+ global_step=self.global_step)
  self.disc_opt.apply_gradients(zip(discB_gradients, self.discB.variables),
-  global_step=self.global_step)
+ global_step=self.global_step)
 
  def save_model(self):
  checkpoint_prefix = os.path.join(self.opt.save_dir, 'checkpoints', 'ckpt')
  checkpoint_path = self.checkpoint.save(file_prefix=checkpoint_prefix)
  print("Checkpoint saved at ", checkpoint_path)
 
+ def test(self):
+ self.fakeA = self.genB2A(self.dataB)
+ self.fakeB = self.genA2B(self.dataA)
+ return [self.dataA, self.fakeA, self.dataB, self.fakeB]
+
  def update_learning_rate(self, batches_per_epoch):
  new_lr = self._get_learning_rate(batches_per_epoch)
  self.learning_rate.assign(new_lr)

src/test.py

@@ -1,64 +1,26 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import os
 import time
 
 import tensorflow as tf
 
-from train import initialize_checkpoint, define_model, restore_from_checkpoint
-from data.dataset import load_test_data, save_images
+from utils.options import Options
+from data.dataset import Dataset
+from models.cyclegan import CycleGANModel
 
 tf.enable_eager_execution()
 
-project_dir = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
-checkpoint_dir = os.path.join(project_dir, 'saved_models', 'checkpoints')
-dataset_id = 'facades'
-
-def test(data, model, checkpoint_info, dataset_id):
- path_to_dataset = os.path.join(project_dir, 'data', 'raw', dataset_id + os.sep)
- generatedA = os.path.join(path_to_dataset, 'generatedA' + os.sep)
- generatedB = os.path.join(path_to_dataset, 'generatedB' + os.sep)
- genA2B = model['genA2B']
- genB2A = model['genB2A']
+if __name__ == "__main__":
+ opt = Options(training=False)
+ # TODO: Test if this is always on CPU:
+ dataset = Dataset(opt)
+ model = CycleGANModel(opt)
 
- checkpoint, checkpoint_dir = checkpoint_info
- restore_from_checkpoint(checkpoint, checkpoint_dir)
- test_datasetA, test_datasetB, testA_size, testB_size = data
- test_datasetA = iter(test_datasetA)
- test_datasetB = iter(test_datasetB)
+ device = ("/gpu:" + str(opt.gpu_id)) if opt.gpu_id != -1 else "/cpu:0"
 
- for imageB in range(testB_size):
+ with tf.device(device):
  start = time.time()
- try:
- # Get next testing image:
- testB = test_datasetB.get_next()
- except tf.errors.OutOfRangeError:
- print("Error, run out of data")
- break
- genB2A_output = genB2A(testB)
- with tf.device("/cpu:0"):
- save_images(genB2A_output, save_dir=generatedA, image_index=imageB)
- print("Generating {} test A images finished in {} sec\n".format(testA_size, time.time()-start))
-
- for imageA in range(testA_size):
- start = time.time()
- try:
- # Get next testing image:
- testA = test_datasetA.get_next()
- except tf.errors.OutOfRangeError:
- print("Error, run out of data")
- break
- genA2B_output = genA2B(testA)
- with tf.device("/cpu:0"):
- save_images(genA2B_output, save_dir=generatedB, image_index=imageA)
- print("Generating {} test B images finished in {} sec\n".format(testB_size, time.time()-start))
-
-if __name__ == "__main__":
- with tf.device("/cpu:0"): # Preprocess data on CPU for significant performance gains.
- data = load_test_data(dataset_id, project_dir)
- with tf.device("/gpu:0"):
- model = define_model(training=False)
- checkpoint_info = initialize_checkpoint(checkpoint_dir, model, training=False)
- test(data, model, checkpoint_info, dataset_id)
+ for index in range(opt.num_test):
+ model.set_input(dataset.data)
+ test_images = model.test()
+ dataset.save_images(test_images, index=index)
+ print("Generating {} test images for both datasets finished in {} sec\n".format(opt.num_test, time.time()-start))

src/utils/options.py

@@ -3,28 +3,27 @@
 
 class Options(object):
 
- def __init__(self, isTrain=True):
+ def __init__(self, training=True):
  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  # basic options
- parser.add_argument('--data_dir', required=True, help='path to folder where the dataset is stored, should have subfolders trainA, trainB, testA, testB')
+ parser.add_argument('--data_dir', required=True, help='path to directory where the dataset is stored, should have subfolders trainA, trainB, testA, testB')
  parser.add_argument('--save_dir', required=True, help='checkpoints and Tensorboard summaries are saved here')
- parser.add_argument('--load_checkpoint', action='store_true', help='if true, loads latest checkpoint')
  parser.add_argument('--gpu_id', type=str, default='0', help='gpu id to run model on, use -1 for CPU, multigpu not supported')
  # model options
  parser.add_argument('--ngf', type=int, default=64, help='number of gen filters in the first conv layer')
  parser.add_argument('--ndf', type=int, default=64, help='number of disc filters in the first conv layer')
  parser.add_argument('--norm', action='store_true', help='if true, uses instance normalisation after each conv layer in D and G')
  parser.add_argument('--init_scale', type=float, default=0.02, help='stddev for weight initialisation; small variance helps prevent colour inversion.')
- parser.add_argument('--use_dropout', action='store_false', help='if true, use dropout for the generator')
  parser.add_argument('--gen_skip', action='store_false', help='if true, use skip connection from first residual block to last in generator')
  parser.add_argument('--resize_conv', action='store_false', help='if true, replace conv2dtranspose in generator with upsample -> conv2d')
+ parser.add_argument('--use_dropout', action='store_false', help='if true, use dropout for the generator')
+ parser.add_argument('--dropout_prob', type=float, default=0.5, help='dropout probability for all layers in generator')
  # dataset options
  cpu_count = multiprocessing.cpu_count()
- parser.add_argument('--num_threads', type=int, default=cpu_count, help='num threads for loading data')
- parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
+ parser.add_argument('--num_threads', type=int, default=cpu_count, help='number of CPU threads to use for loading data')
  parser.add_argument('--img_size', type=int, default=256, help='input image size')
  # get training/testing options
- if isTrain:
+ if training:
  parser = self.get_train_options(parser)
  else:
  parser = self.get_test_options(parser)
@@ -38,13 +37,15 @@ def __init__(self, isTrain=True):
  def get_train_options(self, parser):
  # training specific options
  parser.add_argument('--training', action='store_true', help='boolean for training/testing')
+ parser.add_argument('--load_checkpoint', action='store_true', help='if true, loads latest checkpoint')
  parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')
  parser.add_argument('--summary_freq', type=int, default=100, help='frequency of saving saving tensorboard summaries in training steps')
+ parser.add_argument('--epochs', type=int, default=200, help='number of epochs to train the model; learning rate decays to 0 by epoch 200')
+ parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
  parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
  parser.add_argument('--beta1', type=float, default=0.5, help='momentum term for adam')
  parser.add_argument('--niter', type=int, default=100, help='number of epochs at initial learning rate')
  parser.add_argument('--niter_decay', type=int, default=100, help='number of epochs to linearly decay learning rate to zero')
- parser.add_argument('--epochs', type=int, default=200, help='number of epochs to train the model; learning rate decays to 0 by epoch 200')
  parser.add_argument('--gan_mode', type=str, default='lsgan', help='GAN loss type [vanilla | lsgan | wgangp | rgan]. vanilla is the cross-entropy loss from original paper')
  parser.add_argument('--buffer_size', type=int, default=50, help='size of the image history buffer')
  parser.add_argument('--cyc_lambda', type=float, default=10, help='weight for cycle-consistency loss')
@@ -54,7 +55,7 @@ def get_train_options(self, parser):
  def get_test_options(self, parser):
  # test specific options
  parser.add_argument('--training', action='store_false', help='boolean for training/testing')
- parser.add_argument('--results_dir', required=True, help='directory to save results')
+ parser.add_argument('--results_dir', required=True, help='directory to save result images')
  parser.add_argument('--num_test', type=int, default=50, help='number of test images to generate')
  return parser