Upload app_utils.py (#2)

- Upload app_utils.py (35aa4d8d39070d52eb22e8e048bfa065e986bd78)


Co-authored-by: Kong-Yi <roseDwayane@users.noreply.huggingface.co>

app_utils.py

@@ -0,0 +1,331 @@
+import utils
+import os
+import math
+import json
+import jsbeautifier
+import numpy as np
+import matplotlib.pyplot as plt
+import mne
+from mne.channels import read_custom_montage
+from scipy.interpolate import Rbf
+from scipy.optimize import linear_sum_assignment
+from sklearn.neighbors import NearestNeighbors
+
+def get_matched(tpl_names, tpl_dict):
+	return [name for name in tpl_names if tpl_dict[name]["matched"]==True]
+
+def get_empty_template(tpl_names, tpl_dict):
+	return [name for name in tpl_names if tpl_dict[name]["matched"]==False]
+
+def get_unassigned_input(in_names, in_dict):
+	return [name for name in in_names if in_dict[name]["assigned"]==False]
+
+def read_montage(loc_file):
+	tpl_montage = read_custom_montage("./template_chanlocs.loc")
+	in_montage = read_custom_montage(loc_file)
+	tpl_names = tpl_montage.ch_names
+	in_names = in_montage.ch_names
+	tpl_dict = {}
+	in_dict = {}
+	
+	# convert all channel names to uppercase and store their information
+	for i, name in enumerate(tpl_names):
+		up_name = str.upper(name)
+		tpl_montage.rename_channels({name: up_name})
+		tpl_dict[up_name] = {
+			"index" : i,
+			"coord_3d" : tpl_montage.get_positions()['ch_pos'][up_name],
+			"matched" : False
+		}
+	for i, name in enumerate(in_names):
+		up_name = str.upper(name)
+		in_montage.rename_channels({name: up_name})
+		in_dict[up_name] = {
+			"index" : i,
+			"coord_3d" : in_montage.get_positions()['ch_pos'][up_name],
+			"assigned" : False
+		}
+	return tpl_montage, in_montage, tpl_dict, in_dict
+
+def match_name(stage1_info):
+	# read the location file
+	loc_file = stage1_info["fileNames"]["inputData"]
+	tpl_montage, in_montage, tpl_dict, in_dict = read_montage(loc_file)
+	tpl_names = tpl_montage.ch_names
+	in_names = in_montage.ch_names
+	old_idx = [[None]]*30 # store the indices of the in_channels in the order of tpl_channels
+	is_orig_data = [False]*30
+	
+	alias_dict = {
+		'T3': 'T7',
+		'T4': 'T8',
+		'T5': 'P7',
+		'T6': 'P8'
+	}
+	for i, name in enumerate(tpl_names):
+		if name in alias_dict and alias_dict[name] in in_dict:
+			tpl_montage.rename_channels({name: alias_dict[name]})
+			tpl_dict[alias_dict[name]] = tpl_dict.pop(name)
+			name = alias_dict[name]
+			
+		if name in in_dict:
+			old_idx[i] = [in_dict[name]["index"]]
+			is_orig_data[i] = True
+			tpl_dict[name]["matched"] = True
+			in_dict[name]["assigned"] = True
+	
+	# update the names
+	tpl_names = tpl_montage.ch_names
+	
+	stage1_info.update({
+		"unassignedInput" : get_unassigned_input(in_names, in_dict),
+		"emptyTemplate" : get_empty_template(tpl_names, tpl_dict),
+	 	"mappingResult" : [
+		{
+			"index" : old_idx,
+			"isOriginalData" : is_orig_data
+		}
+		]
+	})
+	channel_info = {
+		"templateNames" : tpl_names,
+		"inputNames" : in_names,
+		"templateDict" : tpl_dict,
+		"inputDict" : in_dict
+	}
+	return stage1_info, channel_info, tpl_montage, in_montage
+
+def align_coords(channel_info, tpl_montage, in_montage):
+	tpl_names = channel_info["templateNames"]
+	in_names = channel_info["inputNames"]
+	tpl_dict = channel_info["templateDict"]
+	in_dict = channel_info["inputDict"]
+	matched_names = get_matched(tpl_names, tpl_dict)
+	
+	# 2D alignment (for visualization purposes)
+	fig = [tpl_montage.plot(), in_montage.plot()]
+	ax = [fig[0].axes[0], fig[1].axes[0]]
+	
+	# extract the displayed 2D coordinates
+	all_tpl = ax[0].collections[0].get_offsets().data
+	all_in= ax[1].collections[0].get_offsets().data
+	matched_tpl = np.array([all_tpl[tpl_dict[name]["index"]] for name in matched_names])
+	matched_in = np.array([all_in[in_dict[name]["index"]] for name in matched_names])
+	plt.close('all')
+	
+	# apply TPS to transform in_channels to align with tpl_channels positions
+	rbf_x = Rbf(matched_in[:,0], matched_in[:,1], matched_tpl[:,0], function='thin_plate')
+	rbf_y = Rbf(matched_in[:,0], matched_in[:,1], matched_tpl[:,1], function='thin_plate')
+	
+	# apply the transformation to all in_channels
+	transformed_in_x = rbf_x(all_in[:,0], all_in[:,1])
+	transformed_in_y = rbf_y(all_in[:,0], all_in[:,1])
+	transformed_in = np.vstack((transformed_in_x, transformed_in_y)).T
+	
+	for i, name in enumerate(tpl_names):
+		tpl_dict[name]["coord_2d"] = all_tpl[i]
+	for i, name in enumerate(in_names):
+		in_dict[name]["coord_2d"] = transformed_in[i].tolist()
+	
+	
+	# 3D alignment
+	all_tpl = np.array([tpl_dict[name]["coord_3d"].tolist() for name in tpl_names])
+	all_in = np.array([in_dict[name]["coord_3d"].tolist() for name in in_names])
+	matched_tpl = np.array([all_tpl[tpl_dict[name]["index"]] for name in matched_names])
+	matched_in = np.array([all_in[in_dict[name]["index"]] for name in matched_names])
+	
+	rbf_x = Rbf(matched_in[:,0], matched_in[:,1], matched_in[:,2], matched_tpl[:,0], function='thin_plate')
+	rbf_y = Rbf(matched_in[:,0], matched_in[:,1], matched_in[:,2], matched_tpl[:,1], function='thin_plate')
+	rbf_z = Rbf(matched_in[:,0], matched_in[:,1], matched_in[:,2], matched_tpl[:,2], function='thin_plate')
+	
+	transformed_in_x = rbf_x(all_in[:,0], all_in[:,1], all_in[:,2])
+	transformed_in_y = rbf_y(all_in[:,0], all_in[:,1], all_in[:,2])
+	transformed_in_z = rbf_z(all_in[:,0], all_in[:,1], all_in[:,2])
+	transformed_in = np.vstack((transformed_in_x, transformed_in_y, transformed_in_z)).T
+	
+	for i, name in enumerate(in_names):
+		in_dict[name]["coord_3d"] = transformed_in[i].tolist()
+	
+	channel_info.update({
+		"templateDict" : tpl_dict,
+		"inputDict" : in_dict
+	})
+	return channel_info
+
+def save_figure(channel_info, tpl_montage, filename1, filename2):
+	tpl_names = channel_info["templateNames"]
+	in_names = channel_info["inputNames"]
+	tpl_dict = channel_info["templateDict"]
+	in_dict = channel_info["inputDict"]
+	
+	tpl_x = [tpl_dict[name]["coord_2d"][0] for name in tpl_names]
+	tpl_y = [tpl_dict[name]["coord_2d"][1] for name in tpl_names]
+	in_x = [in_dict[name]["coord_2d"][0] for name in in_names]
+	in_y = [in_dict[name]["coord_2d"][1] for name in in_names]
+	tpl_coords = np.vstack((tpl_x, tpl_y)).T
+	in_coords = np.vstack((in_x, in_y)).T
+	
+	# extract template's head figure
+	tpl_fig = tpl_montage.plot()
+	tpl_ax = tpl_fig.axes[0]
+	lines = tpl_ax.lines
+	head_lines = []
+	for line in lines:
+		x, y = line.get_data()
+		head_lines.append((x,y))
+	
+	# -------------------------plot input montage------------------------------
+	fig = plt.figure(figsize=(6.4,6.4), dpi=100)
+	ax = fig.add_subplot(111)
+	fig.tight_layout()
+	ax.set_aspect('equal')
+	ax.axis('off')
+	
+	# plot template's head
+	for x, y in head_lines:
+		ax.plot(x, y, color='black', linewidth=1.0)
+	# plot in_channels on it
+	ax.scatter(in_coords[:,0], in_coords[:,1], s=35, color='black')
+	for i, name in enumerate(in_names):
+		ax.text(in_coords[i,0]+0.004, in_coords[i,1], name, color='black', fontsize=10.0, va='center')
+	# save input_montage
+	fig.savefig(filename1)
+	
+	# ---------------------------add indications-------------------------------
+	# plot unmatched input channels in red
+	indices = [in_dict[name]["index"] for name in in_names if in_dict[name]["assigned"]==False]
+	if indices != []:
+		ax.scatter(in_coords[indices,0], in_coords[indices,1], s=35, color='red')
+		for i in indices:
+			ax.text(in_coords[i,0]+0.004, in_coords[i,1], in_names[i], color='red', fontsize=10.0, va='center')
+	# save mapped_montage
+	fig.savefig(filename2)
+	
+	# -------------------------------------------------------------------------
+	# store the tpl and in_channels' display positions (in px)
+	tpl_coords = ax.transData.transform(tpl_coords)
+	in_coords = ax.transData.transform(in_coords)
+	plt.close('all')
+	
+	for i, name in enumerate(tpl_names):
+		left = tpl_coords[i,0]/6.4
+		bottom = tpl_coords[i,1]/6.4
+		tpl_dict[name]["css_position"] = [round(left, 2), round(bottom, 2)]
+	for i, name in enumerate(in_names):
+		left = in_coords[i,0]/6.4
+		bottom = in_coords[i,1]/6.4
+		in_dict[name]["css_position"] = [round(left, 2), round(bottom, 2)]
+	
+	channel_info.update({
+		"templateDict" : tpl_dict,
+		"inputDict" : in_dict
+	})
+	return channel_info
+
+def find_neighbors(channel_info, empty_tpl_names, old_idx):
+	in_names = channel_info["inputNames"]
+	tpl_dict = channel_info["templateDict"]
+	in_dict = channel_info["inputDict"]
+	
+	all_in = [np.array(in_dict[name]["coord_3d"]) for name in in_names]
+	empty_tpl = [np.array(tpl_dict[name]["coord_3d"]) for name in empty_tpl_names]
+	
+	# use KNN to choose k nearest channels
+	k = 4 if len(in_names)>4 else len(in_names)
+	knn = NearestNeighbors(n_neighbors=k, metric='euclidean')
+	knn.fit(all_in)
+	for i, name in enumerate(empty_tpl_names):
+		distances, indices = knn.kneighbors(empty_tpl[i].reshape(1,-1))
+		idx = tpl_dict[name]["index"]
+		old_idx[idx] = indices[0].tolist()
+	
+	return old_idx
+
+def optimal_mapping(channel_info):
+	tpl_names = channel_info["templateNames"]
+	in_names = channel_info["inputNames"]
+	tpl_dict = channel_info["templateDict"]
+	in_dict = channel_info["inputDict"]
+	unass_in_names = get_unassigned_input(in_names, in_dict)
+	# reset all tpl.matched to False
+	for name in tpl_dict:
+		tpl_dict[name]["matched"] = False
+	
+	all_tpl = np.array([tpl_dict[name]["coord_3d"] for name in tpl_names])
+	unass_in = np.array([in_dict[name]["coord_3d"] for name in unass_in_names])
+	
+	# initialize the cost matrix for the Hungarian algorithm
+	if len(unass_in_names) < 30:
+		cost_matrix = np.full((30, 30), 1e6) # add dummy channels to ensure num_col >= num_row
+	else:
+		cost_matrix = np.zeros((30, len(unass_in_names)))
+	# fill the cost matrix with Euclidean distances between tpl and unassigned in_channels
+	for i in range(30):
+		for j in range(len(unass_in_names)):
+			cost_matrix[i][j] = np.linalg.norm((all_tpl[i]-unass_in[j])*1000)
+	
+	# apply the Hungarian algorithm to optimally assign one in_channel to each tpl_channel
+	# by minimizing the total distances between their positions.
+	row_idx, col_idx = linear_sum_assignment(cost_matrix)
+	
+	# store the mapping result
+	old_idx = [[None]]*30
+	is_orig_data = [False]*30
+	for i, j in zip(row_idx, col_idx):
+		if j < len(unass_in_names): # filter out dummy channels
+			tpl_name = tpl_names[i]
+			in_name = unass_in_names[j]
+			
+			old_idx[i] = [in_dict[in_name]["index"]]
+			is_orig_data[i] = True
+			tpl_dict[tpl_name]["matched"] = True
+			in_dict[in_name]["assigned"] = True
+	
+	# fill the remaining empty tpl_channels
+	empty_tpl_names = get_empty_template(tpl_names, tpl_dict)
+	if empty_tpl_names != []:
+		old_idx = find_neighbors(channel_info, empty_tpl_names, old_idx)
+	
+	result = {
+		"index" : old_idx,
+		"isOriginalData" : is_orig_data
+	}
+	channel_info["inputDict"] = in_dict
+	return result, channel_info
+
+def mapping_result(stage1_info, channel_info, filename):
+	unassigned_num = len(stage1_info["unassignedInput"])
+	batch = math.ceil(unassigned_num/30) + 1
+	
+	# map the remaining in_channels
+	results = stage1_info["mappingResult"]
+	for i in range(1, batch):
+		# optimally select 30 in_channels to map to the tpl_channels based on proximity
+		result, channel_info = optimal_mapping(channel_info)
+		results += [result]
+	'''
+	for i in range(batch):
+		results[i]["name"] = {}
+		for j, indices in enumerate(results[i]["index"]):
+			names = [channel_info["inputNames"][idx] for idx in indices] if indices!=[None] else ["zero"]
+			results[i]["name"][channel_info["templateNames"][j]] = names
+	'''
+	data = {
+		#"templateNames" : channel_info["templateNames"],
+		#"inputNames" : channel_info["inputNames"],
+		"channelNum" : len(channel_info["inputNames"]),
+		"batch" : batch,
+		"mappingResult" : results
+	}
+	options = jsbeautifier.default_options()
+	options.indent_size = 4
+	json_data = jsbeautifier.beautify(json.dumps(data), options)
+	with open(filename, 'w') as jsonfile:
+		jsonfile.write(json_data)
+	
+	stage1_info.update({
+		"batch" : batch,
+		"mappingResult" : results
+	})
+	return stage1_info, channel_info
+