summaryrefslogtreecommitdiff
path: root/cli/app/commands/biggan/search_class.py
diff options
context:
space:
mode:
authorJules Laplace <julescarbon@gmail.com>2020-01-07 20:05:57 +0100
committerJules Laplace <julescarbon@gmail.com>2020-01-07 20:05:57 +0100
commit08fa371f49042a2989ec3d494f37b8db63b13c95 (patch)
treedefd020ea3e720e1b58d34c4fb4c16b712cad54c /cli/app/commands/biggan/search_class.py
parent45855c051f415e56306f116a921eefe284139694 (diff)
script to run everything
Diffstat (limited to 'cli/app/commands/biggan/search_class.py')
-rw-r--r--cli/app/commands/biggan/search_class.py9
1 files changed, 8 insertions, 1 deletions
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)