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()
[((1233, 1260), 'numpy.dot', 'np.dot', (['activation', 'weights'], {}), '(activation, weights)\n', (1239, 1260), True, 'import numpy as np\n'), ((1294, 1319), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (1317, 1319), False, 'import argparse\n'), ((1708, 1739), 'os.path.dirname', 'os.path.dirname', (['args.ckpt_path'], {}), '(args.ckpt_path)\n', (1723, 1739), False, 'import os\n'), ((1899, 1936), 'utils.load_model', 'load_model', (['model_dir', 'args.ckpt_path'], {}), '(model_dir, args.ckpt_path)\n', (1909, 1936), False, 'from utils import load_model\n'), ((2280, 2313), 'os.path.basename', 'os.path.basename', (['args.image_path'], {}), '(args.image_path)\n', (2296, 2313), False, 'import os\n'), ((2657, 2684), 'cv2.imread', 'cv2.imread', (['args.image_path'], {}), '(args.image_path)\n', (2667, 2684), False, 'import cv2\n'), ((2693, 2736), 'cv2.resize', 'cv2.resize', (['image', '(image_size, image_size)'], {}), '(image, (image_size, image_size))\n', (2703, 2736), False, 'import cv2\n'), ((2806, 2833), 'numpy.expand_dims', 'np.expand_dims', (['img'], {'axis': '(0)'}), '(img, axis=0)\n', (2820, 2833), True, 'import numpy as np\n'), ((2912, 2959), 'cv2.namedWindow', 'cv2.namedWindow', (['"""ChestXray"""', 'cv2.WINDOW_NORMAL'], {}), "('ChestXray', cv2.WINDOW_NORMAL)\n", (2927, 2959), False, 'import cv2\n'), ((3779, 3802), 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ([], {}), '()\n', (3800, 3802), False, 'import cv2\n'), ((1820, 1832), 'json.load', 'json.load', (['f'], {}), '(f)\n', (1829, 1832), False, 'import json\n'), ((2240, 2263), 'pandas.read_csv', 'pd.read_csv', (['args.bnbox'], {}), '(args.bnbox)\n', (2251, 2263), True, 'import pandas as pd\n'), ((2326, 2376), 'numpy.where', 'np.where', (["(annotation['Image Index'] == image_index)"], {}), "(annotation['Image Index'] == image_index)\n", (2334, 2376), True, 'import numpy as np\n'), ((3239, 3291), 'cv2.resize', 'cv2.resize', (['heatmap', 'image.shape[:2]', 'cv2.INTER_AREA'], {}), '(heatmap, image.shape[:2], cv2.INTER_AREA)\n', (3249, 3291), False, 'import cv2\n'), ((3309, 3343), 'vis.visualization.overlay', 'overlay', (['heatmap', 'image'], {'alpha': '(0.4)'}), '(heatmap, image, alpha=0.4)\n', (3316, 3343), False, 'from vis.visualization import visualize_cam, overlay, visualize_activation\n'), ((3650, 3686), 'cv2.imshow', 'cv2.imshow', (['"""ChestXray"""', 'overlay_img'], {}), "('ChestXray', overlay_img)\n", (3660, 3686), False, 'import cv2\n'), ((3689, 3702), 'cv2.waitKey', 'cv2.waitKey', ([], {}), '()\n', (3700, 3702), False, 'import cv2\n'), ((3706, 3716), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (3714, 3716), True, 'from matplotlib import pyplot as plt\n'), ((1100, 1118), 'keras.backend.learning_phase', 'K.learning_phase', ([], {}), '()\n', (1116, 1118), True, 'import keras.backend as K\n'), ((1752, 1793), 'os.path.join', 'os.path.join', (['model_dir', '"""label_map.json"""'], {}), "(model_dir, 'label_map.json')\n", (1764, 1793), False, 'import os\n'), ((3541, 3640), 'cv2.rectangle', 'cv2.rectangle', (['overlay_img', '(box[0], box[1])', '(box[2], box[3])'], {'color': '(0, 180, 0)', 'thickness': '(2)'}), '(overlay_img, (box[0], box[1]), (box[2], box[3]), color=(0, \n    180, 0), thickness=2)\n', (3554, 3640), False, 'import cv2\n'), ((3102, 3142), 'numpy.where', 'np.where', (['(heatmap < args.threshold * 255)'], {}), '(heatmap < args.threshold * 255)\n', (3110, 3142), True, 'import numpy as np\n')]