update

app.py

@@ -44,14 +44,24 @@ icunet = """
 Electroencephalography (EEG) signals are often contaminated with artifacts. It is imperative to develop a practical and reliable artifact removal method to prevent the misinterpretation of neural signals and the underperformance of brain–computer interfaces. Based on the U-Net architecture, we developed a new artifact removal model, IC-U-Net, for removing pervasive EEG artifacts and reconstructing brain signals. IC-U-Net was trained using mixtures of brain and non-brain components decomposed by independent component analysis. It uses an ensemble of loss functions to model complex signal fluctuations in EEG recordings. The effectiveness of the proposed method in recovering brain activities and removing various artifacts (e.g., eye blinks/movements, muscle activities, and line/channel noise) was demonstrated in a simulation study and four real-world EEG experiments. IC-U-Net can reconstruct a multi-channel EEG signal and is applicable to most artifact types, offering a promising end-to-end solution for automatically removing artifacts from EEG recordings. It also meets the increasing need to image natural brain dynamics in a mobile setting.
 """
 
-chkbox_js = """
+init_js = """
 (app_state, channel_info) => {
 	app_state = JSON.parse(JSON.stringify(app_state));
 	channel_info = JSON.parse(JSON.stringify(channel_info));
-	if(app_state.state == "finished") return;
+	
+	let selector, attribute;
+	let channel, left, bottom;
+	
+	if(app_state.state == "step2-selecting"){
+	    selector = "#radio> div:nth-of-type(2)";
+	    attribute = "value";
+	}else if(app_state.state == "step3-selecting"){
+	    selector = "#chkbox-group> div:nth-of-type(2)";
+	    attribute = "name";
+	}else return;
 	
 	// add figure of in_montage
-	document.querySelector("#chkbox-group> div:nth-of-type(2)").style.cssText = `
+	document.querySelector(selector).style.cssText = `
 		position: relative;
 		width: 560px;
 		height: 560px;
@@ -59,13 +69,31 @@ chkbox_js = """
 	`;
 	
 	
+	// move the radios/checkboxes
+	let all_elem = document.querySelectorAll(selector+"> label");
+	Array.from(all_elem).forEach((item) => {
+		channel = item.querySelector("input").getAttribute(attribute);
+		left = channel_info.inputByName[channel].css_position[0];
+		bottom = channel_info.inputByName[channel].css_position[1];
+		//console.log(`channel: ${channel}, left: ${left}, bottom: ${bottom}`);
+		
+		item.style.cssText = `
+			position: absolute;
+			left: ${left};
+			bottom: ${bottom};
+		`;
+		item.className = "";
+		item.querySelector(":scope> span").innerText = "";
+	});
+	
+	
 	// add indication for the missing channels
-	let channel = channel_info.missingChannels[0]
-	let left = channel_info.templateByName[channel].css_position[0];
-	let bottom = channel_info.templateByName[channel].css_position[1];
+	channel = app_state.missingTemplates[0]
+	left = channel_info.templateByName[channel].css_position[0];
+	bottom = channel_info.templateByName[channel].css_position[1];
 	
 	let rule = `
-		#chkbox-group> div:nth-of-type(2)::after{
+		${selector}::after{
 			content: '';
 			position: absolute;
 			background-color: red;
@@ -81,35 +109,15 @@ chkbox_js = """
 	let exist = 0;
 	const styleSheet = document.styleSheets[0];
 	for(let i=0; i<styleSheet.cssRules.length; i++){
-		if(styleSheet.cssRules[i].selectorText == "#chkbox-group> div:nth-of-type(2)::after"){
+		if(styleSheet.cssRules[i].selectorText == selector+"::after"){
 			exist = 1;
-			console.log('exist!');
+			//console.log('exist!');
 			styleSheet.deleteRule(i);
 			styleSheet.insertRule(rule, styleSheet.cssRules.length);
 			break;
 		}
 	}
 	if(exist == 0) styleSheet.insertRule(rule, styleSheet.cssRules.length);
-
-
-	// move the checkboxes
-	let all_chkbox = document.querySelectorAll("#chkbox-group> div:nth-of-type(2)> label");
-	//all_chkbox = Array.apply(null, all_chkbox);
-	
-	Array.from(all_chkbox).forEach((item, index) => {
-		channel = channel_info.inputByIndex[index];
-		left = channel_info.inputByName[channel].css_position[0];
-		bottom = channel_info.inputByName[channel].css_position[1];
-		console.log(`left: ${left}, bottom: ${bottom}`);
-		
-		item.style.cssText = `
-			position: absolute;
-			left: ${left};
-			bottom: ${bottom};
-		`;
-		item.className = "";
-		item.querySelector(":scope> span").innerText = "";
-	});
 }
 """
 
@@ -117,14 +125,40 @@ indication_js = """
 (app_state, channel_info) => {
 	app_state = JSON.parse(JSON.stringify(app_state));
 	channel_info = JSON.parse(JSON.stringify(channel_info));
-	if(app_state.state == "finished") return;
 	
-	let channel = channel_info.missingChannels[app_state["fillingCount"]-1]
-	let left = channel_info.templateByName[channel].css_position[0];
-	let bottom = channel_info.templateByName[channel].css_position[1];
+	let selector;
+	let channel, left, bottom;
+	
+	if(app_state.state == "step2-selecting"){
+	    selector = "#radio> div:nth-of-type(2)";
+	    
+	    // update the radios
+		let all_elem = document.querySelectorAll(selector+"> label");
+		Array.from(all_elem).forEach((item) => {
+			channel = item.querySelector("input").value;
+			left = channel_info.inputByName[channel].css_position[0];
+			bottom = channel_info.inputByName[channel].css_position[1];
+			//console.log(`channel: ${channel}, left: ${left}, bottom: ${bottom}`);
+			
+			item.style.cssText = `
+				position: absolute;
+				left: ${left};
+				bottom: ${bottom};
+			`;
+			item.className = "";
+			item.querySelector(":scope> span").innerText = "";
+		});
+	}else if(app_state.state == "step3-selecting"){
+	    selector = "#chkbox-group> div:nth-of-type(2)";
+	}else return;
+	
+	// update indication
+	channel = app_state.missingTemplates[app_state["fillingCount"]-1]
+	left = channel_info.templateByName[channel].css_position[0];
+	bottom = channel_info.templateByName[channel].css_position[1];
 	
 	let rule = `
-		#chkbox-group> div:nth-of-type(2)::after{
+		${selector}::after{
 			content: '';
 			position: absolute;
 			background-color: red;
@@ -140,9 +174,9 @@ indication_js = """
 	let exist = 0;
 	const styleSheet = document.styleSheets[0];
 	for(let i=0; i<styleSheet.cssRules.length; i++){
-		if(styleSheet.cssRules[i].selectorText == "#chkbox-group> div:nth-of-type(2)::after"){
+		if(styleSheet.cssRules[i].selectorText == selector+"::after"){
 			exist = 1;
-			console.log('exist!');
+			//console.log('exist!');
 			styleSheet.deleteRule(i);
 			styleSheet.insertRule(rule, styleSheet.cssRules.length);
 			break;
@@ -310,8 +344,8 @@ with gr.Blocks() as demo:
 		#### determine the next step ####
 		#################################
 		
-		in_num = len(channel_info["inputByName"])
-		matched_num = 30 - len(channel_info["missingChannels"])
+		in_num = len(channel_info["inputByIndex"])
+		matched_num = 30 - len(app_state["missingTemplates"])
 		
 		# if the input channels(>=30) has all the 30 template channels
 		# -> Stage2.decode data
@@ -360,23 +394,20 @@ with gr.Blocks() as demo:
 
 	def init_next_step(app_state, channel_info, selected_radio, selected_chkbox):
 		
-		channel_info["missingChannels"] = [channel for channel in channel_info["templateByName"]
-											if channel_info["templateByName"][channel]["matched"]==False]
-		app_state.update({
-			"fillingCount" : 1,
-			"totalFillingNum" : len(channel_info["missingChannels"])
-		})
-		
 		# step1 -> step2
 		if app_state["state"] == "step2-initializing":
-			app_state["state"] = "step2-selecting"
+			print('step1 -> step2')
+			app_state["missingTemplates"] = [channel for channel in channel_info["templateByIndex"]
+											if channel_info["templateByName"][channel]["matched"]==False]
+			app_state.update({
+				"state" : "step2-selecting",
+				"fillingCount" : 1,
+				"totalFillingNum" : len(app_state["missingTemplates"])
+			})
 			
-			name = channel_info["missingChannels"][0]
+			name = app_state["missingTemplates"][0]
 			label = name+' (1/'+str(app_state["totalFillingNum"])+')'
-			
-			unassigned = [channel for channel in channel_info["inputByName"]
-							if channel_info["inputByName"][channel]["assigned"]==False]
-			if len(unassigned)==1 or app_state["totalFillingNum"]==1:
+			if len(app_state["stage1UnassignedInputs"])==1 or app_state["totalFillingNum"]==1:
 				btn_label = "Next step"
 			else:
 				btn_label = "Next"
@@ -386,14 +417,22 @@ with gr.Blocks() as demo:
 					desc_md : gr.Markdown("### step2"),
 					tpl_montage : gr.Image(visible=False),
 					map_montage : gr.Image(visible=False),
-					radio : gr.Radio(choices=unassigned, value=[], label=label, visible=True),
+					radio : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
 					step2_btn : gr.Button(btn_label, visible=True),
 					next_btn : gr.Button(visible=False)}
 		
 		# step1 -> step3
 		elif app_state["state"] == "step3-initializing":
-			
+			print('step1 -> step3')
+			app_state["missingTemplates"] = [channel for channel in channel_info["templateByIndex"]
+											if channel_info["templateByName"][channel]["matched"]==False]
+			app_state.update({
+				"state" : "step3-initializing",
+				"fillingCount" : 1,
+				"totalFillingNum" : len(app_state["missingTemplates"])
+			})
 			return {app_state_json : app_state,
+			        channel_info_json : channel_info,
 					desc_md : gr.Markdown("### step3"),
 					tpl_montage : gr.Image(visible=False),
 					map_montage : gr.Image(visible=False),
@@ -405,22 +444,31 @@ with gr.Blocks() as demo:
 		elif app_state["state"] == "step2-selecting":
 		
 			# save info before clicking on next_btn
-			if selected_radio != []:
-				prev_target_name = channel_info["missingChannels"][app_state["fillingCount"]-1]
-				prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
-				
+			prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
+			prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
+			if selected_radio == []:
+				app_state["stage1NewOrder"][prev_target_idx] = []
+			else:
 				selected_idx = channel_info["inputByName"][selected_radio]["index"]
 				app_state["stage1NewOrder"][prev_target_idx] = [selected_idx]
 			
 				channel_info["templateByName"][prev_target_name]["matched"] = True
-				channel_info["inputByName"][selected_radio]["assigned"] = True
-				print(prev_target_name, '<-', selected_radio)
+				channel_info["inputByName"][selected_radio]["assigned"] = True				
+			print(prev_target_name, '<-', selected_radio)
+			
+			app_state.update({
+				"stage1UnassignedInputs" : [channel for channel in channel_info["inputByIndex"]
+									if channel_info["inputByName"][channel]["assigned"]==False],
+				"missingTemplates" : [channel for channel in channel_info["templateByIndex"]
+									if channel_info["templateByName"][channel]["matched"]==False]
+			})
 			
 			# if all the unmatched template channels were filled by input channels
 			# -> Stage2
-			if app_state["fillingCount"] == app_state["totalFillingNum"]:
-				app_state["state"] = "finished"
+			if len(app_state["missingTemplates"]) == 0:
+				print('step2 -> Stage2')
 				gr.Info('The mapping process is finished!')
+				app_state["state"] = "finished"
 				
 				return {app_state_json : app_state,
 						channel_info_json : channel_info,
@@ -431,8 +479,12 @@ with gr.Blocks() as demo:
 						
 			# -> step3
 			else:
-				app_state["state"] = "step3-initializing"
-				
+				print('step2 -> step3')
+				app_state.update({
+					"state" : "step3-initializing",
+					"fillingCount" : 1,
+					"totalFillingNum" : len(app_state["missingTemplates"])
+				})
 				return {app_state_json : app_state,
 						channel_info_json : channel_info,
 						desc_md : gr.Markdown("### step3"),
@@ -443,18 +495,23 @@ with gr.Blocks() as demo:
 		
 		# step3 -> Stage2.decode data
 		elif app_state["state"] == "step3-selecting":
-			app_state["state"] = "finished"
-			gr.Info('The mapping process is finished!')
 			
 			# save info before clicking on next_btn
-			if selected_chkbox != []:
-				prev_target_name = channel_info["missingChannels"][app_state["fillingCount"]-1]
-				prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
-				
+			prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
+			prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
+			if selected_chkbox == []:
+				app_state["stage1NewOrder"][prev_target_idx] = []
+			else:
 				selected_idx = [channel_info["inputByName"][channel]["index"] for channel in selected_chkbox]
-				app_state["stage1NewOrder"][prev_target_idx] = selected_idx
-				
-				#print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
+				app_state["stage1NewOrder"][prev_target_idx] = selected_idx				
+			#print(f'{prev_target_name}({prev_target_idx}): {selected_chkbox}')
+			
+			gr.Info('The mapping process is finished!')
+			app_state["state"] = "finished"
+			print('step3 -> Stage2')
+			
+			app_state["missingTemplates"] = [channel for channel in channel_info["templateByIndex"]
+											if channel_info["templateByName"][channel]["matched"]==False]
 			
 			return {app_state_json : app_state,
 					desc_md : gr.Markdown(visible=False),
@@ -469,7 +526,7 @@ with gr.Blocks() as demo:
 					chkbox_group, fillmode_btn, step2_btn, next_btn, run_btn]
 	).success(
 	    fn = None,
-	    js = indication_js,
+	    js = init_js,
 	    inputs = [app_state_json, channel_info_json],
 		outputs = []
 	)
@@ -479,36 +536,36 @@ with gr.Blocks() as demo:
 	def update_radio(app_state, channel_info, selected):
 		
 		# save info before clicking on next_btn
-		if selected != []:
-			prev_target_name = channel_info["missingChannels"][app_state["fillingCount"]-1]
-			prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
-			
+		prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
+		prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
+		if selected == []:
+			app_state["stage1NewOrder"][prev_target_idx] = []
+		else:
 			selected_idx = channel_info["inputByName"][selected]["index"]
 			app_state["stage1NewOrder"][prev_target_idx] = [selected_idx]
 		
 			channel_info["templateByName"][prev_target_name]["matched"] = True
 			channel_info["inputByName"][selected]["assigned"] = True
-			print(prev_target_name, '<-', selected)
+		print(prev_target_name, '<-', selected)
 		
 		# update the current round
 		app_state["fillingCount"] += 1
-		unassigned = [channel for channel in channel_info["inputByName"]
-						if channel_info["inputByName"][channel]["assigned"]==False]
+		app_state["stage1UnassignedInputs"] = [channel for channel in channel_info["inputByIndex"]
+												if channel_info["inputByName"][channel]["assigned"]==False]
 		
-		target_name = channel_info["missingChannels"][app_state["fillingCount"]-1]
-				
+		target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
 		radio_label = target_name+' ('+str(app_state["fillingCount"])+'/'+str(app_state["totalFillingNum"])+')'
 				
-		if len(unassigned)==1 or app_state["fillingCount"]==app_state["totalFillingNum"]:
+		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=unassigned, value=[], label=radio_label),
+					radio : 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=unassigned, value=[], label=radio_label)}
+					radio : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=radio_label)}
 	
 	step2_btn.click(
 	    fn = update_radio,
@@ -523,8 +580,6 @@ with gr.Blocks() as demo:
 	)
 	
 	# step3
-	@fillmode_btn.click(inputs = [app_state_json, channel_info_json, in_fill_mode],
-						outputs = [app_state_json, in_fill_mode, fillmode_btn, chkbox_group, step3_btn, run_btn])
 	def fill_value(app_state, channel_info, fill_mode):
 		
 		if fill_mode == 'zero':
@@ -541,7 +596,7 @@ with gr.Blocks() as demo:
 			app_state["state"] = "step3-selecting"
 			app_state = find_neighbors(app_state, channel_info, fill_mode)
 			
-			name = channel_info["missingChannels"][0]
+			name = app_state["missingTemplates"][0]
 			idx = channel_info["templateByName"][name]["index"]
 			value = app_state["stage1NewOrder"][idx]
 			value = [channel_info["inputByIndex"][i] for i in value]
@@ -554,22 +609,33 @@ with gr.Blocks() as demo:
 													value=value, label=label, visible=True),
 		            step3_btn : gr.Button(visible=True)}
 	
+	fillmode_btn.click(
+	    fn = fill_value,
+	    inputs = [app_state_json, channel_info_json, in_fill_mode],
+		outputs = [app_state_json, in_fill_mode, fillmode_btn, chkbox_group, step3_btn, run_btn]
+	).success(
+	    fn = None,
+	    js = init_js,
+	    inputs = [app_state_json, channel_info_json],
+		outputs = []
+	)
+	
 	def update_chkbox(app_state, channel_info, selected):
 	
 		# save info before clicking on next_btn
-		if selected != []:
-			prev_target_name = channel_info["missingChannels"][app_state["fillingCount"]-1]
-			prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
-			
+		prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
+		prev_target_idx = channel_info["templateByName"][prev_target_name]["index"]
+		if selected == []:
+			app_state["stage1NewOrder"][prev_target_idx] = []
+		else:
 			selected_idx = [channel_info["inputByName"][channel]["index"] for channel in selected]
 			app_state["stage1NewOrder"][prev_target_idx] = selected_idx
-			
-			#print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
+		#print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
 		
-		# update next round
+		# update the current round
 		app_state["fillingCount"] += 1
 		
-		target_name = channel_info["missingChannels"][app_state["fillingCount"]-1]
+		target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
 		target_idx = channel_info["templateByName"][target_name]["index"]
 		
 		chkbox_value = app_state["stage1NewOrder"][target_idx]
@@ -606,11 +672,14 @@ with gr.Blocks() as demo:
 	
 	def reset_run(app_state, channel_info, raw_data, model_name):
 		filepath = app_state["filepath"]
+		delete_file(filepath+'mapped.csv')
+		delete_file(filepath+'denoised.csv')
+		
 		input_name = os.path.basename(str(raw_data))
 		output_name = os.path.splitext(input_name)[0]+'_'+model_name+'.csv'
 		
-		in_num = len(channel_info["inputByName"])
-		matched_num = len([channel for channel in channel_info["inputByName"]
+		in_num = len(channel_info["inputByIndex"])
+		matched_num = len([channel for channel in channel_info["inputByIndex"]
 							if channel_info["inputByName"][channel]["matched"]==True])
 		batch_num = math.ceil((in_num-matched_num)/30) + 1
 		
@@ -619,12 +688,16 @@ with gr.Blocks() as demo:
 			"runnigState" : "stage1",
 			"batchCount" : 1,
 			"totalBatchNum" : batch_num,
-			"stage2NewOrder" : [[]]*30
+			"stage2UnassignedInputs" : app_state["stage1UnassignedInputs"],
+			"stage2NewOrder" : [[]]*30,
 		})
 		
-		delete_file(filepath+'mapped.csv')
-		delete_file(filepath+'denoised.csv')
+		# reset in.assigned back to the state after Stage1
+		for channel in app_state["stage1UnassignedInputs"]:
+			channel_info["inputByName"][channel]["assigned"] = False
+		
 		return {app_state_json : app_state,
+				channel_info_json : channel_info,
 				run_btn : gr.Button(interactive=False),
 				batch_md : gr.Markdown(visible=False),
 				out_denoised_data : gr.File(visible=False)}
@@ -644,6 +717,7 @@ with gr.Blocks() as demo:
 			if app_state["batchCount"] > 1:
 				app_state, channel_info = mapping_stage2(app_state, channel_info, fill_mode)
 				if app_state["runnigState"] == "finished":
+					#yield {batch_md : gr.Markdown("error", visible=True)}
 					break
 			
 			reorder_to_template(app_state, raw_data)
@@ -669,7 +743,7 @@ with gr.Blocks() as demo:
 	run_btn.click(
 		fn = reset_run,
 		inputs = [app_state_json, channel_info_json, in_raw_data, in_model_name],
-		outputs = [app_state_json, run_btn, batch_md, out_denoised_data]
+		outputs = [app_state_json, channel_info_json, run_btn, batch_md, out_denoised_data]
 		
 	).success(
 		fn = run_model,

channel_mapping.py

@@ -15,8 +15,8 @@ def reorder_to_template(app_state, filename):
 	old_data = utils.read_train_data(filename) # original raw data
 	new_data = np.zeros((30, old_data.shape[1])) # reordered raw data
 	new_filename = app_state["filepath"]+'mapped.csv'
-	print('new order 1:', app_state["stage1NewOrder"])
-	print('new order 2:', app_state["stage2NewOrder"])
+	#print('new order 1:', app_state["stage1NewOrder"])
+	#print('new order 2:', app_state["stage2NewOrder"])
 	
 	zero_arr = np.zeros((1, old_data.shape[1]))
 	old_data = np.concatenate((old_data, zero_arr), axis=0)
@@ -41,7 +41,7 @@ def reorder_to_origin(app_state, channel_info, filename, new_filename):
 	template_order = channel_info["templateByIndex"]
 	
 	if app_state["runnigState"] == "stage1":
-		new_data = np.zeros((len(channel_info["inputByName"]), old_data.shape[1]))
+		new_data = np.zeros((len(channel_info["inputByIndex"]), old_data.shape[1]))
 	else:
 		new_data = utils.read_train_data(new_filename)
 	
@@ -94,7 +94,7 @@ def align_coords(channel_info, template_montage, input_montage):
 	input_dict = channel_info["inputByName"]
 	template_order = channel_info["templateByIndex"]
 	input_order = channel_info["inputByIndex"]
-	matched = [channel for channel in input_dict if input_dict[channel]["matched"]==True]
+	matched = [channel for channel in input_order if input_dict[channel]["matched"]==True]
 	
 	# 2-d (for the indication of missing template channel's position when fill_mode:'mean_manual')
 	fig = [template_montage.plot(), input_montage.plot()]
@@ -168,8 +168,8 @@ def find_neighbors(app_state, channel_info, fill_mode):
 	template_order = channel_info["templateByIndex"]
 	input_order = channel_info["inputByIndex"]
 	#z_row_idx = channel_info["dataShape"][0]
-	unmatched = [channel for channel in template_dict if template_dict[channel]["matched"]==False]
-	if unmatched == []:
+	missing_channels = app_state["missingTemplates"]
+	if missing_channels == []:
 		return app_state # change mothing
 	
 	
@@ -177,11 +177,11 @@ def find_neighbors(app_state, channel_info, fill_mode):
 	in_coords = np.array([in_coords[i] for i in range(len(in_coords))])
 	
 	# use KNN to choose k nearest channels
-	k = 4 if len(input_dict)>4 else len(input_dict)
+	k = 4 if len(input_order)>4 else len(input_order)
 	knn = NearestNeighbors(n_neighbors=k, metric='euclidean')
 	knn.fit(in_coords)
 	
-	for channel in unmatched:
+	for channel in missing_channels:
 		distances, indices = knn.kneighbors(np.array(template_dict[channel]["coord"]).reshape(1,-1))
 		selected = [input_order[i] for i in indices[0]]
 		print(channel, ':', selected)
@@ -202,8 +202,8 @@ def mapping_stage1(app_state, channel_info, loc_file):
 	
 	template_montage, input_montage, template_dict, input_dict = read_montage_data(loc_file)
 	template_order = template_montage.ch_names
+	input_order = input_montage.ch_names
 	new_idx = [[]]*30
-	missing_channels = []
 	alias = {
 		'T3': 'T7',
 		'T4': 'T8',
@@ -227,19 +227,18 @@ def mapping_stage1(app_state, channel_info, loc_file):
 			template_dict[channel].matched = True
 			input_dict[channel].matched = True
 			input_dict[channel].assigned = True
-		else:
-			missing_channels.append(channel)
 	
 	channel_info.update({
-	    "missingChannels" : missing_channels,
 	    "templateByName" : {k : v.__dict__ for k,v in template_dict.items()},
 	    "inputByName" : {k : v.__dict__ for k,v in input_dict.items()},
-	    "templateByIndex" : template_montage.ch_names,
-		"inputByIndex" : input_montage.ch_names
+	    "templateByIndex" : template_order,
+		"inputByIndex" : input_order
 	})
 	app_state.update({
 	    "stage1NewOrder" : new_idx,
-	    "runnigState" : "stage1"
+	    "runnigState" : "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]
 	})
 	
 	# align input, template's coordinates
@@ -257,7 +256,8 @@ def mapping_stage2(app_state, channel_info, fill_mode):
 	template_dict = channel_info["templateByName"]
 	input_dict = channel_info["inputByName"]
 	template_order = channel_info["templateByIndex"]
-	unassigned = [channel for channel in input_dict if input_dict[channel]["assigned"]==False]
+	input_order = channel_info["inputByIndex"]
+	unassigned = app_state["stage2UnassignedInputs"]
 	if unassigned == []:
 		app_state["runnigState"] = "finished"
 		return app_state, channel_info
@@ -265,7 +265,7 @@ def mapping_stage2(app_state, channel_info, fill_mode):
 	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])
 	
-	# set all tpl.matched to False
+	# reset all tpl.matched to False
 	for channel in template_dict:
 		template_dict[channel]["matched"] = False
 	
@@ -302,10 +302,12 @@ def mapping_stage2(app_state, channel_info, fill_mode):
 	})
 	app_state.update({
 	    "stage2NewOrder" : new_idx,
-	    "runnigState" : "stage2"
+	    "runnigState" : "stage2",
+	    "stage2UnassignedInputs" : [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]
 	})
 	
-	# fill the unmatched channels
+	# fill the missing_channels channels
 	app_state = find_neighbors(app_state, channel_info, fill_mode)
 	
 	second2 = time.time()