--- a
+++ b/create_heatmaps_survival.py
@@ -0,0 +1,431 @@
+from __future__ import print_function
+
+import numpy as np
+
+import argparse
+
+import torch
+import torch.nn as nn
+import pdb
+import os
+import pandas as pd
+from utils.utils import *
+from math import floor
+from utils.eval_utils_survival import initiate_model as initiate_model
+from models.model_amil import AMIL
+from models.resnet_custom import resnet50_baseline
+from types import SimpleNamespace
+from collections import namedtuple
+import h5py
+import yaml
+from wsi_core.batch_process_utils import initialize_df
+from vis_utils.heatmap_utils import initialize_wsi, drawHeatmap, compute_from_patches
+from wsi_core.wsi_utils import sample_rois
+from utils.file_utils import save_hdf5
+
+parser = argparse.ArgumentParser(description='Heatmap inference script')
+parser.add_argument('--save_exp_code', type=str, default=None,
+					help='experiment code')
+parser.add_argument('--overlap', type=float, default=None)
+parser.add_argument('--config_file', type=str, default="heatmap_config_template.yaml")
+args = parser.parse_args()
+
+def infer_single_slide(model, features, label, reverse_label_dict, k=1):
+	features = features.to(device)
+	with torch.no_grad():
+		if isinstance(model, AMIL):
+			model_results_dict = model(features)
+			risk, A, _= model(features)
+
+			risk = risk.item()
+			A = A.view(-1, 1).cpu().numpy()
+
+		else:
+			raise NotImplementedError
+
+	return risk, A
+
+def load_params(df_entry, params):
+	for key in params.keys():
+		if key in df_entry.index:
+			dtype = type(params[key])
+			val = df_entry[key] 
+			val = dtype(val)
+			if isinstance(val, str):
+				if len(val) > 0:
+					params[key] = val
+			elif not np.isnan(val):
+				params[key] = val
+			else:
+				pdb.set_trace()
+
+	return params
+
+def parse_config_dict(args, config_dict):
+	if args.save_exp_code is not None:
+		config_dict['exp_arguments']['save_exp_code'] = args.save_exp_code
+	if args.overlap is not None:
+		config_dict['patching_arguments']['overlap'] = args.overlap
+	return config_dict
+
+if __name__ == '__main__':
+	config_path = os.path.join('heatmaps/configs', args.config_file)
+	config_dict = yaml.safe_load(open(config_path, 'r'))
+	config_dict = parse_config_dict(args, config_dict)
+
+	for key, value in config_dict.items():
+		if isinstance(value, dict):
+			print('\n'+key)
+			for value_key, value_value in value.items():
+				print (value_key + " : " + str(value_value))
+		else:
+			print ('\n'+key + " : " + str(value))
+			
+	decision = input('Continue? Y/N ')
+	if decision in ['Y', 'y', 'Yes', 'yes']:
+		pass
+	elif decision in ['N', 'n', 'No', 'NO']:
+		exit()
+	else:
+		raise NotImplementedError
+
+	args = config_dict
+	patch_args = argparse.Namespace(**args['patching_arguments'])
+	data_args = argparse.Namespace(**args['data_arguments'])
+	model_args = args['model_arguments']
+	model_args.update({'n_classes': args['exp_arguments']['n_classes']})
+	model_args = argparse.Namespace(**model_args)
+	exp_args = argparse.Namespace(**args['exp_arguments'])
+	heatmap_args = argparse.Namespace(**args['heatmap_arguments'])
+	sample_args = argparse.Namespace(**args['sample_arguments'])
+
+	patch_size = tuple([patch_args.patch_size for i in range(2)])
+	step_size = tuple((np.array(patch_size) * (1-patch_args.overlap)).astype(int))
+	print('patch_size: {} x {}, with {:.2f} overlap, step size is {} x {}'.format(patch_size[0], patch_size[1], patch_args.overlap, step_size[0], step_size[1]))
+
+	
+	preset = data_args.preset
+	def_seg_params = {'seg_level': -1, 'sthresh': 15, 'mthresh': 11, 'close': 2, 'use_otsu': False, 
+					  'keep_ids': 'none', 'exclude_ids':'none'}
+	def_filter_params = {'a_t':50.0, 'a_h': 8.0, 'max_n_holes':10}
+	def_vis_params = {'vis_level': -1, 'line_thickness': 250}
+	def_patch_params = {'use_padding': True, 'contour_fn': 'four_pt'}
+
+	if preset is not None:
+		preset_df = pd.read_csv(preset)
+		for key in def_seg_params.keys():
+			def_seg_params[key] = preset_df.loc[0, key]
+
+		for key in def_filter_params.keys():
+			def_filter_params[key] = preset_df.loc[0, key]
+
+		for key in def_vis_params.keys():
+			def_vis_params[key] = preset_df.loc[0, key]
+
+		for key in def_patch_params.keys():
+			def_patch_params[key] = preset_df.loc[0, key]
+
+
+	if data_args.process_list is None:
+		if isinstance(data_args.data_dir, list):
+			slides = []
+			for data_dir in data_args.data_dir:
+				slides.extend(os.listdir(data_dir))
+		else:
+			slides = sorted(os.listdir(data_args.data_dir))
+		slides = [slide for slide in slides if data_args.slide_ext in slide]
+		df = initialize_df(slides, def_seg_params, def_filter_params, def_vis_params, def_patch_params, use_heatmap_args=False)
+		
+	else:
+		df = pd.read_csv(os.path.join('heatmaps/process_lists', data_args.process_list))
+		df = initialize_df(df, def_seg_params, def_filter_params, def_vis_params, def_patch_params, use_heatmap_args=False)
+
+	mask = df['process'] == 1
+	process_stack = df[mask].reset_index(drop=True)
+	total = len(process_stack)
+	print('\nlist of slides to process: ')
+	print(process_stack.head(len(process_stack)))
+
+	print('\ninitializing model from checkpoint')
+	ckpt_path = model_args.ckpt_path
+	print('\nckpt path: {}'.format(ckpt_path))
+	
+	if model_args.initiate_fn == 'initiate_model':
+		settings = {
+			'drop_out': model_args.drop_out,
+			'model_type': model_args.model_type,
+			'model_size': model_args.model_size,
+			'print_model_info': False
+			}
+		model =  initiate_model(settings, model_args.ckpt_path)
+	else:
+		raise NotImplementedError
+
+
+	feature_extractor = resnet50_baseline(pretrained=True)
+	feature_extractor.eval()
+	device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+	print('Done!')
+
+	label_dict =  data_args.label_dict
+	class_labels = list(label_dict.keys())
+	class_encodings = list(label_dict.values())
+	reverse_label_dict = {class_encodings[i]: class_labels[i] for i in range(len(class_labels))} 
+
+	if torch.cuda.device_count() > 1:
+		device_ids = list(range(torch.cuda.device_count()))
+		feature_extractor = nn.DataParallel(feature_extractor, device_ids=device_ids).to('cuda:0')
+	else:
+		feature_extractor = feature_extractor.to(device)
+
+	os.makedirs(exp_args.production_save_dir, exist_ok=True)
+	os.makedirs(exp_args.raw_save_dir, exist_ok=True)
+	blocky_wsi_kwargs = {'top_left': None, 'bot_right': None, 'patch_size': patch_size, 'step_size': patch_size, 
+	'custom_downsample':patch_args.custom_downsample, 'level': patch_args.patch_level, 'use_center_shift': heatmap_args.use_center_shift}
+
+	for i in range(len(process_stack)):
+		slide_name = process_stack.loc[i, 'slide_id']
+		if data_args.slide_ext not in slide_name:
+			slide_name+=data_args.slide_ext
+		print('\nprocessing: ', slide_name)	
+
+		try:
+			label = process_stack.loc[i, 'label']
+		except KeyError:
+			label = 'Unspecified'
+
+		slide_id = slide_name.replace(data_args.slide_ext, '')
+
+		if not isinstance(label, str):
+			grouping = reverse_label_dict[label]
+		else:
+			grouping = label
+
+		p_slide_save_dir = os.path.join(exp_args.production_save_dir, exp_args.save_exp_code, str(grouping))
+		os.makedirs(p_slide_save_dir, exist_ok=True)
+
+		r_slide_save_dir = os.path.join(exp_args.raw_save_dir, exp_args.save_exp_code, str(grouping),  slide_id)
+		os.makedirs(r_slide_save_dir, exist_ok=True)
+
+		if heatmap_args.use_roi:
+			x1, x2 = process_stack.loc[i, 'x1'], process_stack.loc[i, 'x2']
+			y1, y2 = process_stack.loc[i, 'y1'], process_stack.loc[i, 'y2']
+			top_left = (int(x1), int(y1))
+			bot_right = (int(x2), int(y2))
+		else:
+			top_left = None
+			bot_right = None
+		
+		print('slide id: ', slide_id)
+		print('top left: ', top_left, ' bot right: ', bot_right)
+
+		if isinstance(data_args.data_dir, str):
+			slide_path = os.path.join(data_args.data_dir, slide_name)
+		elif isinstance(data_args.data_dir, dict):
+			data_dir_key = process_stack.loc[i, data_args.data_dir_key]
+			slide_path = os.path.join(data_args.data_dir[data_dir_key], slide_name)
+		else:
+			raise NotImplementedError
+
+		mask_file = os.path.join(r_slide_save_dir, slide_id+'_mask.pkl')
+		
+		# Load segmentation and filter parameters
+		seg_params = def_seg_params.copy()
+		filter_params = def_filter_params.copy()
+		vis_params = def_vis_params.copy()
+
+		seg_params = load_params(process_stack.loc[i], seg_params)
+		filter_params = load_params(process_stack.loc[i], filter_params)
+		vis_params = load_params(process_stack.loc[i], vis_params)
+
+		keep_ids = str(seg_params['keep_ids'])
+		if len(keep_ids) > 0 and keep_ids != 'none':
+			seg_params['keep_ids'] = np.array(keep_ids.split(',')).astype(int)
+		else:
+			seg_params['keep_ids'] = []
+
+		exclude_ids = str(seg_params['exclude_ids'])
+		if len(exclude_ids) > 0 and exclude_ids != 'none':
+			seg_params['exclude_ids'] = np.array(exclude_ids.split(',')).astype(int)
+		else:
+			seg_params['exclude_ids'] = []
+
+		for key, val in seg_params.items():
+			print('{}: {}'.format(key, val))
+
+		for key, val in filter_params.items():
+			print('{}: {}'.format(key, val))
+
+		for key, val in vis_params.items():
+			print('{}: {}'.format(key, val))
+		
+		print('Initializing WSI object')
+		wsi_object = initialize_wsi(slide_path, seg_mask_path=mask_file, seg_params=seg_params, filter_params=filter_params)
+		print('Done!')
+
+		wsi_ref_downsample = wsi_object.level_downsamples[patch_args.patch_level]
+
+		# the actual patch size for heatmap visualization should be the patch size * downsample factor * custom downsample factor
+		vis_patch_size = tuple((np.array(patch_size) * np.array(wsi_ref_downsample) * patch_args.custom_downsample).astype(int))
+
+		block_map_save_path = os.path.join(r_slide_save_dir, '{}_blockmap.h5'.format(slide_id))
+		mask_path = os.path.join(r_slide_save_dir, '{}_mask.jpg'.format(slide_id))
+		if vis_params['vis_level'] < 0:
+			best_level = wsi_object.wsi.get_best_level_for_downsample(32)
+			vis_params['vis_level'] = best_level
+		mask = wsi_object.visWSI(**vis_params, number_contours=True)
+		mask.save(mask_path)
+		
+		features_path = os.path.join(r_slide_save_dir, slide_id+'.pt')
+		h5_path = os.path.join(r_slide_save_dir, slide_id+'.h5')
+	
+
+		##### check if h5_features_file exists ######
+		if not os.path.isfile(h5_path) :
+			_, _, wsi_object = compute_from_patches(wsi_object=wsi_object, 
+											model=model, 
+											feature_extractor=feature_extractor, 
+											batch_size=exp_args.batch_size, **blocky_wsi_kwargs, 
+											attn_save_path=None, feat_save_path=h5_path, 
+											ref_scores=None)				
+		
+		##### check if pt_features_file exists ######
+		if not os.path.isfile(features_path):
+			file = h5py.File(h5_path, "r")
+			features = torch.tensor(file['features'][:])
+			torch.save(features, features_path)
+			file.close()
+
+		# load features 
+		features = torch.load(features_path)
+		process_stack.loc[i, 'bag_size'] = len(features)
+		
+		wsi_object.saveSegmentation(mask_file)
+		risk, A = infer_single_slide(model, features, label, reverse_label_dict, exp_args.n_classes)
+		del features
+		
+		if not os.path.isfile(block_map_save_path): 
+			file = h5py.File(h5_path, "r")
+			coords = file['coords'][:]
+			file.close()
+			asset_dict = {'attention_scores': A, 'coords': coords}
+			block_map_save_path = save_hdf5(block_map_save_path, asset_dict, mode='w')
+		
+		# save predictions
+		process_stack.loc[i, 'risk'] = risk
+
+		os.makedirs('heatmaps/results/', exist_ok=True)
+		if data_args.process_list is not None:
+			process_stack.to_csv('heatmaps/results/{}.csv'.format(data_args.process_list.replace('.csv', '')), index=False)
+		else:
+			process_stack.to_csv('heatmaps/results/{}.csv'.format(exp_args.save_exp_code), index=False)
+		
+		file = h5py.File(block_map_save_path, 'r')
+		dset = file['attention_scores']
+		coord_dset = file['coords']
+		scores = dset[:]
+		coords = coord_dset[:]
+		file.close()
+
+		samples = sample_args.samples
+		for sample in samples:
+			if sample['sample']:
+				tag = "label_{}_risk_{}".format(label, risk)
+				sample_save_dir =  os.path.join(exp_args.production_save_dir, exp_args.save_exp_code, 'sampled_patches', str(tag), sample['name'])
+				os.makedirs(sample_save_dir, exist_ok=True)
+				print('sampling {}'.format(sample['name']))
+				sample_results = sample_rois(scores, coords, k=sample['k'], mode=sample['mode'], seed=sample['seed'], 
+					score_start=sample.get('score_start', 0), score_end=sample.get('score_end', 1))
+				for idx, (s_coord, s_score) in enumerate(zip(sample_results['sampled_coords'], sample_results['sampled_scores'])):
+					print('coord: {} score: {:.3f}'.format(s_coord, s_score))
+					patch = wsi_object.wsi.read_region(tuple(s_coord), patch_args.patch_level, (patch_args.patch_size, patch_args.patch_size)).convert('RGB')
+					patch.save(os.path.join(sample_save_dir, '{}_{}_x_{}_y_{}_a_{:.3f}.png'.format(idx, slide_id, s_coord[0], s_coord[1], s_score)))
+
+		wsi_kwargs = {'top_left': top_left, 'bot_right': bot_right, 'patch_size': patch_size, 'step_size': step_size, 
+		'custom_downsample':patch_args.custom_downsample, 'level': patch_args.patch_level, 'use_center_shift': heatmap_args.use_center_shift}
+
+		heatmap_save_name = '{}_blockmap.tiff'.format(slide_id)
+		if os.path.isfile(os.path.join(r_slide_save_dir, heatmap_save_name)):
+			pass
+		else:
+			heatmap = drawHeatmap(scores, coords, slide_path, wsi_object=wsi_object, cmap=heatmap_args.cmap, alpha=heatmap_args.alpha, use_holes=True, binarize=False, vis_level=-1, blank_canvas=False,
+							thresh=-1, patch_size = vis_patch_size, convert_to_percentiles=True)
+		
+			heatmap.save(os.path.join(r_slide_save_dir, '{}_blockmap.png'.format(slide_id)))
+			del heatmap
+
+		save_path = os.path.join(r_slide_save_dir, '{}_{}_roi_{}.h5'.format(slide_id, patch_args.overlap, heatmap_args.use_roi))
+
+		if heatmap_args.use_ref_scores:
+			ref_scores = scores
+		else:
+			ref_scores = None
+		
+		if heatmap_args.calc_heatmap:
+			compute_from_patches(wsi_object=wsi_object, clam_pred=None, model=model, feature_extractor=feature_extractor, batch_size=exp_args.batch_size, **wsi_kwargs, 
+								attn_save_path=save_path,  ref_scores=ref_scores)
+
+		if not os.path.isfile(save_path):
+			print('heatmap {} not found'.format(save_path))
+			if heatmap_args.use_roi:
+				save_path_full = os.path.join(r_slide_save_dir, '{}_{}_roi_False.h5'.format(slide_id, patch_args.overlap))
+				print('found heatmap for whole slide')
+				save_path = save_path_full
+			else:
+				continue
+
+		file = h5py.File(save_path, 'r')
+		dset = file['attention_scores']
+		coord_dset = file['coords']
+		scores = dset[:]
+		coords = coord_dset[:]
+		file.close()
+
+		heatmap_vis_args = {'convert_to_percentiles': True, 'vis_level': heatmap_args.vis_level, 'blur': heatmap_args.blur, 'custom_downsample': heatmap_args.custom_downsample}
+		if heatmap_args.use_ref_scores:
+			heatmap_vis_args['convert_to_percentiles'] = False
+
+		# heatmap_save_name = '{}_{}_roi_{}_blur_{}_rs_{}_bc_{}_a_{}_l_{}_bi_{}_{}.{}'.format(slide_id, float(patch_args.overlap), int(heatmap_args.use_roi),
+		# 																				int(heatmap_args.blur), 
+		# 																				int(heatmap_args.use_ref_scores), int(heatmap_args.blank_canvas), 
+		# 																				float(heatmap_args.alpha), int(heatmap_args.vis_level), 
+		# 																				int(heatmap_args.binarize), float(heatmap_args.binary_thresh), heatmap_args.save_ext)
+
+		heatmap_save_name = '{}_heatmap_risk{:.4f}_label{}.{}'.format(slide_id, risk, label, heatmap_args.save_ext)
+
+
+		if os.path.isfile(os.path.join(p_slide_save_dir, heatmap_save_name)):
+			pass
+		
+		else:                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      
+			heatmap = drawHeatmap(scores, coords, slide_path, wsi_object=wsi_object,  
+						          cmap=heatmap_args.cmap, alpha=heatmap_args.alpha, **heatmap_vis_args, 
+						          binarize=heatmap_args.binarize, 
+						  		  blank_canvas=heatmap_args.blank_canvas,
+						  		  thresh=heatmap_args.binary_thresh,  patch_size = vis_patch_size,
+						  		  overlap=patch_args.overlap, 
+						  		  top_left=top_left, bot_right = bot_right)
+			if heatmap_args.save_ext == 'jpg':
+				heatmap.save(os.path.join(p_slide_save_dir, heatmap_save_name), quality=100)
+			else:
+				heatmap.save(os.path.join(p_slide_save_dir, heatmap_save_name))
+		
+		if heatmap_args.save_orig:
+			if heatmap_args.vis_level >= 0:
+				vis_level = heatmap_args.vis_level
+			else:
+				vis_level = vis_params['vis_level']
+			heatmap_save_name = '{}_orig_{}.{}'.format(slide_id,int(vis_level), heatmap_args.save_ext)
+			if os.path.isfile(os.path.join(p_slide_save_dir, heatmap_save_name)):
+				pass
+			else:
+				heatmap = wsi_object.visWSI(vis_level=vis_level, view_slide_only=True, custom_downsample=heatmap_args.custom_downsample)
+				if heatmap_args.save_ext == 'jpg':
+					heatmap.save(os.path.join(p_slide_save_dir, heatmap_save_name), quality=100)
+				else:
+					heatmap.save(os.path.join(p_slide_save_dir, heatmap_save_name))
+
+	with open(os.path.join(exp_args.raw_save_dir, exp_args.save_exp_code, 'config.yaml'), 'w') as outfile:
+		yaml.dump(config_dict, outfile, default_flow_style=False)
+
+