write json params

1 files changed, 8 insertions, 3 deletions

diff --git a/cli/app/commands/biggan/search.py b/cli/app/commands/biggan/search.py
index 0b31a0d..5844855 100644
--- a/cli/app/commands/biggan/search.py
+++ b/cli/app/commands/biggan/search.py
@@ -17,6 +17,7 @@ import tensorflow_hub as hub
 import shutil
 import h5py
 
+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
 from app.search.vector import truncated_z_sample, truncated_z_single, create_labels
@@ -28,11 +29,13 @@ from app.search.vector import truncated_z_sample, truncated_z_single, create_lab
   help='Dimensions of BigGAN network (128, 256, 512)')
 @click.option('-v', '--video', 'opt_video', is_flag=True,
   help='Export a video for each dataset')
+@click.option('-t', '--tag', 'opt_tag', default='inverse_' + int(time.time() * 1000), 
+  help='Tag this dataset')
 # @click.option('-r', '--recursive', 'opt_recursive', is_flag=True)
 @click.pass_context
 def cli(ctx, opt_fp_in, opt_dims, opt_video):
   """
-  Search for an image in BigGAN using gradient descent
+  Search for an image (class vector) in BigGAN using gradient descent
   """
   generator = hub.Module('https://tfhub.dev/deepmind/biggan-' + str(opt_dims) + '/2')
 
@@ -56,9 +59,11 @@ def cli(ctx, opt_fp_in, opt_dims, opt_video):
   else:
     paths = [opt_fp_in]
 
-  fp_inverses = os.path.join(app_cfg.INVERSES_DIR, 'inverse_' + int(time.time() * 1000))
+  fp_inverses = os.path.join(app_cfg.INVERSES_DIR, opt_tag)
   os.makedirs(fp_inverses, exist_ok=True)
-  out_file = h5py.File(fp_inverses, 'w')
+  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), vocab_size,), dtype='float32')