Improvements to cmd line.

Author: coderforlife

README.md

@@ -30,10 +30,10 @@ out_gpu = hist.histeq_exact(im_gpu)
 
 The code can be run as a standalone program:
 ```sh
-python3 -m hist.main input.png output.png
-python3 -m hist.main input output # converts a 3d image stored in a folder
+python3 -m hist input.png output.png
+python3 -m hist input # converts a 3d image stored in a folder
 ```
-See `python3 -m hist.main --help` for more information.
+See `python3 -m hist --help` for more information.
 
 References
 ----------

hist/__init__.py

@@ -7,3 +7,4 @@
 from .exact import histeq_exact
 from .metrics import (contrast_per_pixel, distortion, count_differences, psnr, ssim,
                       enhancement_measurement)
+

hist/__main__.py

@@ -0,0 +1,78 @@
+"""
+Simple main program to perform histogram equalization.
+"""
+
+from . import histeq, histeq_exact
+
+def main():
+    """Main function that runs histogram equalization on an image."""
+    import argparse
+    import os.path
+    import re
+    import hist._cmd_line_util as cui
+
+    # Extra imports to make sure everything is available now
+    import numpy, scipy.ndimage # pylint: disable=unused-import, multiple-imports
+
+    parser = argparse.ArgumentParser(prog='python3 -m hist',
+                                     description='Perform histogram equalization on an image')
+    cui.add_input_image(parser)
+    parser.add_argument('output', nargs='?',
+                        help='output image file, defaults to input file name with _out before '
+                        'the extension')
+    cui.add_method_arg(parser)
+    cui.add_kwargs_arg(parser)
+    parser.add_argument('--nbins', '-n', type=int, default=256, metavar='N',
+                        help='number of bins in the intermediate histogram, default is 256')
+    args = parser.parse_args()
+
+    # Load image
+    im = cui.open_input_image(args)
+
+    # Run HE
+    if args.method == 'classic':
+        out = histeq(im, args.nbins, **dict(args.kwargs))
+    else:
+        out = histeq_exact(im, args.nbins, method=args.method, **dict(args.kwargs))
+
+    # Save (if not testing)
+    filename = args.output
+    if filename == '': return # hidden feature for "testing" mode, no saving
+    if filename is None:
+        if os.path.isdir(args.input):
+            filename = args.input + '_out'
+        elif os.path.exists(args.input) and ('?' in args.input or '*' in args.input or
+                                             ('[' in args.input and ']' in args.input)):
+            filename = re.sub(r'\*|\?|\[.+\]', '#', args.input)
+        elif args.input.lower().endswith('.npy.gz'):
+            filename = args.input[:-7] + '_out' + args.input[-7:]
+        else:
+            filename = '_out'.join(os.path.splitext(args.input))
+    __save(filename, out)
+
+def __save(filename, out):
+    import gzip
+    import numpy
+    import imageio
+    from .util import is_on_gpu
+    if is_on_gpu(out):
+        out = out.get()
+    if filename.lower().endswith('.npy'):
+        numpy.save(filename, out)
+    elif filename.lower().endswith('.npy.gz'):
+        with gzip.GzipFile(filename, 'wb') as file:
+            numpy.save(filename, file)
+    elif out.ndim == 3 and '#' in filename: 
+        start = filename.index('#')
+        end = start + 1
+        while end < len(filename) and filename[end] == '#':
+            end += 1
+        num_str, fmt_str = filename[start:end], '%0'+str(end-start)+'d'
+        for i in range(out.shape[0]):
+            imageio.imwrite(filename.replace(num_str, fmt_str % i), out[i, :, :])
+    else:
+        imageio.imwrite(filename, out)
+
+
+if __name__ == "__main__":
+    main()

hist/_cmd_line_util.py

@@ -8,25 +8,27 @@
 
 def add_method_arg(parser):
     """Add the method argument to an argument parser object."""
-    parser.add_argument('method', choices=METHODS, help='method of histogram equalization')
+    parser.add_argument('--method', '-m', choices=METHODS, default='va',
+                        help='method of histogram equalization, default is "va"')
 
 def add_kwargs_arg(parser):
-    """Add the kwargs arg to an argument parser object which accepts a series of k=v arguments."""
-    parser.add_argument('kwargs', type=__kwargs_arg, nargs='*', help='any special keyword '+
-                        'arguments to pass to the method, formated as key=value with value being '+
-                        'a valid Python literal or one of the special values nan, inf, -inf, N4, '+
-                        'N8, N8_DIST, N6, N18, N18_DIST, N26, N26_DIST')
+    """Add the kwargs arg to an argument parser object which accepts many k=v arguments."""
+    parser.add_argument('--arg', '-a', type=__kwargs_arg, metavar='K=V', action='append',
+                        dest='kwargs', default=[],
+                        help='any special keyword arguments to pass to the method, formated as '
+                        'key=value with value being a valid Python literal or one of the special '
+                        'values nan, inf, -inf, N4, N8, N8_DIST, N6, N18, N18_DIST, N26, N26_DIST')
 
-def add_input_image(parser):
+def add_input_image(parser, output=False):
     """
     Add a required input argument and an optional --float argument. Use the open_input_image
-    function to read the image. This supports filenames with glob wildcards or directories to read a
-    series of images in as a 3D image.
+    function to read the image. This supports filenames with glob wildcards or directories to read
+    a series of images in as a 3D image.
     """
     parser.add_argument('input', help='input image file (including .npy, .npy.gz, and '
                         'directories/wildcard names for 3D images)')
-    parser.add_argument('--float', action='store_true', help='convert image to float')
-    parser.add_argument('--gpu', action='store_true', help='utilize the GPU when able')
+    parser.add_argument('--float', '-f', action='store_true', help='convert image to float')
+    parser.add_argument('--gpu', '-g', action='store_true', help='utilize the GPU when able')
 
 def open_input_image(args_or_filename, conv_to_float=False, use_gpu=False):
     """
@@ -39,6 +41,7 @@ def open_input_image(args_or_filename, conv_to_float=False, use_gpu=False):
     import os
     from glob import glob
     from numpy import stack
+
     if isinstance(args_or_filename, str):
         filename = args_or_filename
     else:
@@ -66,10 +69,10 @@ def __load_image(filename, conv_to_float=False, use_gpu=False):
     import imageio
     from numpy import load
     from hist.util import as_float
-    if filename.endswith('.npy.gz'):
+    if filename.lower().endswith('.npy.gz'):
         with gzip.GzipFile(filename, 'rb') as file:
             im = load(file)
-    elif filename.endswith('.npy'):
+    elif filename.lower().endswith('.npy'):
         im = load(filename)
     else:
         im = imageio.imread(filename)

hist/main.py

@@ -6,7 +6,7 @@
 from .metrics import (contrast_per_pixel, enhancement_measurement, distortion,
                       contrast_enhancement, count_differences, psnr, ssim)
 
-def metric_battery(original, method, csv=False, plot=False, **kwargs):
+def metric_battery(original, method, nbins=256, csv=False, plot=False, **kwargs):
     """
     Runs a battery of metrics while historgram equalizing the original image and then attempting to
     reconstruct the original image using the given method (either 'classic' or one of the methods
@@ -20,14 +20,14 @@ def metric_battery(original, method, csv=False, plot=False, **kwargs):
     # pylint: disable=too-many-locals, too-many-statements
     hist_orig = imhist(original)
     if method == 'classic':
-        enhanced = histeq(original, 256, **kwargs)
+        enhanced = histeq(original, nbins, **kwargs)
         recon = histeq(enhanced, hist_orig, **kwargs)
         fails_forward = fails_reverse = -1
     else:
         kwargs['method'] = method
         kwargs['return_fails'] = True
         if 'reconstruction' in kwargs: kwargs['reconstruction'] = False
-        enhanced, fails_forward = histeq_exact(original, 256, **kwargs)
+        enhanced, fails_forward = histeq_exact(original, nbins, **kwargs)
         if 'reconstruction' in kwargs: kwargs['reconstruction'] = True
         recon, fails_reverse = histeq_exact(enhanced, hist_orig, **kwargs)
 
@@ -110,6 +110,9 @@ def main():
     parser.add_argument('--csv', action='store_true', help='output data as CSV with no header')
     parser.add_argument('--plot', action='store_true',
                         help='plot original, enhanced, and reconstructed images with histograms')
+    parser.add_argument('--nbins', '-n', type=int, default=256, metavar='N',
+                        help='number of bins in the intermediate histogram, default is 256 '
+                        '(reverse direction always uses a full histogram)')
     cui.add_kwargs_arg(parser)
     args = parser.parse_args()