3 files changed, 59 insertions, 3 deletions
diff --git a/cli/app/commands/biggan/extract_dense_vectors.py b/cli/app/commands/biggan/extract_dense_vectors.py
new file mode 100644
index 0000000..2436ce6
--- /dev/null
+++ b/cli/app/commands/biggan/extract_dense_vectors.py
@@ -0,0 +1,49 @@
+import click
+import os
+
+from app.utils.cortex_utils import fetch_cortex_folder, find_unprocessed_files
+from app.search.search_class import find_nearest_vector_for_images
+from app.search.search_dense import find_dense_embedding_for_images
+from app.search.json import params_dense_dict
+
+@click.command('')
+@click.option('-f', '--folder_id', 'opt_folder_id', type=int,
+  help='Folder ID to process')
+@click.option('-ls', '--latent_steps', 'opt_latent_steps', default=2000, type=int,
+  help='Number of optimization iterations')
+@click.option('-ds', '--dense_steps', 'opt_dense_steps', default=2000, type=int,
+  help='Number of optimization iterations')
+@click.option('-v', '--video', 'opt_video', is_flag=True,
+  help='Export a video for each dataset')
+@click.pass_context
+def cli(ctx, opt_folder_id, opt_latent_steps, opt_dense_steps, opt_video):
+  """
+  The full process:
+  - Fetch new images from the cortex
+  - Extract labels and base latents
+  - Extract dense embeddings
+  - Upload extract images to the cortex
+  """
+  folder = cortex_folder(opt_folder_id)
+  files = download_cortex_files(opt_folder_id)
+  unprocessed_files = find_unprocessed_files(files)
+  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)))
+
+  tag = "folder_{}".format(folder['id'])
+  paths = [file['path'] for file in unprocessed_files]
+  
+  find_nearest_vector_for_images(
+    paths=paths,
+    opt_dims=512,
+    opt_steps=opt_dense_steps,
+    opt_video=opt_video,
+    opt_tag=tag,
+    opt_limit=-1
+  )
+
+  params = params_dense_dict(tag)
+  find_dense_embedding_for_images(params)
diff --git a/cli/app/commands/biggan/fetch.py b/cli/app/commands/biggan/fetch.py
index 39f503c..d454633 100644
--- a/cli/app/commands/biggan/fetch.py
+++ b/cli/app/commands/biggan/fetch.py
@@ -9,9 +9,9 @@ from app.utils.cortex_utils import fetch_cortex_folder, find_unprocessed_files
 @click.pass_context
 def cli(ctx, opt_folder_id):
   """
-  Fetch JSON from the server
+  Fetch new images from the server
   """
-  files = fetch_cortex_folder(opt_folder_id)
+  files = download_cortex_files(opt_folder_id)
   unprocessed_files = find_unprocessed_files(files)
   print("Unprocessed files:")
   for file in unprocessed_files:
diff --git a/cli/app/commands/biggan/search_class.py b/cli/app/commands/biggan/search_class.py
index 0501729..311dc70 100644
--- a/cli/app/commands/biggan/search_class.py
+++ b/cli/app/commands/biggan/search_class.py
@@ -21,4 +21,11 @@ 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
   """
-  find_nearest_vector_for_images(opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag)
+  if 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'))
+  else:
+    paths = [opt_fp_in]
+
+  find_nearest_vector_for_images(paths, opt_dims, opt_steps, opt_video, opt_tag, opt_limit)