import keras 
import numpy as np
import pandas as pd
import cv2
import os
import json
import pdb
import argparse
import math
import copy
from vis.visualization import visualize_cam, overlay, visualize_activation
from vis.utils.utils import apply_modifications
from shutil import rmtree
import matplotlib.cm as cm
from matplotlib import pyplot as plt

from sklearn import metrics

import keras.backend as K
from keras import activations
from keras.applications.inception_v3 import preprocess_input as inception_pre
from keras.applications.mobilenet import preprocess_input as mobilenet_pre
from keras.applications.resnet50 import preprocess_input as resnet_pre
from keras.applications.densenet import preprocess_input as densenet_pre
from datagenerator import ImageDataGenerator

from utils import load_model

def getCAM(model, image):
	# weights of the final fully-connected layer
	weights = model.layers[-1].get_weights()[0]

	# activation before the last global pooling
	for layer in reversed(model.layers):
		if len(layer.output_shape) > 2:
			break
	function = K.function([model.layers[0].input, K.learning_phase()], [layer.output])
	activation = np.squeeze(function([image, 0])[0])

	# weighted sum of the activation map
	CAM = np.dot(activation, weights)

	return CAM


def main():
	ap = argparse.ArgumentParser()
	ap.add_argument('--ckpt_path', help = 'Path to the model checkpoint.')
	ap.add_argument('--image_path', help = 'Path to the image to run inference on.')
	ap.add_argument('--bnbox', help = 'Path to the bounding box annotation, if applies.')
	ap.add_argument('--threshold', default = 0.5, help = 'Threshold for displaying the Class Activation Map.')

	args = ap.parse_args()

	model_dir = os.path.dirname(args.ckpt_path)

	with open(os.path.join(model_dir, 'label_map.json'), 'r') as f:
		label_map = json.load(f)

	num_class = len(list(label_map.keys()))

	model, model_config = load_model(model_dir, args.ckpt_path)

	model_name = model_config['model_name']

	if model_name in ['inception']:
		image_size = 299

	else:
		image_size = 224

	preprocess_input = {
		'inception': inception_pre,
		'resnet': resnet_pre,
		'mobilenet': mobilenet_pre,
		'densenet': densenet_pre
	}

	if args.bnbox is not None:
		annotation = pd.read_csv(args.bnbox)

	image_index = os.path.basename(args.image_path)

	indices = np.where(annotation['Image Index'] == image_index)[0]

	bnbox = {}

	for i in indices:
		disease = annotation['Finding Label'][i]
		x = int(annotation['Bbox [x'][i] + 0.5)
		y = int(annotation['y'][i] + 0.5)
		w = int(annotation['w'][i] + 0.5)
		h = int(annotation['h]'][i] + 0.5)
		bnbox[disease] = [x, y, x + w, y + h]

	image = cv2.imread(args.image_path)

	img = cv2.resize(image, (image_size, image_size))

	img = preprocess_input[model_name](img.astype(np.float32))

	img = np.expand_dims(img, axis = 0)

	predictions = np.squeeze(model.predict(img))

	CAM = getCAM(model, img)

	cv2.namedWindow("ChestXray", cv2.WINDOW_NORMAL)

	for key, value in label_map.items():

		heatmap = CAM[:,:,int(key)]
		heatmap -= heatmap.min()
		heatmap *= 255.0 / heatmap.max()
		heatmap[np.where(heatmap < args.threshold * 255)] *= 0.1

		heatmap = cv2.applyColorMap(heatmap.astype(np.uint8), cv2.COLORMAP_JET)

		heatmap = cv2.resize(heatmap, image.shape[:2], cv2.INTER_AREA)

		overlay_img = overlay(heatmap, image, alpha = 0.4)

		cv2.putText(overlay_img, "{}: {:.2%}".format(value, predictions[int(key)]),
					(30,30), cv2.FONT_HERSHEY_DUPLEX, 1.0, (255,255,255), 2)

		if value in bnbox.keys():
			box = bnbox[value]
			cv2.rectangle(overlay_img, (box[0], box[1]), (box[2], box[3]),
							color = (0, 180, 0), thickness = 2)

		cv2.imshow("ChestXray", overlay_img)
		cv2.waitKey()

		plt.show()

		print('{}: {:.2%}'.format(value, predictions[int(key)]))

	cv2.destroyAllWindows()


if __name__ == "__main__":
	main()
[((41, 7, 41, 34), 'numpy.dot', 'np.dot', ({(41, 14, 41, 24): 'activation', (41, 26, 41, 33): 'weights'}, {}), '(activation, weights)', True, 'import numpy as np\n'), ((47, 6, 47, 31), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((55, 13, 55, 44), 'os.path.dirname', 'os.path.dirname', ({(55, 29, 55, 43): 'args.ckpt_path'}, {}), '(args.ckpt_path)', False, 'import os\n'), ((62, 23, 62, 60), 'utils.load_model', 'load_model', ({(62, 34, 62, 43): 'model_dir', (62, 45, 62, 59): 'args.ckpt_path'}, {}), '(model_dir, args.ckpt_path)', False, 'from utils import load_model\n'), ((82, 15, 82, 48), 'os.path.basename', 'os.path.basename', ({(82, 32, 82, 47): 'args.image_path'}, {}), '(args.image_path)', False, 'import os\n'), ((96, 9, 96, 36), 'cv2.imread', 'cv2.imread', ({(96, 20, 96, 35): 'args.image_path'}, {}), '(args.image_path)', False, 'import cv2\n'), ((98, 7, 98, 50), 'cv2.resize', 'cv2.resize', ({(98, 18, 98, 23): 'image', (98, 25, 98, 49): '(image_size, image_size)'}, {}), '(image, (image_size, image_size))', False, 'import cv2\n'), ((102, 7, 102, 36), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((108, 1, 108, 48), 'cv2.namedWindow', 'cv2.namedWindow', ({(108, 17, 108, 28): '"""ChestXray"""', (108, 30, 108, 47): 'cv2.WINDOW_NORMAL'}, {}), "('ChestXray', cv2.WINDOW_NORMAL)", False, 'import cv2\n'), ((138, 1, 138, 24), 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ({}, {}), '()', False, 'import cv2\n'), ((58, 14, 58, 26), 'json.load', 'json.load', ({(58, 24, 58, 25): 'f'}, {}), '(f)', False, 'import json\n'), ((80, 15, 80, 38), 'pandas.read_csv', 'pd.read_csv', ({(80, 27, 80, 37): 'args.bnbox'}, {}), '(args.bnbox)', True, 'import pandas as pd\n'), ((84, 11, 84, 61), 'numpy.where', 'np.where', ({(84, 20, 84, 60): "(annotation['Image Index'] == image_index)"}, {}), "(annotation['Image Index'] == image_index)", True, 'import numpy as np\n'), ((119, 12, 119, 64), 'cv2.resize', 'cv2.resize', ({(119, 23, 119, 30): 'heatmap', (119, 32, 119, 47): 'image.shape[:2]', (119, 49, 119, 63): 'cv2.INTER_AREA'}, {}), '(heatmap, image.shape[:2], cv2.INTER_AREA)', False, 'import cv2\n'), ((121, 16, 121, 52), 'vis.visualization.overlay', 'overlay', (), '', False, 'from vis.visualization import visualize_cam, overlay, visualize_activation\n'), ((131, 2, 131, 38), 'cv2.imshow', 'cv2.imshow', ({(131, 13, 131, 24): '"""ChestXray"""', (131, 26, 131, 37): 'overlay_img'}, {}), "('ChestXray', overlay_img)", False, 'import cv2\n'), ((132, 2, 132, 15), 'cv2.waitKey', 'cv2.waitKey', ({}, {}), '()', False, 'import cv2\n'), ((134, 2, 134, 12), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'from matplotlib import pyplot as plt\n'), ((37, 47, 37, 65), 'keras.backend.learning_phase', 'K.learning_phase', ({}, {}), '()', True, 'import keras.backend as K\n'), ((57, 11, 57, 52), 'os.path.join', 'os.path.join', ({(57, 24, 57, 33): 'model_dir', (57, 35, 57, 51): '"""label_map.json"""'}, {}), "(model_dir, 'label_map.json')", False, 'import os\n'), ((128, 3, 129, 42), 'cv2.rectangle', 'cv2.rectangle', (), '', False, 'import cv2\n'), ((115, 10, 115, 50), 'numpy.where', 'np.where', ({(115, 19, 115, 49): '(heatmap < args.threshold * 255)'}, {}), '(heatmap < args.threshold * 255)', True, 'import numpy as np\n')]