# A script to improve the contrast of an image using histogram equalization
# Uses the OpenCV library which can be installed using pip install opencv-python

import cv2, argparse
import convertColorSpace

# SPECIFY COMMAND LINE ARGUMENTS
parser = argparse.ArgumentParser(description="image histogram equalization")
parser.add_argument('-i', '--input', help='Enter the input image', required=True, dest="input")
parser.add_argument('-o', '--output', help='Enter the name and location of the output image', required=True, dest="output")
parser.add_argument('-s', '--cspace', help='Enter the color space for the conversion to be applied in (YUV, HSV, LAB, or YCrCb)', required=True, dest="cspace")
args = parser.parse_args()

# Load image
image = cv2.imread(args.input)

# Convert image to specified color space
image = convertColorSpace.convertFromBGR(image, args.cspace)

# Perform histogram equalization on color channel that affects contrast
# Channel is index 2 for HSV, 0 for all others
if args.cspace=='HSV':
	channel = 2
else:
	channel = 0

image[:, :, channel] = cv2.equalizeHist(image[:, :, channel])

# Convert image back to BGR color space
image = convertColorSpace.convertToBGR(image, args.cspace)

# Write image to specified output
cv2.imwrite(args.output,image)