#!/usr/bin/python3
import matplotlib.pyplot as plt
import numpy as np
import cv2, sys

if (len(sys.argv) < 2):
	print("Usage: python3 heatmap.py image.png [colormap]")
	print("If [colormap] is unspecified, the default 'inferno' will be used")
	print("All colormaps from OpenCV2 can be used")
	exit(1)

image = cv2.imread(sys.argv[1], 0)
image = cv2.equalizeHist(image)
if len(sys.argv) < 3:
	map_cmap = "inferno"
else:
	map_cmap = sys.argv[2]
colormap = plt.get_cmap(map_cmap)
heatmap = (colormap(image) * 2**16).astype(np.uint16)[:, :, :3]
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_RGB2BGR)
cv2.imwrite('.'.join(sys.argv[1].split('.')[:-1]) + '_heatmap.png', heatmap)
cv2.waitKey()