2 files changed, 37 insertions, 5 deletions
diff --git a/cli/app/commands/biggan/extract_dense_vectors.py b/cli/app/commands/biggan/extract_dense_vectors.py
index 2436ce6..7dc25bc 100644
--- a/cli/app/commands/biggan/extract_dense_vectors.py
+++ b/cli/app/commands/biggan/extract_dense_vectors.py
@@ -15,8 +15,19 @@ from app.search.json import params_dense_dict
   help='Number of optimization iterations')
 @click.option('-v', '--video', 'opt_video', is_flag=True,
   help='Export a video for each dataset')
+@click.option('-sc', '--stochastic_clipping', 'opt_stochastic_clipping', default=0
+  help='Compute feature loss')
+@click.option('-lc', '--label_clipping', 'opt_label_clipping', default=0,
+  help='Normalize labels every N steps')
+@click.option('-feat', '--use_feature_detector', 'opt_use_feature_detector', is_flag=True,
+  help='Compute feature loss')
+@click.option('-ll', '--feature_layers', 'opt_feature_layers', default="1a,2a,4a,7a"
+  help='Feature layers used for loss')
+@click.option('-snap', '--snapshot_interval', 'opt_snapshot_interval', default=20
+  help='Interval to store sample images')
 @click.pass_context
-def cli(ctx, opt_folder_id, opt_latent_steps, opt_dense_steps, opt_video):
+def cli(ctx, opt_folder_id, opt_latent_steps, opt_dense_steps, opt_video,
+  opt_stochastic_clipping, opt_label_clipping, opt_use_feature_detector, opt_feature_layers, opt_snapshot_interval):
   """
   The full process:
   - Fetch new images from the cortex
@@ -42,7 +53,12 @@ def cli(ctx, opt_folder_id, opt_latent_steps, opt_dense_steps, opt_video):
     opt_steps=opt_dense_steps,
     opt_video=opt_video,
     opt_tag=tag,
-    opt_limit=-1
+    opt_limit=-1,
+    opt_stochastic_clipping=opt_stochastic_clipping,
+    opt_label_clipping=opt_label_clipping,
+    opt_use_feature_detector=opt_use_feature_detector,
+    opt_feature_layers=opt_feature_layers,
+    opt_snapshot_interval=opt_snapshot_interval
   )
 
   params = params_dense_dict(tag)
diff --git a/cli/app/commands/biggan/search_class.py b/cli/app/commands/biggan/search_class.py
index 311dc70..fc27935 100644
--- a/cli/app/commands/biggan/search_class.py
+++ b/cli/app/commands/biggan/search_class.py
@@ -13,11 +13,24 @@ from app.search.search_class import find_nearest_vector_for_images
   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('-d', '--dims', 'opt_dims', default=512, type=int,
+  help='Dimensions of BigGAN network (128, 256, 512)')
 @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.option('-sc', '--stochastic_clipping', 'opt_stochastic_clipping', default=0
+  help='Compute feature loss')
+@click.option('-lc', '--label_clipping', 'opt_label_clipping', default=0,
+  help='Normalize labels every N steps')
+@click.option('-feat', '--use_feature_detector', 'opt_use_feature_detector', is_flag=True,
+  help='Compute feature loss')
+@click.option('-ll', '--feature_layers', 'opt_feature_layers', default="1a,2a,4a,7a"
+  help='Feature layers used for loss')
+@click.option('-snap', '--snapshot_interval', 'opt_snapshot_interval', default=20
+  help='Interval to store sample images')
+
 @click.pass_context
-def cli(ctx, opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag):
+def cli(ctx, opt_fp_in, opt_dims, opt_steps, opt_video, opt_tag,
+  opt_stochastic_clipping, opt_label_clipping, opt_use_feature_detector, opt_feature_layers, opt_snapshot_interval):
   """
   Search for an image (class vector) in BigGAN using gradient descent
   """
@@ -28,4 +41,7 @@ def cli(ctx, opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag):
   else:
     paths = [opt_fp_in]
 
-  find_nearest_vector_for_images(paths, opt_dims, opt_steps, opt_video, opt_tag, opt_limit)
+  opt_feature_layers = opt_feature_layers.split(',')
+
+  find_nearest_vector_for_images(paths, opt_dims, opt_steps, opt_video, opt_tag, opt_limit,
+    opt_stochastic_clipping, opt_label_clipping, opt_use_feature_detector, opt_feature_layers)