implemented selectable interpolation and loop settings

1 files changed, 20 insertions, 9 deletions

diff --git a/gantools/cli.py b/gantools/cli.py
index 0ad3a8d..3e42f2e 100644
--- a/gantools/cli.py
+++ b/gantools/cli.py
@@ -4,26 +4,37 @@ from gantools import biggan
 from gantools import latent_space
 from gantools import image_utils
 
-# create entrypoints for cli tools
-def main():
-    ## handle args
-    parser = argparse.ArgumentParser(description='GAN tools')
+def handle_args(argv=None):
+    parser = argparse.ArgumentParser(
+            description='GAN tools',
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter
+            )
     # load from ganbreeder
-    parser.add_argument('-u', '--username', help='Ganbreeder account email address/username.')
-    parser.add_argument('-p', '--password', help='Ganbreeder account password.')
-    parser.add_argument('-k', '--keys', nargs='+', help='Ganbreeder keys.')
+    ganbreeder_group = parser.add_argument_group(title='GANbreeder login')
+    ganbreeder_group.add_argument('-u', '--username', help='Ganbreeder account email address/username.')
+    ganbreeder_group.add_argument('-p', '--password', help='Ganbreeder account password.')
+    ganbreeder_group.add_argument('-k', '--keys', nargs='+', help='Ganbreeder keys.')
     parser.add_argument('-n', '--nframes', metavar='N', type=int, help='Total number of frames in the final animation.', default=10)
     parser.add_argument('-b', '--nbatch', metavar='N', type=int, help='Number of frames in each \'batch\' \
             (note: the truncation value can only change once per batch. Don\'t fuck with this unless you know \
             what it does.).', default=1)
     parser.add_argument('-o', '--output-dir', help='Directory path for output images.')
     parser.add_argument('--prefix', help='File prefix for output images.')
-    args = parser.parse_args()
+    parser.add_argument('--interp', help='Set interpolation method.', choices=['linear', 'cubic'])
+    group_loop = parser.add_mutually_exclusive_group(required=False)
+    group_loop.add_argument('--loop', dest='loop', action='store_true', default=True, help='Loop the animation.')
+    group_loop.add_argument('--no-loop', dest='loop', action='store_false', help='Don\'t loop the animation.')
+    args = parser.parse_args(argv)
     # validate args
-    if args.keys and not (args.username and args.password):
+    if not (lambda l: (not any(l)) or all(l))(\
+            [e is not None and e is not [] for e in [args.username, args.password, args.keys]]):
         parser.error('The --keys argument requires a --username and --password to login to ganbreeder')
         sys.exit(1)
+    return args
 
+# create entrypoints for cli tools
+def main():
+    handle_args()
     # get animation keyframes from ganbreeder
     print('Downloading keyframe info from ganbreeder...')
     keyframes = ganbreeder.get_info_batch(args.username, args.password, args.keys)