store fns in output

1 files changed, 5 insertions, 1 deletions

diff --git a/cli/app/commands/biggan/search.py b/cli/app/commands/biggan/search.py
index d2b5900..9e43458 100644
--- a/cli/app/commands/biggan/search.py
+++ b/cli/app/commands/biggan/search.py
@@ -29,13 +29,15 @@ from app.search.vector import truncated_z_sample, truncated_z_single, create_lab
   help='Dimensions of BigGAN network (128, 256, 512)')
 @click.option('-s', '--steps', 'opt_steps', default=500, type=int,
   help='Number of optimization iterations')
+@click.option('-l', '--limit', 'opt_limit', default=1000, type=int,
+  help='Limit the number of images to process')
 @click.option('-v', '--video', 'opt_video', is_flag=True,
   help='Export a video for each dataset')
 @click.option('-t', '--tag', 'opt_tag', default='inverse_' + str(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_steps, opt_video, opt_tag):
+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
   """
@@ -70,6 +72,8 @@ def cli(ctx, opt_fp_in, opt_dims, opt_steps, opt_video, opt_tag):
   out_labels = out_file.create_dataset('ytrain', (len(paths), vocab_size,), 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(generator, sess, input_z, input_y, input_trunc, output, path, opt_dims, out_images, out_labels, opt_steps, index)
     if opt_video: