clip latents differently

1 files changed, 5 insertions, 4 deletions

diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py
index ac4dc59..46183c7 100644
--- a/cli/app/search/search_dense.py
+++ b/cli/app/search/search_dense.py
@@ -124,6 +124,8 @@ def find_dense_embedding_for_images(params):
     else:
       gen_img = generator(latent, signature=gen_signature)
 
+  gen_img_orig = gen_img
+  
   # Convert generated image to channels_first.
   gen_img = tf.transpose(gen_img, [0, 3, 1, 2])
 
@@ -393,7 +395,7 @@ def find_dense_embedding_for_images(params):
     ])
 
     # Main optimization loop.
-    print("Total iterations: {}".format(params.inv_it))
+    print("Beginning dense iteration...")
     for _ in range(params.inv_it):
 
       _inv_loss, _mse_loss, _feat_loss,\
@@ -435,15 +437,14 @@ def find_dense_embedding_for_images(params):
     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)
+    gen_images  = sess.run(gen_img_orig)
     images = vs.data2img(gen_images)
 
     # write encoding, latent to pkl file
     for i in range(BATCH_SIZE):
       out_i = out_pos + i
       sample_fn, ext = os.path.splitext(sample_fns[out_i])
-      image = Image.fromarray(images[i])
-      image = vs.grid_transform(image)
+      image = Image.fromarray(image)
       fp = BytesIO()
       image.save(fp, format='png')
       data = upload_bytes_to_cortex(params.folder_id, sample_fn + "-inverse.png", fp, "image/png")