diff --git a/README.md b/README.md index 8e20c00..2be25d4 100644 --- a/README.md +++ b/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 ---------- diff --git a/hist/__init__.py b/hist/__init__.py index f9bb13a..8387390 100644 --- a/hist/__init__.py +++ b/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) + diff --git a/hist/__main__.py b/hist/__main__.py new file mode 100644 index 0000000..7cc29b4 --- /dev/null +++ b/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() diff --git a/hist/_cmd_line_util.py b/hist/_cmd_line_util.py index 8e61f02..4c11aba 100644 --- a/hist/_cmd_line_util.py +++ b/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) diff --git a/hist/main.py b/hist/main.py deleted file mode 100644 index c830395..0000000 --- a/hist/main.py +++ /dev/null @@ -1,60 +0,0 @@ -""" -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 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(description='Perform histogram equalization on an image') - cui.add_input_image(parser) - parser.add_argument('output', help='output image file') - cui.add_method_arg(parser) - cui.add_kwargs_arg(parser) - args = parser.parse_args() - - # Load image - im = cui.open_input_image(args) - - # Run HE - if args.method == 'classic': - out = histeq(im, 256, **dict(args.kwargs)) - else: - out = histeq_exact(im, method=args.method, **dict(args.kwargs)) - - # Save (if not testing) - if args.output != '': - __save(args, out) - -def __save(args, out): - import gzip - import numpy - import imageio - from .util import is_on_gpu - if is_on_gpu(out): - out = out.get() - if args.output.endswith('.npy'): - numpy.save(args.output, out) - elif args.output.endswith('.npy.gz'): - with gzip.GzipFile(args.output, 'wb') as file: - numpy.save(args.output, file) - elif out.ndim == 3 and '#' in args.output: - start = args.output.index('#') - end = start + 1 - while end < len(args.output) and args.output[end] == '#': - end += 1 - num_str, fmt_str = args.output[start:end], '%0'+str(end-start)+'d' - for i in range(out.shape[0]): - imageio.imwrite(args.output.replace(num_str, fmt_str % i), out[i, :, :]) - else: - imageio.imwrite(args.output, out) - - -if __name__ == "__main__": - main() diff --git a/hist/metric_battery.py b/hist/metric_battery.py index 4852395..cf5a25c 100644 --- a/hist/metric_battery.py +++ b/hist/metric_battery.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()