upload bytesio object to cortex

13 files changed, 889 insertions, 256 deletions

diff --git a/cli/app/commands/biggan/fetch.py b/cli/app/commands/biggan/fetch.py
new file mode 100644
index 0000000..94117c9
--- /dev/null
+++ b/cli/app/commands/biggan/fetch.py
@@ -0,0 +1,13 @@
+import click
+
+from app.utils.cortex_utils import fetch_cortex_folder
+
+@click.command('')
+@click.option('-i', '--folder_id', 'opt_folder_id', type=int,
+  help='Folder ID to fetch')
+@click.pass_context
+def cli(ctx, opt_folder_id):
+  """
+  Fetch JSON from the server
+  """
+  fetch_cortex_folder(opt_folder_id)
diff --git a/cli/app/commands/biggan/search_class.py b/cli/app/commands/biggan/search_class.py
index 8ca31d6..0501729 100644
--- a/cli/app/commands/biggan/search_class.py
+++ b/cli/app/commands/biggan/search_class.py
@@ -1,31 +1,6 @@
 import click
 
-from app.utils import click_utils
-from app.settings import app_cfg
-
-import os
-from os.path import join
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
-
-import time
-import numpy as np
-import random
-from subprocess import call
-import cv2 as cv
-from PIL import Image
-from glob import glob
-import tensorflow as tf
-import tensorflow_hub as hub
-import shutil
-import h5py
-
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-
-from app.search.json import save_params_latent, save_params_dense
-from app.search.image import image_to_uint8, imconvert_uint8, imconvert_float32, \
-  imread, imwrite, imgrid, resize_and_crop_image
-from app.search.vector import truncated_z_sample, truncated_z_single, \
-  create_labels, create_labels_uniform
+from app.search.search_class import find_nearest_vector_for_images
 
 @click.command('')
 @click.option('-i', '--input', 'opt_fp_in', required=True, 
@@ -46,139 +21,4 @@ def cli(ctx, opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag):
   """
   Search for an image (class vector) in BigGAN using gradient descent
   """
-
-  sess = tf.compat.v1.Session()
-
-  generator = hub.Module('https://tfhub.dev/deepmind/biggan-512/2')
-
-  if os.path.isdir(opt_fp_in):
-    paths = glob(os.path.join(opt_fp_in, '*.jpg')) + \
-      glob(os.path.join(opt_fp_in, '*.jpeg')) + \
-      glob(os.path.join(opt_fp_in, '*.png'))
-  else:
-    paths = [opt_fp_in]
-
-  fp_inverses = os.path.join(app_cfg.DIR_INVERSES, opt_tag)
-  os.makedirs(fp_inverses, exist_ok=True)
-  save_params_latent(fp_inverses, opt_tag)
-  save_params_dense(fp_inverses, opt_tag)
-  out_file = h5py.File(join(fp_inverses, 'dataset.hdf5'), 'w')
-  out_images = out_file.create_dataset('xtrain', (len(paths), 3, 512, 512,), dtype='float32')
-  out_labels = out_file.create_dataset('ytrain', (len(paths), 1000,), dtype='float32')
-  out_latent = out_file.create_dataset('latent', (len(paths), 128,), dtype='float32')
-  out_fns = out_file.create_dataset('fn', (len(paths),), dtype=h5py.string_dtype())
-  for index, path in enumerate(paths):
-    if index == opt_limit:
-      break
-    out_fns[index] = os.path.basename(path)
-    fp_frames = find_nearest_vector(sess, generator, path, opt_dims, out_images, out_labels, out_latent, opt_steps, index)
-    if opt_video:
-      export_video(fp_frames)
-
-def find_nearest_vector(sess, generator, opt_fp_in, opt_dims, out_images, out_labels, out_latent, opt_steps, index):
-  """
-  Find the closest latent and class vectors for an image. Store the class vector in an HDF5.
-  """
-  gen_signature = 'generator'
-  if 'generator' not in generator.get_signature_names():
-    gen_signature = 'default'
-
-  # inputs = {k: tf.compat.v1.placeholder(v.dtype, v.get_shape().as_list(), k)
-  #         for k, v in generator.get_input_info_dict().items()}
-  batch_size = 1
-  truncation = 1.0
-
-  z_dim = 128
-  vocab_size = 1000
-  img_size = 512
-  num_channels = 3
-
-  z_initial = truncated_z_sample(batch_size, z_dim, truncation/2)
-  y_initial = create_labels_uniform(batch_size, vocab_size)
-
-  z_lr = 0.001
-  y_lr = 0.001
-
-  input_z = tf.compat.v1.Variable(z_initial, dtype=np.float32, constraint=lambda t: tf.clip_by_value(t, -2, 2))
-  input_y = tf.compat.v1.Variable(y_initial, dtype=np.float32, constraint=lambda t: tf.clip_by_value(t, 0, 1))
-  input_trunc = tf.compat.v1.constant(1.0)
-  output = generator({
-    'z': input_z,
-    'y': input_y,
-    'truncation': input_trunc,
-  }, signature=gen_signature)
-
-  layer_name = 'module_apply_' + gen_signature + '/' + "Generator_2/G_Z/Reshape:0"
-  gen_encoding = tf.get_default_graph().get_tensor_by_name(layer_name)
-
-  target = tf.compat.v1.placeholder(tf.float32, shape=(batch_size, img_size, img_size, num_channels))
-
-  # loss = tf.losses.compute_weighted_loss(tf.square(output - target), weights=mask)
-  loss = tf.compat.v1.losses.mean_squared_error(target, output)
-
-  train_step_z = tf.train.AdamOptimizer(z_lr).minimize(loss, var_list=[input_z], name='AdamOpterZ')
-  train_step_y = tf.train.AdamOptimizer(y_lr).minimize(loss, var_list=[input_y], name='AdamOpterY')
-
-  target_im, fp_frames = load_target_image(opt_fp_in)
-
-  # crop image and convert to format for next script
-  phi_target_for_inversion = resize_and_crop_image(target_im, 512)
-  b = np.dsplit(phi_target_for_inversion, 3)
-  phi_target_for_inversion = np.stack(b).reshape((3, 512, 512))
-
-  # create phi target for the latent / label pass
-  phi_target = resize_and_crop_image(target_im, opt_dims)
-  phi_target = np.expand_dims(phi_target, 0)
-  phi_target = np.repeat(phi_target, batch_size, axis=0)
-
-  # IMPORTANT: initialize variables before running the session
-  sess.run(tf.compat.v1.global_variables_initializer())
-  sess.run(tf.compat.v1.tables_initializer())
-
-  feed_dict = {
-    target: phi_target,
-  }
-
-  try:
-    print("Preparing to iterate...")
-    for i in range(opt_steps):
-      curr_loss, _, _ = sess.run([loss, train_step_z, train_step_y], feed_dict=feed_dict)
-
-      if i % 20 == 0:
-        phi_guess = sess.run(output)
-        guess_im = imgrid(imconvert_uint8(phi_guess), cols=1)
-        imwrite(join(app_cfg.DIR_OUTPUTS, fp_frames, 'frame_{:04d}.png'.format(int(i / 20))), guess_im)
-        print('iter: {}, loss: {}'.format(i, curr_loss))
-  except KeyboardInterrupt:
-    pass
-
-  z_guess, y_guess = sess.run([input_z, input_y])
-  out_images[index] = phi_target_for_inversion
-  out_labels[index] = y_guess
-  out_latent[index] = z_guess
-  return fp_frames
-
-def export_video(fp_frames):
-  print("Exporting video...")
-  cmd = [
-    '/home/lens/bin/ffmpeg',
-    '-y', # '-v', 'quiet',
-    '-r', '30',
-    '-i', join(app_cfg.DIR_OUTPUTS, fp_frames, 'frame_%04d.png'),
-    '-pix_fmt', 'yuv420p',
-    join(app_cfg.DIR_OUTPUTS, fp_frames + '.mp4')
-  ]
-  # print(' '.join(cmd))
-  call(cmd)
-  shutil.rmtree(join(app_cfg.DIR_OUTPUTS, fp_frames))
-
-def load_target_image(opt_fp_in):
-  print("Loading {}".format(opt_fp_in))
-  fn = os.path.basename(opt_fp_in)
-  fbase, ext = os.path.splitext(fn)
-  fp_frames = "frames_{}_{}".format(fbase, int(time.time() * 1000))
-  fp_frames_fullpath = join(app_cfg.DIR_OUTPUTS, fp_frames)
-  print("Output to {}".format(fp_frames_fullpath))
-  os.makedirs(fp_frames_fullpath, exist_ok=True)
-  target_im = imread(opt_fp_in)
-  return target_im, fp_frames
+  find_nearest_vector_for_images(opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag)
diff --git a/cli/app/commands/biggan/search_dense.py b/cli/app/commands/biggan/search_dense.py
new file mode 100644
index 0000000..0cf56a3
--- /dev/null
+++ b/cli/app/commands/biggan/search_dense.py
@@ -0,0 +1,15 @@
+import click
+
+from app.search.params import Params
+from app.search.search_dense import find_dense_embedding_for_images
+
+@click.command('')
+@click.option('-i', '--input', 'opt_fp_in', required=True, 
+  help='Path to input image')
+@click.pass_context
+def cli(ctx, opt_fp_in):
+  """
+  Search for an image (class vector) in BigGAN using gradient descent
+  """
+  params = Params(opt_fp_in)
+  find_dense_embedding_for_images(params)
diff --git a/cli/app/commands/cortex/upload.py b/cli/app/commands/cortex/upload.py
new file mode 100644
index 0000000..c1b0a46
--- /dev/null
+++ b/cli/app/commands/cortex/upload.py
@@ -0,0 +1,15 @@
+import click
+
+from app.utils.cortex_utils import upload_to_cortex
+
+@click.command('')
+@click.option('-i', '--input', 'opt_fp_in', required=True, 
+  help='Path to input image')
+@click.option('-f', '--folder_id', 'opt_folder_id', required=True, 
+  help='ID of folder on Cortex')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_folder_id):
+  """
+  Test uploading a file to Cortex
+  """
+  upload_to_cortex(opt_folder_id, opt_fp_in)
diff --git a/cli/app/commands/cortex/upload_bytes.py b/cli/app/commands/cortex/upload_bytes.py
new file mode 100644
index 0000000..5260b8e
--- /dev/null
+++ b/cli/app/commands/cortex/upload_bytes.py
@@ -0,0 +1,33 @@
+import click
+
+import numpy as np
+from PIL import Image
+from io import BytesIO
+from app.utils.cortex_utils import upload_bytes_to_cortex
+
+@click.command('')
+@click.option('-f', '--folder_id', 'opt_folder_id', required=True, 
+  help='ID of folder on Cortex')
+@click.pass_context
+def cli(ctx, opt_folder_id):
+  """
+  Test uploading a BytesIO file to Cortex
+  """
+  arr = get_gradation_3d(256, 256, (0, 0, 192), (255, 255, 64), (True, False, False))
+  arr = arr.astype(np.uint8)
+  image = Image.fromarray(arr)
+  fp = BytesIO()
+  image.save(fp, format='png')
+  upload_bytes_to_cortex(opt_folder_id, "test.png", fp, "image/png")
+
+def get_gradation_2d(start, stop, width, height, is_horizontal):
+  if is_horizontal:
+    return np.tile(np.linspace(start, stop, width), (height, 1))
+  else:
+    return np.tile(np.linspace(start, stop, height), (width, 1)).T
+
+def get_gradation_3d(width, height, start_list, stop_list, is_horizontal_list):
+  result = np.zeros((height, width, len(start_list)), dtype=np.float)
+  for i, (start, stop, is_horizontal) in enumerate(zip(start_list, stop_list, is_horizontal_list)):
+    result[:, :, i] = get_gradation_2d(start, stop, width, height, is_horizontal)
+  return result
diff --git a/cli/app/search/json.py b/cli/app/search/json.py
index 0f6c71c..beaed3b 100644
--- a/cli/app/search/json.py
+++ b/cli/app/search/json.py
@@ -1,8 +1,9 @@
 from os.path import join
 from app.utils.file_utils import write_json
+from app.search.params import ParamsDict
 
-def save_params_latent(fp_out_dir, tag):
-  data = {
+def make_params_latent(fp_out_dir, tag):
+  return {
     "tag": tag,
     "decay_n": 2, 
     "features": True, 
@@ -43,12 +44,19 @@ def save_params_latent(fp_out_dir, tag):
     "dataset": "inverses/{}/dataset.encodings.hdf5".format(tag),
     "save_progress": True,
   }
+
+def params_latent(fp_out_dir, tag):
+  return ParamsDict(make_params_latent(fp_out_dir, tag))
+
+def save_params_latent(fp_out_dir, tag):
+  data = make_params_latent(fp_out_dir, tag)
   fp_out_fn = join(fp_out_dir, "params_latent.json")
   write_json(data, fp_out_fn)
 
-def save_params_dense(fp_out_dir, tag):
-  data = {
+def make_params_dense(fp_out_dir, tag, folder_id=None):
+  return {
     "tag": tag,
+    "folder_id": folder_id,
     "decay_n": 2, 
     "features": True, 
     "clip": 1.0, 
@@ -56,8 +64,8 @@ def save_params_dense(fp_out_dir, tag):
     "clipping": False, 
     "dataset": "inverses/{}/dataset.encodings.hdf5".format(tag),
     "inv_layer": "Generator_2/G_Z/Reshape:0", 
-    "decay_lr": True, 
-    "inv_it": 15000, 
+    "decay_lr": False, 
+    "inv_it": 6000, 
     "generator_path": "https://tfhub.dev/deepmind/biggan-512/2", 
     "attention_map_layer": "Generator_2/attention/Softmax:0",
     "pre_trained_latent": True, 
@@ -85,8 +93,15 @@ def save_params_dense(fp_out_dir, tag):
     "lambda_reg": 0.1, 
     "dist_loss": True, 
     "sample_size": 4, 
-    "dataset": "inverses/{}/dataset.encodings.dense.hdf5".format(tag),
+    "out_dataset": "inverses/{}/dataset.encodings.dense.hdf5".format(tag),
     "save_progress": True,
+    "max_batches": 0,
   }
+
+def params_dense_dict(fp_out_dir, tag, folder_id=None):
+  return ParamsDict(make_params_dense(fp_out_dir, tag, folder_id))
+
+def save_params_dense(fp_out_dir, tag, folder_id):
+  data = make_params_dense(fp_out_dir, tag, folder_id)
   fp_out_fn = join(fp_out_dir, "params_dense.json")
   write_json(data, fp_out_fn)
diff --git a/cli/app/search/params.py b/cli/app/search/params.py
new file mode 100644
index 0000000..8972436
--- /dev/null
+++ b/cli/app/search/params.py
@@ -0,0 +1,56 @@
+# ------------------------------------------------------------------------------
+# Util class for hyperparams.
+# ------------------------------------------------------------------------------
+
+import json
+
+class Params():
+  """Class that loads hyperparameters from a json file."""
+
+  def __init__(self, json_path):
+    self.update(json_path)
+
+  def save(self, json_path):
+    """Saves parameters to json file."""
+    with open(json_path, 'w') as f:
+      json.dump(self.__dict__, f, indent=4)
+
+  def update(self, json_path):
+    """Loads parameters from json file."""
+    with open(json_path) as f:
+      params = json.load(f)
+      self.__dict__.update(params)
+
+  @property
+  def dict(self):
+    """Gives dict-like access to Params instance."""
+    return self.__dict__
+
+class ParamsDict():
+  """Class that loads hyperparameters from a json file."""
+
+  def __init__(self, data):
+    self.update(data)
+
+  def __setitem__(self, key, item):
+    self.__dict__[key] = item
+
+  def __getitem__(self, key):
+    return self.__dict__[key]
+
+  def __repr__(self):
+    return repr(self.__dict__)
+
+  def __len__(self):
+    return len(self.__dict__)
+
+  def __delitem__(self, key):
+    del self.__dict__[key]
+
+  def save(self, json_path):
+    """Saves parameters to json file."""
+    with open(json_path, 'w') as f:
+      json.dump(self.__dict__, f, indent=4)
+
+  def update(self, *args, **kwargs):
+    return self.__dict__.update(*args, **kwargs)
diff --git a/cli/app/search/search_class.py b/cli/app/search/search_class.py
new file mode 100644
index 0000000..eb9ff42
--- /dev/null
+++ b/cli/app/search/search_class.py
@@ -0,0 +1,153 @@
+from app.settings import app_cfg
+
+import os
+from os.path import join
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
+
+import time
+import numpy as np
+import random
+from subprocess import call
+import cv2 as cv
+from PIL import Image
+from glob import glob
+import tensorflow as tf
+import tensorflow_hub as hub
+import shutil
+import h5py
+
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+from app.search.json import save_params_latent, save_params_dense
+from app.search.image import image_to_uint8, imconvert_uint8, imconvert_float32, \
+  imread, imwrite, imgrid, resize_and_crop_image
+from app.search.vector import truncated_z_sample, truncated_z_single, \
+  create_labels, create_labels_uniform
+
+def find_nearest_vector_for_images(opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag):
+  sess = tf.compat.v1.Session()
+
+  generator = hub.Module('https://tfhub.dev/deepmind/biggan-512/2')
+
+  if os.path.isdir(opt_fp_in):
+    paths = glob(os.path.join(opt_fp_in, '*.jpg')) + \
+      glob(os.path.join(opt_fp_in, '*.jpeg')) + \
+      glob(os.path.join(opt_fp_in, '*.png'))
+  else:
+    paths = [opt_fp_in]
+
+  fp_inverses = os.path.join(app_cfg.DIR_INVERSES, opt_tag)
+  os.makedirs(fp_inverses, exist_ok=True)
+  save_params_latent(fp_inverses, opt_tag)
+  save_params_dense(fp_inverses, opt_tag)
+  out_file = h5py.File(join(fp_inverses, 'dataset.hdf5'), 'w')
+  out_images = out_file.create_dataset('xtrain', (len(paths), 3, 512, 512,), dtype='float32')
+  out_labels = out_file.create_dataset('ytrain', (len(paths), 1000,), dtype='float32')
+  out_latent = out_file.create_dataset('latent', (len(paths), 128,), dtype='float32')
+  out_fns = out_file.create_dataset('fn', (len(paths),), dtype=h5py.string_dtype())
+  for index, path in enumerate(paths):
+    if index == opt_limit:
+      break
+    out_fns[index] = os.path.basename(path)
+    fp_frames = find_nearest_vector(sess, generator, path, opt_dims, out_images, out_labels, out_latent, opt_steps, index)
+    if opt_video:
+      export_video(fp_frames)
+
+def find_nearest_vector(sess, generator, opt_fp_in, opt_dims, out_images, out_labels, out_latent, opt_steps, index):
+  """
+  Find the closest latent and class vectors for an image. Store the class vector in an HDF5.
+  """
+  batch_size = 1
+  truncation = 1.0
+
+  z_dim = 128
+  vocab_size = 1000
+  img_size = 512
+  num_channels = 3
+
+  z_initial = truncated_z_sample(batch_size, z_dim, truncation/2)
+  y_initial = create_labels_uniform(batch_size, vocab_size)
+
+  z_lr = 0.001
+  y_lr = 0.001
+
+  input_z = tf.compat.v1.Variable(z_initial, dtype=np.float32, constraint=lambda t: tf.clip_by_value(t, -2, 2))
+  input_y = tf.compat.v1.Variable(y_initial, dtype=np.float32, constraint=lambda t: tf.clip_by_value(t, 0, 1))
+  input_trunc = tf.compat.v1.constant(1.0)
+  output = generator({
+    'z': input_z,
+    'y': input_y,
+    'truncation': input_trunc,
+  })
+
+  target = tf.compat.v1.placeholder(tf.float32, shape=(batch_size, img_size, img_size, num_channels))
+
+  # loss = tf.losses.compute_weighted_loss(tf.square(output - target), weights=mask)
+  loss = tf.compat.v1.losses.mean_squared_error(target, output)
+
+  train_step_z = tf.train.AdamOptimizer(z_lr).minimize(loss, var_list=[input_z], name='AdamOpterZ')
+  train_step_y = tf.train.AdamOptimizer(y_lr).minimize(loss, var_list=[input_y], name='AdamOpterY')
+
+  target_im, fp_frames = load_target_image(opt_fp_in)
+
+  # crop image and convert to format for next script
+  phi_target_for_inversion = resize_and_crop_image(target_im, 512)
+  b = np.dsplit(phi_target_for_inversion, 3)
+  phi_target_for_inversion = np.stack(b).reshape((3, 512, 512))
+
+  # create phi target for the latent / label pass
+  phi_target = resize_and_crop_image(target_im, opt_dims)
+  phi_target = np.expand_dims(phi_target, 0)
+  phi_target = np.repeat(phi_target, batch_size, axis=0)
+
+  # IMPORTANT: initialize variables before running the session
+  sess.run(tf.compat.v1.global_variables_initializer())
+  sess.run(tf.compat.v1.tables_initializer())
+
+  feed_dict = {
+    target: phi_target,
+  }
+
+  try:
+    print("Preparing to iterate...")
+    for i in range(opt_steps):
+      curr_loss, _, _ = sess.run([loss, train_step_z, train_step_y], feed_dict=feed_dict)
+
+      if i % 20 == 0:
+        phi_guess = sess.run(output)
+        guess_im = imgrid(imconvert_uint8(phi_guess), cols=1)
+        imwrite(join(app_cfg.DIR_OUTPUTS, fp_frames, 'frame_{:04d}.png'.format(int(i / 20))), guess_im)
+        print('iter: {}, loss: {}'.format(i, curr_loss))
+  except KeyboardInterrupt:
+    pass
+
+  z_guess, y_guess = sess.run([input_z, input_y])
+  out_images[index] = phi_target_for_inversion
+  out_labels[index] = y_guess
+  out_latent[index] = z_guess
+  return fp_frames
+
+def export_video(fp_frames):
+  print("Exporting video...")
+  cmd = [
+    '/home/lens/bin/ffmpeg',
+    '-y', # '-v', 'quiet',
+    '-r', '30',
+    '-i', join(app_cfg.DIR_OUTPUTS, fp_frames, 'frame_%04d.png'),
+    '-pix_fmt', 'yuv420p',
+    join(app_cfg.DIR_OUTPUTS, fp_frames + '.mp4')
+  ]
+  # print(' '.join(cmd))
+  call(cmd)
+  shutil.rmtree(join(app_cfg.DIR_OUTPUTS, fp_frames))
+
+def load_target_image(opt_fp_in):
+  print("Loading {}".format(opt_fp_in))
+  fn = os.path.basename(opt_fp_in)
+  fbase, ext = os.path.splitext(fn)
+  fp_frames = "frames_{}_{}".format(fbase, int(time.time() * 1000))
+  fp_frames_fullpath = join(app_cfg.DIR_OUTPUTS, fp_frames)
+  print("Output to {}".format(fp_frames_fullpath))
+  os.makedirs(fp_frames_fullpath, exist_ok=True)
+  target_im = imread(opt_fp_in)
+  return target_im, fp_frames
diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py
new file mode 100644
index 0000000..0086db5
--- /dev/null
+++ b/cli/app/search/search_dense.py
@@ -0,0 +1,510 @@
+import glob
+import h5py
+import itertools
+import numpy as np
+from io import BytesIO
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
+
+from PIL import Image
+import scipy
+import sys
+import tensorflow as tf
+import tensorflow_probability as tfp
+import tensorflow_hub as hub
+import time
+import visualize as vs
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+from app.search.params import Params
+from app.settings import app_cfg
+from app.utils.file_utils import write_pickle
+from app.utils.cortex_utils import upload_bytes_to_cortex
+
+# --------------------------
+# Hyper-parameters.
+# --------------------------
+# Expected parameters:
+#  generator_path: path to generator module.
+#  generator_fixed_inputs: dictionary of fixed generator's input parameters.
+#  dataset: name of the dataset (hdf5 file).
+#  dataset_out: name for the output inverted dataset (hdf5 file).
+# General parameters:
+#  batch_size: number of images inverted at the same time.
+#  inv_it: number of iterations to invert an image.
+#  inv_layer: 'latent' or name of the tensor of the custom layer to be inverted.
+#  lr: learning rate.
+#  decay_lr: exponential decay on the learning rate.
+#  decay_n: number of exponential decays on the learning rate.
+#  custom_grad_relu: replace relus with custom gradient.
+# Logging:
+#  sample_size: number of images included in sampled images.
+#  save_progress: whether to save intermediate images during optimization.
+#  log_z_norm: log the norm of different sections of z.
+#  log_activation_layer: log the percentage of active neurons in this layer.
+# Losses:
+#  mse: use the mean squared error on pixels for image comparison.
+#  features: use features extracted by a feature extractor for image comparison.
+#  feature_extractor_path: path to feature extractor module.
+#  feature_extractor_output: output name from feature extractor.
+#  likeli_loss: regularization loss on the log likelihood of encodings.
+#  norm_loss: regularization loss on the norm of encodings.
+#  dist_loss: whether to include a loss on the dist between g1(z) and enc.
+#  lambda_mse: coefficient for mse loss.
+#  lambda_feat: coefficient for features loss.
+#  lambda_reg: coefficient for regularization loss on latent.
+#  lambda_dist: coefficient for l1 regularization on delta.
+# Latent:
+#  clipping: whether to clip encoding values after every update.
+#  stochastic_clipping: whether to consider stochastic clipping.
+#  clip: clipping bound.
+#  pretrained_latent: load pre trained fixed latent.
+#  fixed_z: do not train the latent vector.
+# Initialization:
+#  init_gen_dist: initialize encodings from the generated distribution.
+#  init_lo: init min value.
+#  init_hi: init max value.
+
+def find_dense_embedding_for_images(params):
+  # --------------------------
+  # Global directories.
+  # --------------------------
+  LATENT_TAG = 'latent' if params.inv_layer == 'latent' else 'dense'
+  BATCH_SIZE = params.batch_size
+  SAMPLE_SIZE = params.sample_size
+  LOGS_DIR = os.path.join('inverses', params.tag, LATENT_TAG, 'logs')
+  SAMPLES_DIR = os.path.join('inverses', params.tag, LATENT_TAG, 'samples')
+  INVERSES_DIR = os.path.join('inverses', params.tag)
+  if not os.path.exists(LOGS_DIR):
+    os.makedirs(LOGS_DIR)
+  if not os.path.exists(SAMPLES_DIR):
+    os.makedirs(SAMPLES_DIR)
+  if not os.path.exists(INVERSES_DIR):
+    os.makedirs(INVERSES_DIR)
+
+  # --------------------------
+  # Util functions.
+  # --------------------------
+  # One hot encoding for classes.
+  def one_hot(values):
+    return np.eye(N_CLASS)[values]
+
+  # --------------------------
+  # Logging.
+  # --------------------------
+  summary_writer = tf.summary.FileWriter(LOGS_DIR)
+  def log_stats(name, val, it):
+    summary = tf.Summary(value=[tf.Summary.Value(tag=name, simple_value=val)])
+    summary_writer.add_summary(summary, it)
+
+  # --------------------------
+  # 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))
+
+  # Step counter.
+  inv_step = tf.get_variable('inv_step', initializer=0, trainable=False)
+
+  # Define target image.
+  IMG_SHAPE = gen_img.get_shape().as_list()[1:]
+  target = tf.get_variable(name='target', dtype=tf.float32,  # normally this is the real [0-255]image
+                           shape=[BATCH_SIZE,] + IMG_SHAPE)
+  # target_img = (tf.cast(target, tf.float32) / 255.) * 2.0 - 1. # Norm to [-1, 1].
+  target_img = target
+
+  # Custom Gradient for Relus.
+  if params.custom_grad_relu:
+    grad_lambda = tf.train.exponential_decay(0.1, inv_step, params.inv_it / 5,
+                                             0.1, staircase=False)
+    @tf.custom_gradient
+    def relu_custom_grad(x):
+      def grad(dy):
+        return tf.where(x >= 0, dy,
+            grad_lambda*tf.where(dy < 0, dy, tf.zeros_like(dy)))
+      return tf.nn.relu(x), grad
+
+    gen_scope = 'module_apply_' + gen_signature + '/'
+    for op in tf.get_default_graph().get_operations():
+      if 'Relu' in op.name and gen_scope in op.name:
+        assert len(op.inputs) == 1
+        assert len(op.outputs) == 1
+        new_out = relu_custom_grad(op.inputs[0])
+        tf.contrib.graph_editor.swap_ts(op.outputs[0], new_out)
+
+  # Operations to clip the values of the encodings.
+  if params.clipping or params.stochastic_clipping:
+    assert params.clip >= 0
+    if params.stochastic_clipping:
+      new_enc = tf.where(tf.abs(latent) >= params.clip,
+          tf.random.uniform([BATCH_SIZE, Z_DIM], minval=-params.clip,
+                            maxval=params.clip), latent)
+    else:
+      new_enc = tf.clip_by_value(latent, -params.clip, params.clip)
+    clip_latent = tf.assign(latent, new_enc)
+
+  # Monitor relu's activation.
+  if params.log_activation_layer:
+    gen_scope = 'module_apply_' + gen_signature + '/'
+    activation_rate = 1.0 - tf.nn.zero_fraction(tf.get_default_graph()\
+        .get_tensor_by_name(gen_scope + params.log_activation_layer))
+
+  # --------------------------
+  # Reconstruction losses.
+  # --------------------------
+  # Mse loss for image comparison.
+  if params.mse:
+    pix_square_diff = tf.square((target_img - gen_img) / 2.0)
+    mse_loss = tf.reduce_mean(pix_square_diff)
+    img_mse_err = tf.reduce_mean(pix_square_diff, axis=[1,2,3])
+  else:
+    mse_loss = tf.constant(0.0)
+    img_mse_err = tf.constant(0.0)
+
+  # Use custom features for image comparison.
+  if params.features:
+    feature_extractor = hub.Module(str(params.feature_extractor_path))
+
+    # Convert images from range [-1, 1] channels_first to [0, 1] channels_last.
+    gen_img_1 = tf.transpose(gen_img / 2.0 + 0.5, [0, 2, 3, 1])
+    target_img_1 = tf.transpose(target_img / 2.0 + 0.5, [0, 2, 3, 1])
+
+    # Convert images to appropriate size for feature extraction.
+    height, width = hub.get_expected_image_size(feature_extractor)
+    gen_img_1 = tf.image.resize_images(gen_img_1, [height, width])
+    target_img_1 = tf.image.resize_images(target_img_1, [height, width])
+
+    gen_feat_ex = feature_extractor(dict(images=gen_img_1), as_dict=True, signature='image_feature_vector')
+    target_feat_ex = feature_extractor(dict(images=target_img_1), as_dict=True, signature='image_feature_vector')
+
+    # gen_feat = gen_feat_ex["InceptionV3/Mixed_7a"]
+    # target_feat = target_feat_ex["InceptionV3/Mixed_7a"]
+    # feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    # feat_loss = tf.reduce_mean(feat_square_diff) * 0.334
+    # img_feat_err = tf.reduce_mean(feat_square_diff, axis=1) * 0.334
+
+    # gen_feat = gen_feat_ex["InceptionV3/Mixed_7b"]
+    # target_feat = target_feat_ex["InceptionV3/Mixed_7b"]
+    # feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    # feat_loss += tf.reduce_mean(feat_square_diff) * 0.333
+    # img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.333
+
+    # gen_feat = gen_feat_ex["InceptionV3/Mixed_7c"]
+    # target_feat = target_feat_ex["InceptionV3/Mixed_7c"]
+    # feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    # feat_loss += tf.reduce_mean(feat_square_diff) * 0.333
+    # img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.333
+
+    # # gen_feat = gen_feat_ex["InceptionV3/Mixed_5a"]
+    # # target_feat = target_feat_ex["InceptionV3/Mixed_5a"]
+    # # feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    # # feat_loss += tf.reduce_mean(feat_square_diff) * 0.16
+    # # img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.16
+
+    # gen_feat = gen_feat_ex["InceptionV3/Mixed_7b"]
+    # target_feat = target_feat_ex["InceptionV3/Mixed_7b"]
+    # feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    # feat_loss += tf.reduce_mean(feat_square_diff) * 0.33
+    # img_feat_err += tf.reduce_mean(feat_square_diff, axis=1)
+
+    # # gen_feat = gen_feat_ex["InceptionV3/Mixed_7c"]
+    # # target_feat = target_feat_ex["InceptionV3/Mixed_7c"]
+    # # feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    # # feat_loss += tf.reduce_mean(feat_square_diff) * 0.17
+    # # img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.17
+
+    # conv1 1, conv1 2, conv3 2 and conv4 2
+    gen_feat = gen_feat_ex["InceptionV3/Conv2d_1a_3x3"]
+    target_feat = target_feat_ex["InceptionV3/Conv2d_1a_3x3"]
+    feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    feat_loss = tf.reduce_mean(feat_square_diff) * 0.15
+    img_feat_err = tf.reduce_mean(feat_square_diff, axis=1) * 0.15
+
+    gen_feat = gen_feat_ex["InceptionV3/Conv2d_2a_3x3"]
+    target_feat = target_feat_ex["InceptionV3/Conv2d_2a_3x3"]
+    feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    feat_loss += tf.reduce_mean(feat_square_diff) * 0.15
+    img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.15
+
+    gen_feat = gen_feat_ex["InceptionV3/Conv2d_3b_1x1"]
+    target_feat = target_feat_ex["InceptionV3/Conv2d_3b_1x1"]
+    feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    feat_loss += tf.reduce_mean(feat_square_diff) * 0.15
+    img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.15
+
+    gen_feat = gen_feat_ex["InceptionV3/Conv2d_4a_3x3"]
+    target_feat = target_feat_ex["InceptionV3/Conv2d_4a_3x3"]
+    feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    feat_loss += tf.reduce_mean(feat_square_diff) * 0.15
+    img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.15
+
+    gen_feat = gen_feat_ex["InceptionV3/Mixed_7a"]
+    target_feat = target_feat_ex["InceptionV3/Mixed_7a"]
+    feat_square_diff = tf.reshape(tf.square(gen_feat - target_feat), [BATCH_SIZE, -1])
+    feat_loss += tf.reduce_mean(feat_square_diff) * 0.4
+    img_feat_err += tf.reduce_mean(feat_square_diff, axis=1) * 0.4
+
+  else:
+    feat_loss = tf.constant(0.0)
+    img_feat_err = tf.constant(0.0)
+
+  # --------------------------
+  # Regularization losses.
+  # --------------------------
+  # Loss on the norm of the encoding.
+  if params.norm_loss:
+    dim = 20
+    chi2_dist = tfp.distributions.Chi2(dim)
+    mode = dim - 2
+    mode_log_prob = chi2_dist.log_prob(mode)
+    norm_loss = 0.0
+    for i in range(int(Z_DIM / dim)):
+      squared_l2 = tf.reduce_sum(tf.square(latent[:,i*dim:(i+1)*dim]), axis=1)
+      over_mode = tf.nn.relu(squared_l2 - mode)
+      norm_loss -= tf.reduce_mean(chi2_dist.log_prob(mode + over_mode))
+      norm_loss += mode_log_prob
+  else:
+    norm_loss = tf.constant(0.0)
+
+  # Loss on the likelihood of the encoding.
+  if params.likeli_loss:
+    norm_dist = tfp.distributions.Normal(0.0, 1.0)
+    likeli_loss = - tf.reduce_mean(norm_dist.log_prob(latent))
+    mode_log_prob = norm_dist.log_prob(0.0)
+    likeli_loss += mode_log_prob
+  else:
+    likeli_loss = tf.constant(0.0)
+
+  # Regularization loss.
+  reg_loss = norm_loss + likeli_loss
+
+  # Loss on the l1 distance between gen_encoding and inverted encoding.
+  if params.dist_loss:
+    dist_loss = tf.reduce_mean(tf.abs(encoding - gen_encoding))
+  else:
+    dist_loss = tf.constant(0.0)
+
+  # Per image reconstruction error.
+  img_rec_err = params.lambda_mse * img_mse_err\
+      + params.lambda_feat * img_feat_err
+
+  # Batch reconstruction error.
+  rec_loss = params.lambda_mse * mse_loss + params.lambda_feat * feat_loss
+
+  # Total inversion loss.
+  inv_loss = rec_loss + params.lambda_reg * reg_loss\
+      + params.lambda_dist * dist_loss
+
+  # --------------------------
+  # Optimizer.
+  # --------------------------
+  if params.decay_lr:
+    lrate = tf.train.exponential_decay(params.lr, inv_step,
+        params.inv_it / params.decay_n, 0.1, staircase=True)
+  else:
+    lrate = tf.constant(params.lr)
+  trained_params = [encoding] if params.fixed_z else [latent, encoding]
+  optimizer = tf.train.AdamOptimizer(learning_rate=lrate, beta1=0.9, beta2=0.999)
+  inv_train_op = optimizer.minimize(inv_loss, var_list=trained_params,
+                                    global_step=inv_step)
+  reinit_optimizer = tf.variables_initializer(optimizer.variables())
+
+  # --------------------------
+  # Noise source.
+  # --------------------------
+  def noise_sampler():
+    return np.random.normal(size=[BATCH_SIZE, Z_DIM])
+
+  def small_init(shape=[BATCH_SIZE, Z_DIM]):
+    return np.random.uniform(low=params.init_lo, high=params.init_hi, size=shape)
+
+  # --------------------------
+  # Dataset.
+  # --------------------------
+  if params.dataset.endswith('.hdf5'):
+    in_file = h5py.File(params.dataset, 'r')
+    sample_images = in_file['xtrain'][()]
+    sample_labels = in_file['ytrain'][()]
+    sample_fns = in_file['fn'][()]
+    NUM_IMGS = sample_images.shape[0] # number of images to be inverted.
+    print("Number of images: {}".format(NUM_IMGS))
+    print("Batch size: {}".format(BATCH_SIZE))
+    def sample_images_gen():
+      for i in range(int(NUM_IMGS / BATCH_SIZE)):
+        i_1, i_2 = i*BATCH_SIZE, (i+1)*BATCH_SIZE
+        yield sample_images[i_1:i_2], sample_labels[i_1:i_2]
+    image_gen = sample_images_gen()
+    sample_latents = in_file['latent']
+    def sample_latent_gen():
+      for i in range(int(NUM_IMGS / BATCH_SIZE)):
+        i_1, i_2 = i*BATCH_SIZE, (i+1)*BATCH_SIZE
+        yield sample_latents[i_1:i_2]
+    latent_gen = sample_latent_gen()
+    if NUM_IMGS % BATCH_SIZE != 0:
+      REMAINDER = BATCH_SIZE - (NUM_IMGS % BATCH_SIZE)
+      NUM_IMGS += REMAINDER
+      sample_images = np.append(sample_images, sample_images[-REMAINDER:,...], axis=0)
+      sample_labels = np.append(sample_labels, sample_labels[-REMAINDER:,...], axis=0)
+      sample_latents = np.append(sample_latents, sample_latents[-REMAINDER:,...], axis=0)
+      sample_fns = np.append(sample_fns, sample_fns[-REMAINDER:], axis=0)
+    assert(NUM_IMGS % BATCH_SIZE == 0)
+  else:
+    sys.exit('Unknown dataset {}.'.format(params.dataset))
+
+  # --------------------------
+  # Training.
+  # --------------------------
+  # Start session.
+  sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
+  sess.run(tf.global_variables_initializer())
+  sess.run(tf.tables_initializer())
+
+  if params.max_batches > 0:
+    NUM_IMGS_TO_PROCESS = params.max_batches * BATCH_SIZE
+  else:
+    NUM_IMGS_TO_PROCESS = NUM_IMGS
+
+  # Output file.
+  out_file = h5py.File(os.path.join(INVERSES_DIR, params.out_dataset), 'w')
+  out_images = out_file.create_dataset('xtrain', [NUM_IMGS_TO_PROCESS,] + IMG_SHAPE, dtype='float32')
+  out_enc = out_file.create_dataset('encoding', [NUM_IMGS_TO_PROCESS,] + ENC_SHAPE, dtype='float32')
+  out_lat = out_file.create_dataset('latent', [NUM_IMGS_TO_PROCESS, Z_DIM], dtype='float32')
+  out_fns = out_file.create_dataset('fn', [NUM_IMGS_TO_PROCESS], dtype=h5py.string_dtype())
+  if COND_GAN:
+    out_labels = out_file.create_dataset('ytrain', (NUM_IMGS_TO_PROCESS, N_CLASS,), dtype='float32')
+  out_err = out_file.create_dataset('err', (NUM_IMGS_TO_PROCESS,))
+
+  out_fns[:] = sample_fns[:NUM_IMGS_TO_PROCESS]
+
+  # Gradient descent w.r.t. generator's inputs.
+  it = 0
+  out_pos = 0
+  start_time = time.time()
+
+  for image_batch, label_batch in image_gen:
+    sess.run([
+      target.assign(image_batch),
+      label.assign(label_batch),
+      latent.assign(next(latent_gen)),
+      inv_step.assign(0),
+    ])
+    sess.run([
+      encoding.assign(gen_encoding),
+      reinit_optimizer,
+    ])
+
+    # Main optimization loop.
+    print("Total iterations: {}".format(params.inv_it))
+    for _ in range(params.inv_it):
+
+      _inv_loss, _mse_loss, _feat_loss, _rec_loss, _reg_loss, _dist_loss,\
+          _lrate, _ = sess.run([inv_loss, mse_loss, feat_loss,
+          rec_loss, reg_loss, dist_loss, lrate, inv_train_op])
+
+      if params.clipping or params.stochastic_clipping:
+        sess.run(clip_latent)
+
+      # Save logs with training information.
+      if it % 500 == 0:
+        # Log losses.
+        etime = time.time() - start_time
+        print('It [{:8d}] time [{:5.1f}] total [{:.4f}] mse [{:.4f}] '
+              'feat [{:.4f}] rec [{:.4f}] reg [{:.4f}] dist [{:.4f}] '
+              'lr [{:.4f}]'.format(it, etime, _inv_loss, _mse_loss,
+              _feat_loss, _rec_loss, _reg_loss, _dist_loss, _lrate))
+
+        sys.stdout.flush()
+
+        # Save target images and reconstructions.
+        if params.save_progress:
+          assert SAMPLE_SIZE <= BATCH_SIZE
+          gen_time = time.time()
+          gen_images  = sess.run(gen_img)
+          print("Generation time: {:.1f}s".format(time.time() - gen_time))
+          inv_batch = vs.interleave(vs.data2img(image_batch[BATCH_SIZE - SAMPLE_SIZE:]),
+                          vs.data2img(gen_images[BATCH_SIZE - SAMPLE_SIZE:]))
+          inv_batch = vs.grid_transform(inv_batch)
+          vs.save_image('{}/progress_{}_{:04d}.png'.format(SAMPLES_DIR, params.tag, int(it / 500)), inv_batch)
+
+      it += 1
+
+    # Save images that are ready.
+    latent_batch, enc_batch, rec_err_batch = sess.run([latent, encoding, img_rec_err])
+    out_lat[out_pos:out_pos+BATCH_SIZE] = latent_batch
+    out_enc[out_pos:out_pos+BATCH_SIZE] = enc_batch
+    out_images[out_pos:out_pos+BATCH_SIZE] = image_batch
+    out_labels[out_pos:out_pos+BATCH_SIZE] = label_batch
+    out_err[out_pos:out_pos+BATCH_SIZE] = rec_err_batch
+
+    gen_images  = sess.run(gen_img)
+    images = vs.data2img(gen_images)
+
+    # write encoding, latent to pkl file
+    for i in range(BATCH_SIZE):
+      out_i = out_pos + i
+      fn, ext = os.path.splitext(sample_fns[out_i])
+      fp_out_pkl = os.path.join(app_cfg.INVERSES_DIR, fn  ".pkl")
+      out_data = {
+        'id': fn,
+        'latent': out_lat[out_i],
+        'encoding': out_enc[out_i],
+        'label': out_labels[out_i],
+      }
+      write_pickle(out_data, fp_out_pkl)
+      image = Image.fromarray(images[i])
+      fp = BytesIO()
+      image.save(fp, format='png')
+      upload_bytes_to_cortex(params.folder_id, fn, fp, 'image/png')
+
+    out_pos += BATCH_SIZE
+    if params.max_batches > 0 and (out_pos / BATCH_SIZE) >= params.max_batches:
+      break
+
+  print('Mean reconstruction error: {}'.format(np.mean(out_err)))
+  print('Stdev reconstruction error: {}'.format(np.std(out_err)))
+  print('End of inversion.')
+  out_file.close()
+  sess.close()
diff --git a/cli/app/search/search_km.py b/cli/app/search/search_km.py
deleted file mode 100644
index bdffbe4..0000000
--- a/cli/app/search/search_km.py
+++ /dev/null
@@ -1,86 +0,0 @@
-import cStringIO
-import numpy as np
-import PIL.Image
-from scipy.stats import truncnorm
-import tensorflow as tf
-import tensorflow_hub as hub
-import cv2
-
-module_path = 'https://tfhub.dev/deepmind/biggan-128/2'  # 128x128 BigGAN
-# module_path = 'https://tfhub.dev/deepmind/biggan-256/2'  # 256x256 BigGAN
-# module_path = 'https://tfhub.dev/deepmind/biggan-512/2'  # 512x512 BigGAN
-
-tf.reset_default_graph()
-module = hub.Module(module_path)
-inputs = {k: tf.placeholder(v.dtype, v.get_shape().as_list(), k)
-          for k, v in module.get_input_info_dict().iteritems()}
-output = module(inputs)
-
-input_z = inputs['z']
-input_y = inputs['y']
-input_trunc = inputs['truncation']
-
-dim_z = input_z.shape.as_list()[1]
-vocab_size = input_y.shape.as_list()[1]
-
-initializer = tf.global_variables_initializer()
-sess = tf.Session()
-sess.run(initializer)
-
-y = 259 # pomeranian
-n_samples = 9
-truncation = 0.5
-
-# phi_target = imread(uploaded.keys()[0])
-# phi_target = imconvert_float32(phi_target)
-# phi_target = np.expand_dims(phi_target, 0)
-# phi_target = phi_target[:128,:128]
-# phi_target = np.repeat(phi_target, n_samples, axis=0)
-
-label = one_hot([y] * n_samples, vocab_size)
-
-# use z from manifold
-if uploaded is not None:
-  z_target = np.repeat(truncated_z_sample(1, truncation, 0), n_samples, axis=0)
-  feed_dict = {input_z: z_target, input_y: label, input_trunc: truncation}
-  phi_target = sess.run(output, feed_dict=feed_dict)
-
-target_im = imgrid(imconvert_uint8(phi_target), cols=3)
-cost = tf.reduce_sum(tf.pow(output - phi_target, 2))
-dc_dz, = tf.gradients(cost, [input_z])
-
-lr = 0.0001
-z_guess = np.asarray(truncated_z_sample(n_samples, truncation/2, 1))
-feed_dict = {input_z: z_guess, input_y: label, input_trunc: truncation}
-phi_impostor = sess.run(output, feed_dict=feed_dict)
-impostor_im = imgrid(imconvert_uint8(phi_impostor), cols=3)
-comparison = None
-
-try:
-  for i in range(1000):
-    feed_dict = {input_z: z_guess, input_y: label, input_trunc: truncation}
-    grad = dc_dz.eval(session=sess, feed_dict=feed_dict)
-    z_guess -= grad * lr
-
-    # decay/attenuate learning rate to 0.05 of the original over 1000 frames
-    lr *= 0.997
-
-    indices = np.logical_or(z_guess <= -2*truncation, z_guess >= +2*truncation)
-    z_guess[indices] = np.random.randn(np.count_nonzero(indices))
-
-    feed_dict = {input_z: z_guess, input_y: label, input_trunc: truncation}
-    phi_guess = sess.run(output, feed_dict=feed_dict)
-    guess_im = imgrid(imconvert_uint8(phi_guess), cols=3)
-
-    imwrite('frames/{:06d}.png'.format(i), guess_im)
-
-    # display the progress every 10 frames
-    if i % 10 == 0:
-      comparison = imgrid(np.asarray([impostor_im, guess_im, target_im]), cols=3, pad=10)
-
-      # clear_output(wait=True)
-      print('lr: {}, iter: {}, grad_std: {}'.format(lr, i, np.std(grad)))
-      imshow(comparison, format='jpeg')
-except KeyboardInterrupt:
-  pass
-
diff --git a/cli/app/settings/app_cfg.py b/cli/app/settings/app_cfg.py
index 6739da7..7e739df 100644
--- a/cli/app/settings/app_cfg.py
+++ b/cli/app/settings/app_cfg.py
@@ -20,6 +20,7 @@ codecs.register(lambda name: codecs.lookup('utf8') if name == 'utf8mb4' else Non
 CLICK_GROUPS = {
   'biggan': 'app/commands/biggan',
   'bigbigan': 'app/commands/bigbigan',
+  'cortex': 'app/commands/cortex',
   'process': 'app/commands/process',
 }
 
@@ -31,6 +32,7 @@ SELF_CWD = os.path.dirname(os.path.realpath(__file__))  # Script CWD
 DIR_APP = str(Path(SELF_CWD).parent.parent.parent)
 DIR_IMAGENET = join(DIR_APP, 'data_store/imagenet')
 DIR_INVERSES = join(DIR_APP, 'data_store/inverses')
+DIR_INPUTS = join(DIR_APP, 'data_store/inputs')
 DIR_OUTPUTS = join(DIR_APP, 'data_store/outputs')
 FP_MODELZOO = join(DIR_APP, 'modelzoo/modelzoo.yaml')
 
diff --git a/cli/app/utils/cortex_utils.py b/cli/app/utils/cortex_utils.py
new file mode 100644
index 0000000..328b93a
--- /dev/null
+++ b/cli/app/utils/cortex_utils.py
@@ -0,0 +1,64 @@
+import os
+from os.path import join
+import requests
+import urllib3
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+from app.settings import app_cfg
+
+def api_url(path):
+  return "https://lens.neural.garden/api/{}/".format(path)
+
+def fetch_cortex_folder(opt_folder_id):
+  rows = fetch_json(api_url('file'), folder_id=opt_folder_id)
+  fp_out_dir = join(app_cfg.DIR_INPUTS, "cortex", str(opt_folder_id))
+  os.makedirs(fp_out_dir, exist_ok=True)
+  for row in rows:
+    if row['generated'] == 0 and row['processed'] != 1:
+      fn, ext = os.path.splitext(row['name'])
+      fp_out_image = join(fp_out_dir, "{}{}".format(row['id'], ext))
+      if not os.path.exists(fp_out_image):
+        fetch_file(row['url'], fp_out_image)
+
+def fetch_json(url, **kwargs):
+  resp = requests.get(url, params=kwargs, verify=False, timeout=10)
+  return None if resp.status_code != 200 else resp.json()
+
+def fetch_file(url, fn, **kwargs):
+  print("Fetch {} => {}".format(url, fn))
+  try:
+    resp = requests.get(url, params=kwargs, verify=False, timeout=10)
+    if resp.status_code != 200:
+      return None
+  except:
+    return None
+  size = 0
+  with open(fn, 'wb') as f:
+    for chunk in resp.iter_content(chunk_size=1024):
+      if chunk:
+        size += len(chunk)
+        f.write(chunk)
+  return size
+
+def upload_fp_to_cortex(opt_folder_id, fp):
+  files = {
+    'file': fp
+  }
+  data = {
+    'folder_id': opt_folder_id,
+    'generated': 'true',
+    'module': 'biggan',
+    'activity': 'invert',
+    'datatype': 'image',
+  }
+  url = os.path.join(api_url('folder'), opt_folder_id, 'upload/')
+  print(url)
+  r = requests.post(url, files=files, data=data)
+  print(r.json())
+
+def upload_bytes_to_cortex(opt_folder_id, fn, fp, mimetype):
+  upload_fp_to_cortex(opt_folder_id, (fn, fp.getvalue(), mimetype,))
+
+def upload_file_to_cortex(opt_folder_id, fn):
+  with open(fn, 'rb') as fp:
+    upload_fp_to_cortex(opt_folder_id, fp)
diff --git a/inversion/image_inversion_placeholder.py b/inversion/image_inversion_placeholder.py
index 64929cc..88af8ef 100644
--- a/inversion/image_inversion_placeholder.py
+++ b/inversion/image_inversion_placeholder.py
@@ -9,8 +9,8 @@ import itertools
 import numpy as np
 import os
 import params
-import PIL
 import scipy
+import pickle
 import sys
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -485,6 +485,9 @@ for image_batch, label_batch in image_gen:
   out_labels[out_pos:out_pos+BATCH_SIZE] = label_batch
   out_err[out_pos:out_pos+BATCH_SIZE] = rec_err_batch
   out_pos += BATCH_SIZE
+  # upload image
+  # write encoding, latent to pkl file
+
   if params.max_batches > 0 and (out_pos / BATCH_SIZE) >= params.max_batches:
     break