update

app.py

@@ -1,13 +1,16 @@
 import gradio as gr
-import numpy as np
+
 import os
 import random
 import math
+import numpy as np
+import matplotlib.pyplot as plt
+import mne
+from mne.channels import read_custom_montage
+
 import utils
 from channel_mapping import mapping_stage1, mapping_stage2, reorder_to_template, reorder_to_origin, find_neighbors
 
-import mne
-from mne.channels import read_custom_montage
 
 quickstart = """
 
@@ -56,51 +59,52 @@ init_js = """
 	app_state = JSON.parse(JSON.stringify(app_state));
 	channel_info = JSON.parse(JSON.stringify(channel_info));
 	
-	let selector, attribute;
+	let selector, classname, attribute;
 	let channel, left, bottom;
 	
 	if(app_state.stage1State == "step2-selecting"){
-	    selector = "#radio > div:nth-of-type(2)";
+	    selector = "#radio-group > div:nth-of-type(2)";
+	    //classname = "radio";
 	    attribute = "value";
 	}else if(app_state.stage1State == "step3-selecting"){
 	    selector = "#chkbox-group > div:nth-of-type(2)";
+	     //classname = "chkbox";
 	    attribute = "name";
 	}else return;
 	
-	// add figure of mapped_montage
+	
+	// add figure of the mapping result
 	document.querySelector(selector).style.cssText = `
 		position: relative;
-		width: 560px;
-		height: 560px;
+		width: 100%;
+		aspect-ratio: 1;
+		//width: 560px;
+		//height: 560px;
 		background: url("file=${app_state.filenames.raw_montage}");
+		background-size: contain;
+
 	`;
 	
-	
 	// move the radios/checkboxes
 	let all_elem = document.querySelectorAll(selector+" > label");
-	Array.from(all_elem).forEach((item) => {
+	Array.from(all_elem).forEach(item => {
 		channel = item.querySelector("input").getAttribute(attribute);
 		left = channel_info.inputDict[channel].css_position[0];
 		bottom = channel_info.inputDict[channel].css_position[1];
-		//console.log(`channel: ${channel}, left: ${left}, bottom: ${bottom}`);
 		
-		item.style.cssText = `
-			position: absolute;
-			left: ${left};
-			bottom: ${bottom};
-		`;
-		item.className = "";
+		item.style.cssText = `position: absolute; left: ${left}; bottom: ${bottom};`;
+		item.className = ""; //classname;
 		item.querySelector(":scope > span").innerText = "";
 	});
 	
 	
 	// add indication for the missing channels
-	channel = app_state.missingTemplates[0]
+	channel = app_state.missingTemplates[0];
 	left = channel_info.templateDict[channel].css_position[0];
 	bottom = channel_info.templateDict[channel].css_position[1];
 	
 	let dot_rule = `
-		${selector}::before{
+		${selector}::before {
 			content: '';
 			position: absolute;
 			background-color: red;
@@ -112,12 +116,12 @@ init_js = """
 		}
 	`;
 	
-	left = parseFloat(left.slice(0, -1))+2.5
-	left = left.toString()+"%"
-	bottom = parseFloat(bottom.slice(0, -1))-1
-	bottom = bottom.toString()+"%"
+	left = parseFloat(left.slice(0, -1))+2.5;
+	left = left.toString()+"%";
+	bottom = parseFloat(bottom.slice(0, -1))-1;
+	bottom = bottom.toString()+"%";
 	let txt_rule = `
-		${selector}::after{
+		${selector}::after {
 			content: "${channel}";
 			position: absolute;
 			color: red;
@@ -149,21 +153,16 @@ update_js = """
 	let channel, left, bottom;
 	
 	if(app_state.stage1State == "step2-selecting"){
-	    selector = "#radio > div:nth-of-type(2)";
+	    selector = "#radio-group > div:nth-of-type(2)";
 	    
 	    // update the radios
 		let all_elem = document.querySelectorAll(selector+" > label");
-		Array.from(all_elem).forEach((item) => {
+		Array.from(all_elem).forEach(item => {
 			channel = item.querySelector("input").value;
 			left = channel_info.inputDict[channel].css_position[0];
 			bottom = channel_info.inputDict[channel].css_position[1];
-			//console.log(`channel: ${channel}, left: ${left}, bottom: ${bottom}`);
 			
-			item.style.cssText = `
-				position: absolute;
-				left: ${left};
-				bottom: ${bottom};
-			`;
+			item.style.cssText = `position: absolute; left: ${left}; bottom: ${bottom};`;
 			item.className = "";
 			item.querySelector(":scope > span").innerText = "";
 		});
@@ -172,12 +171,12 @@ update_js = """
 	}else return;
 	
 	// update indication
-	channel = app_state.missingTemplates[app_state["fillingCount"]-1]
+	channel = app_state.missingTemplates[app_state["fillingCount"]-1];
 	left = channel_info.templateDict[channel].css_position[0];
 	bottom = channel_info.templateDict[channel].css_position[1];
 	
 	let dot_rule = `
-		${selector}::before{
+		${selector}::before {
 			content: "";
 			position: absolute;
 			background-color: red;
@@ -189,12 +188,12 @@ update_js = """
 		}
 	`;
 	
-	left = parseFloat(left.slice(0, -1))+2.5
-	left = left.toString()+"%"
-	bottom = parseFloat(bottom.slice(0, -1))-1
-	bottom = bottom.toString()+"%"
+	left = parseFloat(left.slice(0, -1))+2.5;
+	left = left.toString()+"%";
+	bottom = parseFloat(bottom.slice(0, -1))-1;
+	bottom = bottom.toString()+"%";
 	let txt_rule = `
-		${selector}::after{
+		${selector}::after {
 			content: "${channel}";
 			position: absolute;
 			color: red;
@@ -208,7 +207,6 @@ update_js = """
 	for(let i=0; i<styleSheet.cssRules.length; i++){
 		let tmp = styleSheet.cssRules[i].selectorText;
 		if(tmp==selector+"::before" || tmp==selector+"::after"){
-			console.log('exist!!', tmp);
 			styleSheet.deleteRule(i);
 			i--;
 		}
@@ -249,10 +247,10 @@ with gr.Blocks() as demo:
 			# stage1-1 : mapping result
 			with gr.Row():
 				tpl_montage = gr.Image("./template_montage.png", label="Template montage", visible=False)
-				mapped_montage = gr.Image(elem_id="input-montage", label="Input channels", visible=False)
+				mapped_montage = gr.Image(label="Input channels", visible=False)
 			
 			# stage1-2 : assign unmatched input channels to empty template channels
-			radio = gr.Radio(elem_id="radio", visible=False)
+			radio_group = gr.Radio(elem_id="radio-group", visible=False)
 			
 			# stage1-3 : select a way to fill the empty template channels
 			with gr.Row():
@@ -291,10 +289,6 @@ with gr.Blocks() as demo:
 			batch_md = gr.Markdown(visible=False)
 			out_denoised_data = gr.File(label="Denoised data", visible=False)
 			
-			#files = []
-			#for i in range():
-				#f = gr.File()
-				#files.append(f)
 			# -------------------------------------------------------	
 				
 	with gr.Row():
@@ -309,7 +303,7 @@ with gr.Blocks() as demo:
 		with gr.Tab("QuickStart"):
 			gr.Markdown(quickstart)
     
-	#demo.load(js=js)
+	#demo.load(js=tmp_js)
 	
 	# -------------------------stage1: channel mapping-------------------------------
 	def reset_all(raw_data, raw_loc, samplerate):
@@ -330,17 +324,14 @@ with gr.Blocks() as demo:
 			os.mkdir(filepath+"/temp_data/")
 			#print(e)
 		
-		# initialize app_state, channel_info
-		#data = utils.read_train_data(raw_data)
+		# initialize channel_info, app_state
+		channel_info = {}
 		app_state = {
 			"filepath": filepath+"/temp_data/",
 			"filenames": {},
 			"sampleRate": int(samplerate),
 			"stage1State" : "step1"
 		}
-		channel_info = {
-			#"dataShape" : data.shape
-		}
 		
 		# reset layout
 		return {app_state_json : app_state,
@@ -349,7 +340,7 @@ with gr.Blocks() as demo:
 				desc_md : gr.Markdown("### Step1: Mapping result", visible=False),
 				tpl_montage : gr.Image(visible=False),
 				mapped_montage : gr.Image(value=None, visible=False),
-				radio : gr.Radio(choices=[], value=[], label="", visible=False),
+				radio_group : gr.Radio(choices=[], value=[], label="", visible=False),
 				in_fill_mode : gr.Dropdown(value="mean", visible=False),
 				chkbox_group : gr.CheckboxGroup(choices=[], value=[], label="", visible=False),
 				fillmode_btn : gr.Button(visible=False),
@@ -364,28 +355,94 @@ with gr.Blocks() as demo:
 	
 	
 	# ---------------------------stage1-1-------------------------------
-	def mapping_result(app_state, channel_info, raw_loc):
+	def save_figures(channel_info, filename1, filename2):
+		
+		template_montage = read_custom_montage("./template_chanlocs.loc")
+		template_dict = channel_info["templateDict"]
+		input_dict = channel_info["inputDict"]
+		template_order = channel_info["templateOrder"]
+		input_order = channel_info["inputOrder"]
+		
+		# get template's head figure
+		tpl_fig = template_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))
+		plt.close()
+		
+		# get template's and input's 2d coords
+		tpl_x = [template_dict[channel]["coord_2d"][0] for channel in template_order]
+		tpl_y = [template_dict[channel]["coord_2d"][1] for channel in template_order]
+		in_x = [input_dict[channel]["coord_2d"][0] for channel in input_order]
+		in_y = [input_dict[channel]["coord_2d"][1] for channel in input_order]
+		tpl_coords = np.vstack((tpl_x, tpl_y)).T
+		in_coords = np.vstack((in_x, in_y)).T
+		
+		# -------------------------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 input channels
+		ax.scatter(in_coords[:,0], in_coords[:,1], s=35, color='black')
+		for i, channel in enumerate(input_order):
+			ax.text(in_coords[i,0]+0.003, in_coords[i,1], channel, color='black', fontsize=10.0, va='center')
+		
+		# save raw_montage
+		fig.savefig(filename1)
+		
+		# ---------------------------add indications-------------------------------
+		indices = [input_dict[channel]["index"] for channel in input_order if input_dict[channel]["assigned"]==False]
+		
+		# plot unmatched input channels in red
+		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.003, in_coords[i,1], input_order[i], color='red', fontsize=10.0, va='center')
+		
+		# save mapped_montage
+		fig.savefig(filename2)
+		plt.close()
+		
+		# -------------------------------------------------------------------------
+		# save the template and input channels' display position (in px).
+		tpl_coords = ax.transData.transform(tpl_coords)
+		in_coords = ax.transData.transform(in_coords)
+		
+		for i, channel in enumerate(template_order):
+			css_left = (tpl_coords[i,0]-11)/6.4
+			css_bottom = (tpl_coords[i,1]-7)/6.4
+			template_dict[channel]["css_position"] = [str(round(css_left, 2))+"%", str(round(css_bottom, 2))+"%"]
+
+		for i, channel in enumerate(input_order):
+			css_left = (in_coords[i,0]-11)/6.4
+			css_bottom = (in_coords[i,1]-7)/6.4
+			input_dict[channel]["css_position"] = [str(round(css_left, 2))+"%", str(round(css_bottom, 2))+"%"]
+		
+		channel_info.update({
+			"templateDict" : template_dict,
+			"inputDict" : input_dict
+		})
+		return channel_info
+	
+	def mapping_result(app_state, channel_info):
 		filepath = app_state["filepath"]
 		filename1 = filepath+"raw_montage_"+str(random.randint(1,10000))+".png"
 		filename2 = filepath+"mapped_montage_"+str(random.randint(1,10000))+".png"
+		channel_info = save_figures(channel_info, filename1, filename2)
 		
 		app_state["filenames"].update({
 		    "raw_montage" : filename1,
 		    "mapped_montage" : filename2
 		})
 		
-		raw_montage = read_custom_montage(raw_loc)
-		raw_fig = raw_montage.plot()
-		raw_fig.set_size_inches(5.6, 5.6)
-		raw_fig.savefig(filename1, pad_inches=0)
-		
-		# plot red dots on unmatched input channels
-		ax = raw_fig.axes[0]
-		coords = ax.collections[0].get_offsets().data
-		idx = [channel_info["inputDict"][channel]["index"] for channel in app_state["stage1UnassignedInputs"]]
-		ax.scatter(coords[idx,0]-0.0001, coords[idx,1]+0.0001, color='red')
-		raw_fig.savefig(filename2, pad_inches=0)
-		
 		# ------------------determine the next step-----------------------
 		
 		in_num = len(channel_info["inputOrder"])
@@ -398,6 +455,7 @@ with gr.Blocks() as demo:
 			gr.Info('The mapping process has been finished.')
 			
 			return {app_state_json : app_state,
+					channel_info_json : channel_info,
 					desc_md : gr.Markdown("### Mapping result", visible=True),
 					tpl_montage : gr.Image(visible=True),
 					mapped_montage : gr.Image(value=filename2, visible=True),
@@ -414,6 +472,7 @@ with gr.Blocks() as demo:
 			app_state["stage1State"] = "step3-initializing"
 		
 		return {app_state_json : app_state,
+				channel_info_json : channel_info,
 				desc_md : gr.Markdown("### Step1: Mapping result", visible=True),
 				tpl_montage : gr.Image(visible=True),
 				mapped_montage : gr.Image(value=filename2, visible=True),
@@ -422,8 +481,9 @@ with gr.Blocks() as demo:
 	map_btn.click(
 	    fn = reset_all,
 	    inputs = [in_raw_data, in_raw_loc, in_samplerate],
-	    outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio, in_fill_mode, chkbox_group,
-	               fillmode_btn, clear_btn, step2_btn, step3_btn, next_btn, run_btn, batch_md, out_denoised_data]
+	    outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio_group,
+	                in_fill_mode, chkbox_group, fillmode_btn, clear_btn, step2_btn, step3_btn, next_btn,
+	                run_btn, batch_md, out_denoised_data]
 	).success(
 	    fn = mapping_stage1,
 	    inputs = [app_state_json, channel_info_json, in_raw_loc], 
@@ -431,8 +491,8 @@ with gr.Blocks() as demo:
 	    
 	).success(
 	    fn = mapping_result,
-	    inputs = [app_state_json, channel_info_json, in_raw_loc],
-	    outputs = [app_state_json, desc_md, tpl_montage, mapped_montage, next_btn, run_btn]
+	    inputs = [app_state_json, channel_info_json],
+	    outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, next_btn, run_btn]
 	)
 	
 
@@ -458,7 +518,7 @@ with gr.Blocks() as demo:
 					    desc_md : gr.Markdown("### Step2: Assign unmatched input channels"),
 					    tpl_montage : gr.Image(visible=False),
 					    mapped_montage : gr.Image(visible=False),
-					    radio : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
+					    radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
 					    clear_btn : gr.Button(visible=True),
 					    next_btn : gr.Button("Next step")}
 			else:
@@ -467,7 +527,7 @@ with gr.Blocks() as demo:
 					    desc_md : gr.Markdown("### Step2: Assign unmatched input channels"),
 					    tpl_montage : gr.Image(visible=False),
 					    mapped_montage : gr.Image(visible=False),
-					    radio : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
+					    radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
 					    clear_btn : gr.Button(visible=True),
 					    step2_btn : gr.Button(visible=True),
 					    next_btn : gr.Button(visible=False)}
@@ -524,7 +584,7 @@ with gr.Blocks() as demo:
 				return {app_state_json : app_state,
 						channel_info_json : channel_info,
 						desc_md : gr.Markdown(visible=False),
-						radio : gr.Radio(visible=False),
+						radio_group : gr.Radio(visible=False),
 						clear_btn : gr.Button(visible=False),
 						next_btn : gr.Button(visible=False),
 						run_btn : gr.Button(interactive=True)}
@@ -540,7 +600,7 @@ with gr.Blocks() as demo:
 				return {app_state_json : app_state,
 						channel_info_json : channel_info,
 						desc_md : gr.Markdown("### Step3: Fill the remaining template channels"),
-						radio : gr.Radio(visible=False),
+						radio_group : gr.Radio(visible=False),
 						in_fill_mode : gr.Dropdown(visible=True),
 						fillmode_btn : gr.Button(visible=True),
 						clear_btn : gr.Button(visible=False),
@@ -574,9 +634,9 @@ with gr.Blocks() as demo:
 	
 	next_btn.click(
 		fn = init_next_step,
-		inputs = [app_state_json, channel_info_json, radio, chkbox_group],
-		outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio, in_fill_mode,
-					chkbox_group, fillmode_btn, clear_btn, step2_btn, next_btn, run_btn]
+		inputs = [app_state_json, channel_info_json, radio_group, chkbox_group],
+		outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio_group,
+		            in_fill_mode, chkbox_group, fillmode_btn, clear_btn, step2_btn, next_btn, run_btn]
 	).success(
 	    fn = None,
 	    js = init_js,
@@ -611,18 +671,20 @@ with gr.Blocks() as demo:
 		if len(app_state["stage1UnassignedInputs"])==1 or app_state["fillingCount"]==app_state["totalFillingNum"]:
 			return {app_state_json : app_state,
 					channel_info_json : channel_info,
-					radio : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=radio_label),
+					radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"],
+					                        value=[], label=radio_label),
 					step2_btn : gr.Button(visible=False),
 					next_btn : gr.Button("Next step", visible=True)}
 		else:
 			return {app_state_json : app_state,
 					channel_info_json : channel_info,
-					radio : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=radio_label)}
+					radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"],
+					                        value=[], label=radio_label)}
 	
 	step2_btn.click(
 	    fn = update_radio,
-	    inputs = [app_state_json, channel_info_json, radio],
-	    outputs = [app_state_json, channel_info_json, radio, step2_btn, next_btn]
+	    inputs = [app_state_json, channel_info_json, radio_group],
+	    outputs = [app_state_json, channel_info_json, radio_group, step2_btn, next_btn]
 	    
 	).success(
 	    fn = None,
@@ -634,7 +696,7 @@ with gr.Blocks() as demo:
 	clear_btn.click(
 	    fn = lambda : gr.Radio(value=[]),
 	    inputs = [],
-	    outputs = radio
+	    outputs = radio_group
 	)
 	
 	
@@ -776,9 +838,6 @@ with gr.Blocks() as demo:
 		new_filename = app_state["filenames"]["denoised"]
 		
 		while app_state["runningState"] != "finished":
-			#if app_state["batchCount"] > app_state["totalBatchNum"]:
-				#app_state["runningState"] = "finished"
-				#break
 			md = 'Running model('+str(app_state["batchCount"])+'/'+str(app_state["totalBatchNum"])+')...'
 			yield {batch_md : gr.Markdown(md, visible=True)}
 			

channel_mapping.py

@@ -57,17 +57,6 @@ def reorder_to_origin(app_state, channel_info, new_filename):
 	utils.save_data(new_data, new_filename)
 	return
 
-class Channel:
-
-	def __init__(self, index, name=None, matched=False, assigned=False, coord=None, css_position=None):
-		self.name = name
-		self.index = index
-		self.matched = matched
-		self.assigned = assigned # for input channels
-		self.coord = coord
-		self.css_position = css_position
-
-
 def read_montage_data(loc_file):
 		
 	template_montage = read_custom_montage("./template_chanlocs.loc")
@@ -75,17 +64,26 @@ def read_montage_data(loc_file):
 	template_dict = {}
 	input_dict = {}
 	
-	montages = [template_montage, input_montage]
-	dicts = [template_dict, input_dict]
-	num = [30, len(input_montage.ch_names)]
-	
-	for i in range(2):
-		for j in range(num[i]):
-			channel = montages[i].ch_names[j]
-			montages[i].rename_channels({channel: str.upper(channel)}) # convert all channel names to uppercase
-	        
-			channel = str.upper(channel)
-			dicts[i][channel] = Channel(index=j, name=channel, coord=montages[i].get_positions()['ch_pos'][channel])
+	for i in range(30):
+		channel = template_montage.ch_names[i]
+		template_montage.rename_channels({channel: str.upper(channel)}) # convert all channel names to uppercase
+		
+		channel = str.upper(channel)
+		template_dict[channel] = {
+			"index" : i,
+			"coord_3d" : template_montage.get_positions()['ch_pos'][channel],
+			"matched" : False
+		}
+	for i in range(len(input_montage.ch_names)):
+		channel = input_montage.ch_names[i]
+		input_montage.rename_channels({channel: str.upper(channel)}) # convert all channel names to uppercase
+		
+		channel = str.upper(channel)
+		input_dict[channel] = {
+			"index" : i,
+			"coord_3d" : input_montage.get_positions()['ch_pos'][channel],
+			"assigned" : False
+		}
 	
 	return template_montage, input_montage, template_dict, input_dict
 
@@ -95,49 +93,43 @@ def align_coords(channel_info, template_montage, input_montage):
 	input_dict = channel_info["inputDict"]
 	template_order = channel_info["templateOrder"]
 	input_order = channel_info["inputOrder"]
-	matched = [channel for channel in input_order if input_dict[channel]["matched"]==True]
+	matched = [channel for channel in template_order if template_dict[channel]["matched"]==True]
+	
 	
-	# 2-d (for the indication of missing template channel's position when fill_mode:'mean')
-	fig = [template_montage.plot(), input_montage.plot()]
-	fig[0].set_size_inches(5.6, 5.6)
-	fig[1].set_size_inches(5.6, 5.6)
+	# --------------------------------2-D------------------------------------
+	# (for the indication of missing template channel's position when fill_mode:'mean')
 	
+	fig = [template_montage.plot(), input_montage.plot()]	
 	ax = [fig[0].axes[0], fig[1].axes[0]]
-	ax[0].set_aspect('equal')
-	ax[1].set_aspect('equal')
-	ax[0].figure.canvas.draw() #update the figure
-	ax[1].figure.canvas.draw()
 	
 	# get the original coords
-	all_tpl = ax[0].transData.transform(ax[0].collections[0].get_offsets().data) # display coords (px)
-	all_in= ax[1].transData.transform(ax[1].collections[0].get_offsets().data)
+	#all_tpl = ax[0].transData.transform(ax[0].collections[0].get_offsets().data) # displayed coords (px)
+	#all_in= ax[1].transData.transform(ax[1].collections[0].get_offsets().data)
+	all_tpl = ax[0].collections[0].get_offsets().data
+	all_in= ax[1].collections[0].get_offsets().data
 	matched_tpl = np.array([all_tpl[template_dict[channel]["index"]] for channel in matched])
 	matched_in = np.array([all_in[input_dict[channel]["index"]] for channel in matched])
 	
-	# transform the xy axis (template's -> input's)
-	rbf_x = Rbf(matched_tpl[:,0], matched_tpl[:,1], matched_in[:,0], function='thin_plate')
-	rbf_y = Rbf(matched_tpl[:,0], matched_tpl[:,1], matched_in[:,1], function='thin_plate')
+	# transform the xy axis (input's -> template's)
+	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 to all template channels
-	transformed_tpl_x = rbf_x(all_tpl[:,0], all_tpl[:,1])
-	transformed_tpl_y = rbf_y(all_tpl[:,0], all_tpl[:,1])
-	#transformed_tpl = np.vstack((transformed_tpl_x, transformed_tpl_y)).T
+	# apply to all input 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
 	
-	# update input, template's position
+	# save template's and input's 2d position
 	for i, channel in enumerate(template_order):
-		css_left = (transformed_tpl_x[i]-11)/560
-		css_bottom = (transformed_tpl_y[i]-7)/560
-		template_dict[channel]["css_position"] = [str(round(css_left*100, 2))+"%", str(round(css_bottom*100, 2))+"%"]
+		template_dict[channel]["coord_2d"] = all_tpl[i]
 	for i, channel in enumerate(input_order):
-		css_left = (all_in[i][0]-11)/560
-		css_bottom = (all_in[i][1]-7)/560
-		input_dict[channel]["css_position"] = [str(round(css_left*100, 2))+"%", str(round(css_bottom*100, 2))+"%"]
+		input_dict[channel]["coord_2d"] = transformed_in[i].tolist()
 	
 	
 	# 3-d (to use KNN)
 	# get the original coords
-	all_tpl = np.array([template_dict[channel]["coord"].tolist() for channel in template_order])
-	all_in = np.array([input_dict[channel]["coord"].tolist() for channel in input_order])
+	all_tpl = np.array([template_dict[channel]["coord_3d"].tolist() for channel in template_order])
+	all_in = np.array([input_dict[channel]["coord_3d"].tolist() for channel in input_order])
 	matched_tpl = np.array([all_tpl[template_dict[channel]["index"]] for channel in matched])
 	matched_in = np.array([all_in[input_dict[channel]["index"]] for channel in matched])
 	
@@ -152,13 +144,13 @@ def align_coords(channel_info, template_montage, input_montage):
 	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
 	
-	# update input's position
+	# update input's 3d position
 	for i, channel in enumerate(input_order):
-		input_dict[channel]["coord"] = transformed_in[i].tolist()
+		input_dict[channel]["coord_3d"] = transformed_in[i].tolist()
 	
 	channel_info.update({
 	    "templateDict" : template_dict,
-	    "inputDict" : input_dict,
+	    "inputDict" : input_dict
 	})
 	return channel_info
 
@@ -168,13 +160,12 @@ def find_neighbors(app_state, channel_info):
 	input_dict = channel_info["inputDict"]
 	template_order = channel_info["templateOrder"]
 	input_order = channel_info["inputOrder"]
-	#z_row_idx = channel_info["dataShape"][0]
 	missing_channels = app_state["missingTemplates"]
 	if missing_channels == []:
 		return app_state # change nothing
 	
 	
-	in_coords = [input_dict[channel]["coord"] for channel in input_order]
+	in_coords = [input_dict[channel]["coord_3d"] for channel in input_order]
 	in_coords = np.array([in_coords[i] for i in range(len(in_coords))])
 	
 	# use KNN to choose k nearest channels
@@ -183,7 +174,7 @@ def find_neighbors(app_state, channel_info):
 	knn.fit(in_coords)
 	
 	for channel in missing_channels:
-		distances, indices = knn.kneighbors(np.array(template_dict[channel]["coord"]).reshape(1,-1))
+		distances, indices = knn.kneighbors(np.array(template_dict[channel]["coord_3d"]).reshape(1,-1))
 		selected = [input_order[i] for i in indices[0]]
 		#print(channel, ':', selected)
 		
@@ -205,51 +196,44 @@ def mapping_stage1(app_state, channel_info, loc_file):
 	template_order = template_montage.ch_names
 	input_order = input_montage.ch_names
 	new_idx = [[]]*30
-	alias = {
+	alias_dict = {
 		'T3': 'T7',
 		'T4': 'T8',
 		'T5': 'P7',
-		'T6': 'P8',
-		#'TP7': 'T5\'',
-		#'TP8': 'T6\'',
+		'T6': 'P8'
 	}
 	
 	# match the names of input channels -> template channels
 	for i, channel in enumerate(template_order):
-		if channel in alias and alias[channel] in input_dict:
-			template_montage.rename_channels({channel: alias[channel]})
-			template_dict[alias[channel]] = template_dict.pop(channel)
-			template_dict[alias[channel]].name = alias[channel]
-			channel = alias[channel]
+		if channel in alias_dict and alias_dict[channel] in input_dict:
+			template_montage.rename_channels({channel: alias_dict[channel]})
+			template_dict[alias_dict[channel]] = template_dict.pop(channel)
+			channel = alias_dict[channel]
 		
 		if channel in input_dict:
-			new_idx[i] = [input_dict[channel].index]
-			
-			template_dict[channel].matched = True
-			input_dict[channel].matched = True
-			input_dict[channel].assigned = True
+			new_idx[i] = [input_dict[channel]["index"]]
+			template_dict[channel]["matched"] = True
+			input_dict[channel]["assigned"] = True
 	
 	# update names
 	template_order = template_montage.ch_names
 	input_order = input_montage.ch_names
 	
 	channel_info.update({
-	    "templateDict" : {k : v.__dict__ for k,v in template_dict.items()},
-	    "inputDict" : {k : v.__dict__ for k,v in input_dict.items()},
+	    "templateDict" : template_dict,
+	    "inputDict" : input_dict,
 	    "templateOrder" : template_order,
 		"inputOrder" : input_order
 	})
 	app_state.update({
 	    "stage1NewOrder" : new_idx,
 	    "runningState" : "stage1",
-	    "stage1UnassignedInputs" : [channel for channel in input_order if input_dict[channel].assigned==False],
-	    "missingTemplates" : [channel for channel in template_order if template_dict[channel].matched==False]
+	    "stage1UnassignedInputs" : [channel for channel in input_order if input_dict[channel]["assigned"]==False],
+	    "missingTemplates" : [channel for channel in template_order if template_dict[channel]["matched"]==False]
 	})
 	
 	# align input, template's coordinates
 	channel_info = align_coords(channel_info, template_montage, input_montage)
-	# fill the unmatched channels
-	#app_state = fill_channels(app_state, channel_info, fill_mode)
 	
 	second2 = time.time()
 	print('Mapping (stage1) finished in',second2 - second1,'s.')
@@ -267,8 +251,8 @@ def mapping_stage2(app_state, channel_info):
 		app_state["runningState"] = "finished"
 		return app_state, channel_info
 	
-	tpl_coords = np.array([template_dict[channel]["coord"] for channel in template_order])
-	unassigned_coords = np.array([input_dict[channel]["coord"] for channel in unassigned])
+	tpl_coords = np.array([template_dict[channel]["coord_3d"] for channel in template_order])
+	unassigned_coords = np.array([input_dict[channel]["coord_3d"] for channel in unassigned])
 	
 	# reset all tpl.matched to False
 	for channel in template_dict:
@@ -276,7 +260,7 @@ def mapping_stage2(app_state, channel_info):
 	
 	# initialize the cost matrix
 	if len(unassigned) < 30:
-		cost_matrix = np.full((30, 30), 1e6) # add dummy channels to ensure num_col > num_row
+		cost_matrix = np.full((30, 30), 1e6) # add dummy channels to ensure num_col >= num_row
 	else:
 		cost_matrix = np.zeros((30, len(unassigned)))
 	for i in range(30):
@@ -287,17 +271,16 @@ def mapping_stage2(app_state, channel_info):
 	# use Hungarian Algorithm to find the minimum sum of distance of (input's coord to template's coord)...?
 	row_idx, col_idx = linear_sum_assignment(cost_matrix)
 	
-	matches = []
 	new_idx = [[]]*30
 	for i in range(30):
 		if col_idx[i] < len(unassigned): # filter out dummy channels
-			matches.append([row_idx[i], col_idx[i]])
+			print(f'({row_idx[i]}, {col_idx[i]})')
 			
 			tpl_channel = template_order[row_idx[i]]
 			in_channel = unassigned[col_idx[i]]
 			template_dict[tpl_channel]["matched"] = True
 			input_dict[in_channel]["assigned"] = True
-			new_idx[i] = [input_dict[in_channel]["index"]]
+			new_idx[row_idx[i]] = [input_dict[in_channel]["index"]]
 			
 			print(template_order[row_idx[i]], '<-', unassigned[col_idx[i]])
 	
@@ -312,7 +295,7 @@ def mapping_stage2(app_state, channel_info):
 	    "missingTemplates" : [channel for channel in template_order if template_dict[channel]["matched"]==False]
 	})
 	
-	# fill the missing_channels channels
+	# fill the missing_channels
 	app_state = find_neighbors(app_state, channel_info)
 	
 	second2 = time.time()