cleanup

1 files changed, 4 insertions, 1 deletions

diff --git a/cli/app/commands/biggan/search.py b/cli/app/commands/biggan/search.py
index 3479680..aeda9f5 100644
--- a/cli/app/commands/biggan/search.py
+++ b/cli/app/commands/biggan/search.py
@@ -15,6 +15,7 @@ from PIL import Image
 from glob import glob
 import tensorflow as tf
 import tensorflow_hub as hub
+import shutil
 
 def image_to_uint8(x):
   """Converts [-1, 1] float array to [0, 255] uint8."""
@@ -136,6 +137,7 @@ def find_nearest_vector(generator, sess, input_z, input_y, input_trunc, output,
   y = create_labels(batch_size, vocab_size, num_classes)
 
   if opt_fp_in:
+    print("Processing {}".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))
@@ -156,6 +158,7 @@ def find_nearest_vector(generator, sess, input_z, input_y, input_trunc, output,
     phi_target = np.expand_dims(phi_target, 0)
     phi_target = np.repeat(phi_target, batch_size, axis=0)
   else:
+    print("Processing random vector")
     fp_frames = "frames_{}".format(int(time.time() * 1000))
     os.makedirs(join(app_cfg.DIR_OUTPUTS, fp_frames), exist_ok=True)
     z_target = np.repeat(truncated_z_single(z_dim, truncation), batch_size, axis=0)
@@ -224,5 +227,5 @@ def find_nearest_vector(generator, sess, input_z, input_y, input_trunc, output,
   ]
   print(' '.join(cmd))
   call(cmd)
-  print("Done")
+  shutil.rmtree(join(app_cfg.DIR_OUTPUTS, fp_frames))