inversion json files

1 files changed, 190 insertions, 0 deletions

diff --git a/inversion/interpolation.py b/inversion/interpolation.py
new file mode 100644
index 0000000..7133337
--- /dev/null
+++ b/inversion/interpolation.py
@@ -0,0 +1,190 @@
+# ------------------------------------------------------------------------------
+# Linear interpolation between inverted images and generated images.
+# ------------------------------------------------------------------------------
+
+import functools
+import h5py
+import itertools
+import numpy as np
+import os
+import pickle
+import params
+import scipy
+import sys
+import tensorflow as tf
+import tensorflow_hub as hub
+import time
+import visualize as vs
+
+# --------------------------
+# Hyper-parameters.
+# --------------------------
+if len(sys.argv) < 2:
+  sys.exit('Must provide a configuration file.')
+params = params.Params(sys.argv[1])
+
+# --------------------------
+# Global variables.
+# --------------------------
+BATCH_SIZE = params.batch_size
+SAMPLE_SIZE = params.sample_size
+SAMPLES_DIR = 'interpolation'
+INVERSES_DIR = 'inverses'
+if not os.path.exists(SAMPLES_DIR):
+  os.makedirs(SAMPLES_DIR)
+if not os.path.exists(INVERSES_DIR):
+  os.makedirs(INVERSES_DIR)
+
+# --------------------------
+# Util functions.
+# --------------------------
+def interpolate(A, B, num_interps):
+  alphas = np.linspace(0, 1., num_interps)
+  if A.shape != B.shape:
+    raise ValueError('A and B must have the same shape to interpolate.')
+  return np.array([(1-a)*A + a*B for a in alphas])
+
+# One hot encoding for classes.
+def one_hot(values):
+  return np.eye(N_CLASS)[values]
+
+# Random sampler for classes.
+def label_sampler(size=[BATCH_SIZE]):
+  return np.random.random_integers(low=0, high=N_CLASS-1, size=size)
+
+def label_hot_sampler(size=[BATCH_SIZE]):
+  return one_hot(label_sampler(size=size))
+
+# --------------------------
+# Load Graph.
+# --------------------------
+generator = hub.Module(str(params.generator_path))
+
+gen_signature = 'generator'
+if 'generator' not in generator.get_signature_names():
+  gen_signature = 'default'
+
+input_info = generator.get_input_info_dict(gen_signature)
+COND_GAN = 'y' in input_info
+
+if COND_GAN:
+  Z_DIM = input_info['z'].get_shape().as_list()[1]
+  latent = tf.get_variable(name='latent', dtype=tf.float32,
+                           shape=[BATCH_SIZE, Z_DIM])
+  N_CLASS = input_info['y'].get_shape().as_list()[1]
+  label = tf.get_variable(name='label', dtype=tf.float32,
+                          shape=[BATCH_SIZE, N_CLASS])
+  gen_in = dict(params.generator_fixed_inputs)
+  gen_in['z'] = latent
+  gen_in['y'] = label
+  gen_img = generator(gen_in, signature=gen_signature)
+else:
+  Z_DIM = input_info['default'].get_shape().as_list()[1]
+  latent = tf.get_variable(name='latent', dtype=tf.float32,
+                           shape=[BATCH_SIZE, Z_DIM])
+  if (params.generator_fixed_inputs):
+    gen_in = dict(params.generator_fixed_inputs)
+    gen_in['z'] = latent
+    gen_img = generator(gen_in, signature=gen_signature)
+  else:
+    gen_img = generator(latent, signature=gen_signature)
+
+# Convert generated image to channels_first.
+gen_img = tf.transpose(gen_img, [0, 3, 1, 2])
+
+# Override intermediate layer.
+if params.inv_layer == 'latent':
+  encoding = latent
+  ENC_SHAPE = [Z_DIM]
+else:
+  layer_name = 'module_apply_' + gen_signature + '/' + params.inv_layer
+  gen_encoding = tf.get_default_graph().get_tensor_by_name(layer_name)
+  ENC_SHAPE = gen_encoding.get_shape().as_list()[1:]
+  encoding = tf.get_variable(name='encoding', dtype=tf.float32,
+                             shape=[BATCH_SIZE,] + ENC_SHAPE)
+  tf.contrib.graph_editor.swap_ts(gen_encoding, tf.convert_to_tensor(encoding))
+
+# Define image shape.
+IMG_SHAPE = gen_img.get_shape().as_list()[1:]
+
+# --------------------------
+# Noise source.
+# --------------------------
+def noise_sampler():
+  return np.random.normal(size=[BATCH_SIZE, Z_DIM])
+
+# --------------------------
+# Dataset.
+# --------------------------
+in_file = h5py.File(os.path.join(INVERSES_DIR, params.out_dataset), 'r')
+in_images = in_file['xtrain']
+if COND_GAN:
+  in_labels = in_file['ytrain']
+in_encoding = in_file['encoding']
+in_latent = in_file['latent']
+NUM_IMGS = in_images.shape[0] # number of images.
+
+# --------------------------
+# Training.
+# --------------------------
+# Start session.
+sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
+sess.run(tf.global_variables_initializer())
+sess.run(tf.tables_initializer())
+
+for i in range(0, NUM_IMGS, BATCH_SIZE):
+  # Set label.
+  if COND_GAN:
+    sess.run(label.assign(one_hot(in_labels[i:i+BATCH_SIZE])))
+
+  # Linear interpolation between G_1(z*) and G_1(z*)+delta*.
+  sample_enc_1 = in_encoding[i:i+BATCH_SIZE]
+  out_batch = np.ndarray(shape=[8,BATCH_SIZE]+IMG_SHAPE, dtype='uint8')
+  out_batch[0] = in_images[i:i+BATCH_SIZE]
+  sess.run(latent.assign(in_latent[i:i+BATCH_SIZE]))
+  sample_enc_2 = sess.run(gen_encoding)
+  sample_enc = interpolate(sample_enc_1, sample_enc_2, 7)
+  for j in range(0,7):
+    sess.run(encoding.assign(sample_enc[j]))
+    gen_images = sess.run(gen_img)
+    gen_images = vs.data2img(gen_images)
+    out_batch[j+1] = gen_images
+
+  out_batch = np.transpose(out_batch, [1, 0, 2, 3, 4])
+  for k in range(BATCH_SIZE):
+    out_batch_k = vs.seq_transform(out_batch[k])
+    # Add white padding.
+    pad = 20
+    out_batch_kk = np.ndarray(shape=[IMG_SHAPE[1], IMG_SHAPE[1]*8+pad, 3],
+                              dtype='uint8')
+    out_batch_kk[:,:IMG_SHAPE[1],:] = out_batch_k[:,:IMG_SHAPE[1],:]
+    out_batch_kk[:,IMG_SHAPE[1]:IMG_SHAPE[1]+pad,:] = 255
+    out_batch_kk[:,IMG_SHAPE[1]+pad:,:] = out_batch_k[:,IMG_SHAPE[1]:,:]
+
+    vs.save_image('{}/interpolation_delta_{}.png'.format(SAMPLES_DIR, i+k), out_batch_kk)
+    print('Saved delta interpolation for img: {}.'.format(i+k))
+
+  # Linear interpolation between G_1(z_random) and G_1(z*)+delta*.
+  sample_enc_1 = in_encoding[i:i+BATCH_SIZE]
+  sample_z_1 = in_latent[i:i+BATCH_SIZE]
+  out_batch = np.ndarray(shape=[8*8,BATCH_SIZE]+IMG_SHAPE, dtype='uint8')
+  for k in range(8):
+    sample_z_2 = noise_sampler()
+    sess.run(latent.assign(sample_z_2))
+    sample_enc_2 = sess.run(gen_encoding)
+    sample_z = interpolate(sample_z_1, sample_z_2, 8)
+    sample_enc = interpolate(sample_enc_1, sample_enc_2, 8)
+    for j in range(8):
+      sess.run(latent.assign(sample_z[j]))
+      sess.run(encoding.assign(sample_enc[j]))
+      gen_images = sess.run(gen_img)
+      gen_images = vs.data2img(gen_images)
+      out_batch[k*8+j] = gen_images
+
+  out_batch = np.transpose(out_batch, [1, 0, 2, 3, 4])
+  for k in range(BATCH_SIZE):
+    out_batch_k = vs.grid_transform(out_batch[k])
+    vs.save_image('{}/interpolation_rand_{}.png'.format(SAMPLES_DIR, i+k), out_batch_k)
+    print('Saved rand interpolation for img: {}.'.format(i+k))
+
+sess.close()