search_class_for_folders

1 files changed, 13 insertions, 2 deletions

diff --git a/cli/app/commands/biggan/search_class.py b/cli/app/commands/biggan/search_class.py
index fc27935..60b3542 100644
--- a/cli/app/commands/biggan/search_class.py
+++ b/cli/app/commands/biggan/search_class.py
@@ -3,6 +3,8 @@ import click
 from app.search.search_class import find_nearest_vector_for_images
 
 @click.command('')
+@click.option('-f', '--folder_id', 'opt_folder_id', type=int, default=0,
+  help='Folder ID to process')
 @click.option('-i', '--input', 'opt_fp_in', required=True, 
   help='Path to input image')
 @click.option('-d', '--dims', 'opt_dims', default=512, type=int,
@@ -29,12 +31,21 @@ from app.search.search_class import find_nearest_vector_for_images
   help='Interval to store sample images')
 
 @click.pass_context
-def cli(ctx, opt_fp_in, opt_dims, opt_steps, opt_video, opt_tag,
+def cli(ctx, opt_folder_id, 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
   """
-  if os.path.isdir(opt_fp_in):
+  if opt_folder_id != 0:
+    folder = cortex_folder(opt_folder_id)
+    files = download_cortex_files(opt_folder_id)
+    unprocessed_files = [file for file in files if file['generated'] == 0 and file['type'] == 'image']
+    if len(unprocessed_files) == 0:
+      print("All files processed, nothing to do")
+      return
+    print("Processing folder {} ({}), {} new files".format(folder['name'], folder['id'], len(unprocessed_files)))
+    paths = [file['path'] for file in unprocessed_files]
+  elif os.path.isdir(opt_fp_in):
     paths = glob(os.path.join(opt_fp_in, '*.jpg')) + \
       glob(os.path.join(opt_fp_in, '*.jpeg')) + \
       glob(os.path.join(opt_fp_in, '*.png'))