Replace click with argparse

README.md

@@ -20,7 +20,6 @@ requested output format and saved to the specified location.
   * MXNet
   * Pillow
   * Numpy
-  * Click
 
 ## Installation
 The application is available on `nvidia-pyindex` and can be downloaded and

imagine/imagine.py

@@ -15,7 +15,7 @@
 import os
 import re
 import numpy
-import click
+from argparse import ArgumentParser
 from PIL import Image
 from multiprocessing.pool import Pool
 try:
@@ -35,21 +35,61 @@
     TFRecordWriter = None
 
 
+STANDARD_IMAGE = 'create-images'
+TFRECORD = 'create-tfrecord'
+RECORDIO = 'create-recordio'
 SUPPORTED_IMAGE_FORMATS = {"jpg": "jpg", "jpeg": "jpg", "bmp": "bmp",
                            "bitmap": "bmp", "png": "png"}
 
 
-@click.group()
-def main():
-    """
+def parse_args():
+    message = """
     CLI for generating a fake dataset of various quantities at different
     resolutions.
 
     Supported file types: .bmp, .png, and .jpg.
     Supported record types: TFRecords, and RecordIO.
     TFRecords requires an external index file creation step.
     """
-    pass
+    parser = ArgumentParser(message)
+    # Required positional command subparser which should be specified first
+    commands = parser.add_subparsers(dest='command', metavar='command')
+    commands_parent = ArgumentParser(add_help=False)
+
+    # Options specific to record types
+    commands_parent.add_argument('source_path', metavar='source-path',
+                                 help='Path containing valid input images to '
+                                 'convert to records')
+    commands_parent.add_argument('dest_path', metavar='dest-path',
+                                 help='Path to save record files to')
+    commands_parent.add_argument('name', help='Name to prepend files with, '
+                                 'such as "sample_record_"')
+    commands_parent.add_argument('--img-per-file', type=int, default=1000)
+    commands.add_parser(TFRECORD, help='Create TFRecords from input images',
+                        parents=[commands_parent])
+    commands.add_parser(RECORDIO, help='Create RecordIO from input images',
+                        parents=[commands_parent])
+
+    # Options specific to generating standard images
+    standard = commands.add_parser(STANDARD_IMAGE, help='Generate random '
+                                   'images')
+    standard.add_argument('path', help='Path to save images to')
+    standard.add_argument('name', help='Name to prepend files with, such as '
+                          '"sample_image_"')
+    standard.add_argument('image_format', metavar='image-format', help='The '
+                          'image format to generate',
+                          choices=SUPPORTED_IMAGE_FORMATS.keys())
+    standard.add_argument('--width', help='The image width in pixels',
+                          type=int, default=1920)
+    standard.add_argument('--height', help='The image height in pixels',
+                          type=int, default=1080)
+    standard.add_argument('--count', help='The number of images to generate',
+                          type=int, default=1)
+    standard.add_argument('--seed', help='The seed to use while generating '
+                          'random image data', type=int, default=0)
+    standard.add_argument('--size', help='Display the first image size and '
+                          'the directory size for the images')
+    return parser.parse_args()
 
 
 def try_create_directory(directory):
@@ -63,21 +103,11 @@ def check_directory_exists(directory):
                            'contains valid images.')
 
 
-@main.command()
-@click.option('--path', required=True)
-@click.option('--name', required=True)
-@click.option('--width', default=1920, required=True)
-@click.option('--height', default=1080, required=True)
-@click.option('--count', default=1, required=True)
-@click.option('--image_format', default='png', required=True)
-@click.option('--seed', default=0)
-@click.option('--size', is_flag=True, default=False)
-@click.option('--chunksize', default=64)
 def create_images(path, name, width, height, count, image_format, seed, size,
-                  chunksize):
-    click.echo("Creating {} {} files located at {} of {}x{} resolution with a "
-               "base filename of {}".format(count, image_format, path, width,
-                                            height, name))
+                  chunksize=64):
+    print('Creating {} {} files located at {} of {}x{} resolution with a base '
+          'base filename of {}'.format(count, image_format, path, width,
+                                       height, name))
     try_create_directory(path)
     combined_path = os.path.join(path, name)
 
@@ -102,8 +132,7 @@ def create_images(path, name, width, height, count, image_format, seed, size,
     if size:
         print_image_information(path)
 
-    click.echo("Created {} files in {} seconds".format(count,
-                                                       stop_time-start_time))
+    print('Created {} files in {} seconds'.format(count, stop_time-start_time))
 
 
 def record_slice(source_path, dest_path, name, image_files, images_per_file,
@@ -117,17 +146,12 @@ def record_slice(source_path, dest_path, name, image_files, images_per_file,
                num)
 
 
-@main.command()
-@click.option('--source_path', required=True)
-@click.option('--dest_path', required=True)
-@click.option('--name', required=True)
-@click.option('--img_per_file', default=1000)
 def create_recordio(source_path, dest_path, name, img_per_file):
-    click.echo("Creating RecordIO files at {} from {} targeting {} files per "
-               "record with a base filename of {}".format(dest_path,
-                                                          source_path,
-                                                          img_per_file,
-                                                          name))
+    print('Creating RecordIO files at {} from {} targeting {} files per '
+          'record with a base filename of {}'.format(dest_path,
+                                                     source_path,
+                                                     img_per_file,
+                                                     name))
     if not IRHeader:
         raise ImportError('MXNet not found! Please install MXNet dependency '
                           'using "pip install nvidia-imageinary[\'mxnet\']".')
@@ -159,20 +183,15 @@ def create_recordio(source_path, dest_path, name, img_per_file):
         pool.join()
 
     stop_time = perf_counter()
-    click.echo("Completed in {} seconds".format(stop_time-start_time))
+    print('Completed in {} seconds'.format(stop_time-start_time))
 
 
-@main.command()
-@click.option('--source_path', required=True)
-@click.option('--dest_path', required=True)
-@click.option('--name', required=True)
-@click.option('--img_per_file', default=1000)
 def create_tfrecords(source_path, dest_path, name, img_per_file):
-    click.echo("Creating TFRecord files at {} from {} targeting {} files per "
-               "TFRecord with a base filename of {}".format(dest_path,
-                                                            source_path,
-                                                            img_per_file,
-                                                            name))
+    print('Creating TFRecord files at {} from {} targeting {} files per '
+          'TFRecord with a base filename of {}'.format(dest_path,
+                                                       source_path,
+                                                       img_per_file,
+                                                       name))
     if not TFRecordWriter:
         raise ImportError('TensorFlow not found! Please install TensorFlow '
                           'dependency using "pip install '
@@ -212,7 +231,7 @@ def create_tfrecords(source_path, dest_path, name, img_per_file):
     writer.close()
     stop_time = perf_counter()
 
-    click.echo("Completed in {} seconds".format(stop_time-start_time))
+    print('Completed in {} seconds'.format(stop_time-start_time))
 
 
 def print_image_information(path):
@@ -227,9 +246,9 @@ def print_image_information(path):
         if is_first_image:
             first_image_size = directory_size
             is_first_image = False
-    click.echo("First image size from {}, in bytes: {}".format(path,
-               first_image_size))
-    click.echo("Directory {} size, in bytes: {}".format(path, directory_size))
+    print('First image size from {}, in bytes: {}'.format(path,
+                                                          first_image_size))
+    print('Directory {} size, in bytes: {}'.format(path, directory_size))
 
 
 def recordio_creation(source_path, dest_path, name, image_files, n):
@@ -264,5 +283,18 @@ def image_creation(combined_path, width, height, seed, image_format, n):
     im_out.save('%s%d.%s' % (combined_path, n, file_ext))
 
 
+def main():
+    args = parse_args()
+    if args.command == STANDARD_IMAGE:
+        create_images(args.path, args.name, args.width, args.height,
+                      args.count, args.image_format, args.seed, args.size)
+    elif args.command == TFRECORD:
+        create_tfrecords(args.source_path, args.dest_path, args.name,
+                         args.img_per_file)
+    elif args.command == RECORDIO:
+        create_recordio(args.source_path, args.dest_path, args.name,
+                        args.img_per_file)
+
+
 if __name__ == "__main__":
     main()

requirements.txt

@@ -4,7 +4,6 @@ attrs==20.2.0
 cachetools==4.1.0
 certifi==2020.4.5.1
 chardet==3.0.4
-click==7.1.2
 coverage==5.3
 gast==0.3.3
 google-auth==1.14.3

requirements.txt.tf1

@@ -3,7 +3,6 @@ astor==0.8.1
 attrs==20.2.0
 certifi==2020.4.5.2
 chardet==3.0.4
-click==7.1.2
 coverage==5.3
 gast==0.3.3
 google-pasta==0.2.0

setup.py

@@ -27,7 +27,6 @@
         'console_scripts': ['imagine=imagine.imagine:main']
     },
     install_requires=[
-        'click >= 7.1.2',
         'numpy >= 1.18.0',
         'Pillow >= 7.1.2'
     ],

tests/functional/test_jpg_creation.py

@@ -14,31 +14,32 @@
 import pytest
 import re
 import os
-from click.testing import CliRunner
 from glob import glob
-from imagine.imagine import main
+from imagine.imagine import create_images
 from PIL import Image
 
 
 class TestJPGCreation:
     @pytest.fixture(autouse=True)
     def setup(self, tmpdir):
         self.tmpdir = tmpdir.mkdir('jpg_files')
-        self.runner = CliRunner()
 
     def teardown_method(self):
         for image in glob(f'{str(self.tmpdir)}/*'):
             os.remove(image)
         os.rmdir(str(self.tmpdir))
 
     def test_creating_one_hundred_images(self):
-        self.runner.invoke(main, ['create-images',
-                                  '--path', str(self.tmpdir),
-                                  '--name', 'tmp_',
-                                  '--width', 1920,
-                                  '--height', 1080,
-                                  '--count', 100,
-                                  '--image_format', 'jpg'])
+        create_images(
+            str(self.tmpdir),
+            'tmp_',
+            1920,
+            1080,
+            100,
+            'jpg',
+            0,
+            False
+        )
 
         images = glob(f'{str(self.tmpdir)}/*')
 
@@ -49,13 +50,16 @@ def test_creating_one_hundred_images(self):
                 assert im.size == (1920, 1080)
 
     def test_creating_one_hundred_4K_images(self):
-        self.runner.invoke(main, ['create-images',
-                                  '--path', str(self.tmpdir),
-                                  '--name', 'tmp_',
-                                  '--width', 3840,
-                                  '--height', 2160,
-                                  '--count', 100,
-                                  '--image_format', 'jpg'])
+        create_images(
+            str(self.tmpdir),
+            'tmp_',
+            3840,
+            2160,
+            100,
+            'jpg',
+            0,
+            False
+        )
 
         images = glob(f'{str(self.tmpdir)}/*')
 

tests/functional/test_png_creation.py

@@ -14,31 +14,32 @@
 import pytest
 import re
 import os
-from click.testing import CliRunner
 from glob import glob
-from imagine.imagine import main
+from imagine.imagine import create_images
 from PIL import Image
 
 
 class TestPNGCreation:
     @pytest.fixture(autouse=True)
     def setup(self, tmpdir):
         self.tmpdir = tmpdir.mkdir('png_files')
-        self.runner = CliRunner()
 
     def teardown_method(self):
         for image in glob(f'{str(self.tmpdir)}/*'):
             os.remove(image)
         os.rmdir(str(self.tmpdir))
 
     def test_creating_one_hundred_images(self):
-        self.runner.invoke(main, ['create-images',
-                                  '--path', str(self.tmpdir),
-                                  '--name', 'tmp_',
-                                  '--width', 1920,
-                                  '--height', 1080,
-                                  '--count', 100,
-                                  '--image_format', 'png'])
+        create_images(
+            str(self.tmpdir),
+            'tmp_',
+            1920,
+            1080,
+            100,
+            'png',
+            0,
+            False
+        )
 
         images = glob(f'{str(self.tmpdir)}/*')
 
@@ -49,13 +50,16 @@ def test_creating_one_hundred_images(self):
                 assert im.size == (1920, 1080)
 
     def test_creating_one_hundred_4K_images(self):
-        self.runner.invoke(main, ['create-images',
-                                  '--path', str(self.tmpdir),
-                                  '--name', 'tmp_',
-                                  '--width', 3840,
-                                  '--height', 2160,
-                                  '--count', 100,
-                                  '--image_format', 'png'])
+        create_images(
+            str(self.tmpdir),
+            'tmp_',
+            3840,
+            2160,
+            100,
+            'png',
+            0,
+            False
+        )
 
         images = glob(f'{str(self.tmpdir)}/*')