Skip to content

Commit

Permalink
Improvements to cmd line.
Browse files Browse the repository at this point in the history
  • Loading branch information
coderforlife committed Jan 31, 2021
1 parent 60aa77e commit a323334
Show file tree
Hide file tree
Showing 6 changed files with 104 additions and 79 deletions.
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand Down
1 change: 1 addition & 0 deletions hist/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,3 +7,4 @@
from .exact import histeq_exact
from .metrics import (contrast_per_pixel, distortion, count_differences, psnr, ssim,
enhancement_measurement)

78 changes: 78 additions & 0 deletions hist/__main__.py
Original file line number Diff line number Diff line change
@@ -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()
29 changes: 16 additions & 13 deletions hist/_cmd_line_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
60 changes: 0 additions & 60 deletions hist/main.py

This file was deleted.

9 changes: 6 additions & 3 deletions hist/metric_battery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -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()

Expand Down

0 comments on commit a323334

Please sign in to comment.