update v4

app.py

@@ -1,16 +1,8 @@
+import utils
+import app_utils
 import gradio as gr
-
 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_input_data, restore_original_order, find_neighbors
-
 
 readme = """
 
@@ -35,8 +27,14 @@ Your unmatched channels, previously highlighted in red, will be shown on your mo
 ### Step3: Filling Remaining Template Channels
 To run the models successfully, we need to ensure that all 30 template channels are filled. In this step, you are required to select one of the methods provided below to fill the remaining empty template channels:
 - **Mean** method: Each empty template channel is filled with the average value of data from the nearest input channels. By default, the 4 closest input channels (determined after aligning your montage to the template's scale using TPS) are selected for this averaging process. On the interface, you will see checkboxes displayed above each of your channel. The 4 nearest channels are pre-selected by default for each empty template channels, but you can modify these selections as needed. If you uncheck all the checkboxes for a particular template channel, it will be filled with zeros.
-- **Zero** method: All empty template channels are filled with zeros.
-Choose he method that best suits your needs, considering that the model's performance may vary depending on the method used. Once all template channels are filled, you can proceed to run the models.
+- **Zero** method: All empty template channels are filled with zeros.  
+Choose the method that best suits your needs, considering that the model's performance may vary depending on the method used.
+
+### Step4: Auto-mapping Remaining Channels
+After completing the initial mapping steps, any channels that are not yet assigned to a template will be processed in this step. These remaining channels will be automatically mapped in batches, with a batch size of up to 30 channels. If the final batch contains fewer than 30 channels, the **Mean** method from Step3 will be applied to fill the remaining template channels.
+
+
+### Mapping Result
 
 ## 2. Decode data
 In this phase, you can select which model to use for denoising your EEG data. Detailed information about the models can be found in the other tabs.
@@ -64,7 +62,7 @@ init_js = """
 	    selector = "#radio-group > div:nth-of-type(2)";
 	    //classname = "radio";
 	    attribute = "value";
-	}else if(stage1_info.state == "step3-selecting"){
+	}else if(stage1_info.state == "step3-2-selecting"){
 	    selector = "#chkbox-group > div:nth-of-type(2)";
 	     //classname = "chkbox";
 	    attribute = "name";
@@ -78,7 +76,7 @@ init_js = """
 		aspect-ratio: 1;
 		//width: 560px;
 		//height: 560px;
-		background: url("file=${stage1_info.filenames.input_montage}");
+		background: url("file=${stage1_info.fileNames.input_montage}");
 		background-size: contain;
 
 	`;
@@ -165,7 +163,7 @@ update_js = """
 			item.className = "";
 			item.querySelector(":scope > span").innerText = "";
 		});
-	}else if(stage1_info.state == "step3-selecting"){
+	}else if(stage1_info.state == "step3-2-selecting"){
 	    selector = "#chkbox-group > div:nth-of-type(2)";
 	}else return;
 	
@@ -247,15 +245,14 @@ with gr.Blocks() as demo:
 				map_btn = gr.Button("Mapping", interactive=False, scale=1)
 			
 			# ------------------------mapping------------------------
-			# description for stage1-123
 			desc_md = gr.Markdown(visible=False)
-			# stage1-1 : mapping result
+			# step1 : initial mapping abd rescaling
 			with gr.Row():
 				tpl_img = gr.Image("./template_montage.png", label="Template channels", visible=False)
 				mapped_img = gr.Image(label="Input channels", visible=False)
-			# stage1-2 : assign unmatched input channels to empty template channels
+			# step2 : forward unmatched input channels to empty template channels
 			radio_group = gr.Radio(elem_id="radio-group", visible=False)
-			# stage1-3 : select a way to fill the empty template channels
+			# step3 : fill the remaining template channels
 			with gr.Row():
 				in_fillmode = gr.Dropdown(choices=["mean", "zero"],
 									value="mean",
@@ -264,6 +261,8 @@ with gr.Blocks() as demo:
 									scale=2)
 				fillmode_btn = gr.Button("OK", visible=False, scale=1)
 			chkbox_group = gr.CheckboxGroup(elem_id="chkbox-group", visible=False)
+			# step4 : mapping result
+			out_json_file = gr.File(label="Mapping result", visible=False)
 			
 			with gr.Row():
 				clear_btn = gr.Button("Clear", visible=False)
@@ -280,7 +279,9 @@ with gr.Blocks() as demo:
 									("ART", "EEGART"),
 									("IC-U-Net", "ICUNet"),
 									("IC-U-Net++", "UNetpp"),
-									("IC-U-Net-Attn", "AttUnet")],
+									("IC-U-Net-Attn", "AttUnet"),
+									"(mapped data)",
+									"(denoised data)"],
 									value="EEGART",
 									label="Model",
 									scale=2)
@@ -303,7 +304,7 @@ with gr.Blocks() as demo:
 		with gr.Tab("README"):
 			gr.Markdown(readme)
     
-	#demo.load(js=tmp_js)
+	#demo.load(js=js)
 	
 	# verify that all required inputs have been provided
 	@gr.on(triggers = [in_data_file.upload, in_data_file.clear, in_loc_file.upload, in_loc_file.clear, in_samplerate.change],
@@ -316,317 +317,233 @@ with gr.Blocks() as demo:
 	
 	
 	# +========================================================================================+
-	# |                                stage1: channel mapping                                 |
+	# |                                Stage1: channel mapping                                 |
 	# +========================================================================================+
 	def reset_all(in_data, in_loc, samplerate):
 		# establish a new folder for the current session
-		filepath = os.path.dirname(str(in_data))
+		rootpath = os.path.dirname(str(in_data))
 		try:
-			os.mkdir(filepath+"/session_data/")
+			os.mkdir(rootpath+"/session_data/")
 		except OSError as e:
-			utils.dataDelete(filepath+"/session_data/")
-			os.mkdir(filepath+"/session_data/")
+			utils.dataDelete(rootpath+"/session_data/")
+			os.mkdir(rootpath+"/session_data/")
 			print(e)
 		# establish new folders for stage1 and stage2
-		os.mkdir(filepath+"/session_data/stage1/")
-		os.mkdir(filepath+"/session_data/stage2/")
+		os.mkdir(rootpath+"/session_data/stage1/")
+		os.mkdir(rootpath+"/session_data/stage2/")
 		
 		# initialize channel_info, app_info
 		channel_info = {}
 		app_info = {
-			"rootFilepath" : filepath+"/session_data/",
+			"rootPath" : rootpath+"/session_data/",
 			"sampleRate" : int(samplerate),
-			#"currentStage" : "stage1",
 			"stage1" : {
-				"filepath" : filepath+"/session_data/stage1/",
-				"filenames" : {
+				"filePath" : rootpath+"/session_data/stage1/",
+				"fileNames" : {
 					"input_data" : in_data,
 					"input_loc" : in_loc,
 					"input_montage" : "",
 		    		"mapped_montage" : ""
 				},
-				"state" : None,
+				"state" : "step1-initializing",
 				"fillingCount" : None,
 				"totalFillingNum" : None,
-				"newOrder" : None,
 	    		"unassignedInputs" : None,
-	    		"missingTemplates" : None
+	    		"missingTemplates" : None,
+	    		"mappingData" : [
+	    		{
+					"newOrder" : None,
+					"fillFlags" : None,
+					#"channelUsageNum" : None
+				}
+	    		]
 			},
 			"stage2" : {
-				"filepath" : filepath+"/session_data/stage2/",
-				"filenames" : {
+				"filePath" : rootpath+"/session_data/stage2/",
+				"fileNames" : {
 					"output_data" : ""
 				},
-				#"state" : None,
-				"totalBatchNum" : None,
-				"newOrder" : None,
-	    		"unassignedInputs" : None
+				"totalBatchNum" : None
 			}
 		}
 		# reset layout
 		return {app_info_json : app_info,
 				channel_info_json : channel_info,
-				# ------------------stage1-----------------------
+				# --------------------Stage1-------------------------
 				map_btn : gr.Button(interactive=False),
 				desc_md : gr.Markdown(visible=False),
+				next_btn : gr.Button(visible=False),
 				tpl_img : gr.Image(visible=False),
 				mapped_img : gr.Image(value=None, visible=False),
 				radio_group : gr.Radio(choices=[], value=[], label="", visible=False),
-				in_fillmode : gr.Dropdown(value="mean", visible=False),
-				chkbox_group : gr.CheckboxGroup(choices=[], value=[], label="", visible=False),
-				fillmode_btn : gr.Button(visible=False),
 				clear_btn : gr.Button(visible=False),
 				step2_btn : gr.Button(visible=False),
+				in_fillmode : gr.Dropdown(value="mean", visible=False),
+				fillmode_btn : gr.Button(visible=False),
+				chkbox_group : gr.CheckboxGroup(choices=[], value=[], label="", visible=False),
 				step3_btn : gr.Button(visible=False),
-				next_btn : gr.Button("Next step", visible=False),
-				# ------------------stage2-----------------------
+				out_json_file : gr.File(value=None, visible=False),
+				# --------------------Stage2-------------------------
 				run_btn : gr.Button(interactive=False),
 				batch_md : gr.Markdown(visible=False),
 				out_data_file : gr.File(visible=False)}
 	
 	
 	# +========================================================================================+
-	# |                                        stage1-1                                        |
+	# |                                 manage step transition                                 |
 	# +========================================================================================+
-	def save_figures(channel_info, filename1, filename2):
-		
-		tpl_montage = read_custom_montage("./template_chanlocs.loc")
-		tpl_dict = channel_info["templateDict"]
-		in_dict = channel_info["inputDict"]
-		tpl_order = channel_info["templateOrder"]
-		in_order = channel_info["inputOrder"]
-		
-		# get template and input's 2d coords
-		tpl_x = [tpl_dict[channel]["coord_2d"][0] for channel in tpl_order]
-		tpl_y = [tpl_dict[channel]["coord_2d"][1] for channel in tpl_order]
-		in_x = [in_dict[channel]["coord_2d"][0] for channel in in_order]
-		in_y = [in_dict[channel]["coord_2d"][1] for channel in in_order]
-		tpl_coords = np.vstack((tpl_x, tpl_y)).T
-		in_coords = np.vstack((in_x, in_y)).T
-		
-		# get 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))
-		plt.close()
-		
-		# -------------------------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 on it
-		ax.scatter(in_coords[:,0], in_coords[:,1], s=35, color='black')
-		for i, channel in enumerate(in_order):
-			ax.text(in_coords[i,0]+0.003, in_coords[i,1], channel, color='black', fontsize=10.0, va='center')
-		# save input_montage
-		fig.savefig(filename1)
-		
-		# ---------------------------add indications-------------------------------
-		indices = [in_dict[channel]["index"] for channel in in_order if in_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], in_order[i], color='red', fontsize=10.0, va='center')
-		# save mapped_montage
-		fig.savefig(filename2)
-		
-		# -------------------------------------------------------------------------
-		# store the template and input channels' display position (in px).
-		tpl_coords = ax.transData.transform(tpl_coords)
-		in_coords = ax.transData.transform(in_coords)
-		plt.close()
-		
-		for i, channel in enumerate(tpl_order):
-			css_left = (tpl_coords[i,0]-11)/6.4
-			css_bottom = (tpl_coords[i,1]-7)/6.4
-			tpl_dict[channel]["css_position"] = [str(round(css_left, 2))+"%", str(round(css_bottom, 2))+"%"]
-
-		for i, channel in enumerate(in_order):
-			css_left = (in_coords[i,0]-11)/6.4
-			css_bottom = (in_coords[i,1]-7)/6.4
-			in_dict[channel]["css_position"] = [str(round(css_left, 2))+"%", str(round(css_bottom, 2))+"%"]
-		
-		channel_info.update({
-			"templateDict" : tpl_dict,
-			"inputDict" : in_dict
-		})
-		return channel_info
-	
-	def mapping_result(app_info, channel_info):
+	def init_next_step(app_info, channel_info, fillmode, selected_radio, selected_chkbox):
 		stage1_info = app_info["stage1"]
-		filepath = stage1_info["filepath"]
-		
-		# generate and save figures of the input montage and the mapped montage
-		filename1 = filepath+"input_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)
-		stage1_info["filenames"].update({
-		    "input_montage" : filename1,
-		    "mapped_montage" : filename2
-		})
-		
-		# -------------------------determine the next step--------------------------
-		
-		in_num = len(channel_info["inputOrder"])
-		matched_num = 30 - len(stage1_info["missingTemplates"])
-		
-		# if the in_channels has all the 30 tpl_channels (in_num>=30)
-		# -> stage2
-		if matched_num == 30:
-			stage1_info["state"] = "finished"
-			gr.Info('The mapping process has been finished.')
-			
-			if in_num == 30:
-			    md = """
-			    ---
-			    ### Step1: Initial Matching and Rescaling
-			    Below is the result of mapping your channels to our template channels based on their names.
-			    """
-			else:
-			    md = """
-			    ---
-			    ### Step1: Initial Matching and Rescaling
-			    Below is the result of mapping your channels to our template channels based on their names.  
-			    - channels highlighted in red are those that do not match any template channels.
-			    """
+		stage2_info = app_info["stage2"]
+		filepath = stage1_info["filePath"]
+		
+		# ========================================step0=========================================
+		# step0 to step1
+		if stage1_info["state"] == "step1-initializing":
+			#print('step0 -> step1')
+		    
+			# 1. match the names of in_channels and tpl_channels
+			yield {desc_md : gr.Markdown("Mapping...", visible=True)}
+			stage1_info, channel_info, tpl_montage, in_montage = app_utils.match_names(stage1_info, channel_info)
 			
-			app_info["stage1"] = stage1_info
-			return {app_info_json : app_info,
-					channel_info_json : channel_info,
-					map_btn : gr.Button(interactive=True),
-					desc_md : gr.Markdown(md, visible=True),
-					tpl_img : gr.Image(visible=True),
-					mapped_img : gr.Image(value=filename2, visible=True),
-					run_btn : gr.Button(interactive=True)}
-		
-		else:
-			# if matched_num < 30, and there're still some unmatched in_channels
-			# -> assign these in_channels to nearby unmatched tpl_channels
-			if in_num > matched_num:
-				stage1_info["state"] = "step2-initializing"
+			# 2. rescale coordinates
+			yield {desc_md : gr.Markdown("Rescaling...")}
+			channel_info = app_utils.align_coords(channel_info, tpl_montage, in_montage)
+		    
+			# 3. generate and save figures of the montages
+			filename1 = filepath+"input_montage_"+str(random.randint(1,10000))+".png"
+			filename2 = filepath+"mapped_montage_"+str(random.randint(1,10000))+".png"
+			channel_info = app_utils.save_figures(channel_info, tpl_montage, filename1, filename2)
+			stage1_info["fileNames"].update({
+				"input_montage" : filename1,
+				"mapped_montage" : filename2
+			})
+		    
+			# 4. matching result
+			# check if there are red dots (unmatched in_channels) on the input montage
+			unassigned_num = len(stage1_info["unassignedInputs"])
+			if unassigned_num == 0:
 				md = """
 				---
 				### Step1: Initial Matching and Rescaling
-				Below is the result of mapping your channels to our template channels based on their names.  
-				- channels highlighted in red are those that do not match any template channels.
+				Below is the result of mapping your channels to our template channels based on their names.
 				"""
-			
-			# if in_num < 30, but all of them can match to some tpl_channels
-			# -> directly use fillmode to fill the remaining tpl_channels
-			elif in_num == matched_num:
-				stage1_info["state"] = "step3-initializing"
+			else:
 				md = """
 				---
 				### Step1: Initial Matching and Rescaling
-				Below is the result of mapping your channels to our template channels based on their names.
+				Below is the result of mapping your channels to our template channels based on their names.  
+				- channels highlighted in red are those that do not match any template channels.
 				"""
 			
+			stage1_info["state"] = "step1-finished"
 			app_info["stage1"] = stage1_info
-			return {app_info_json : app_info,
+			yield {app_info_json : app_info,
 					channel_info_json : channel_info,
 					map_btn : gr.Button(interactive=True),
-					desc_md : gr.Markdown(md, visible=True),
+					desc_md : gr.Markdown(md),
 					tpl_img : gr.Image(visible=True),
 					mapped_img : gr.Image(value=filename2, visible=True),
 					next_btn : gr.Button(visible=True)}
-	
-	start_stage1 = map_btn.click(
-	    fn = reset_all,
-	    inputs = [in_data_file, in_loc_file, in_samplerate],
-	    outputs = [app_info_json, channel_info_json, map_btn, desc_md, tpl_img, mapped_img, radio_group,
-	                in_fillmode, chkbox_group, fillmode_btn, clear_btn, step2_btn, step3_btn, next_btn,
-	                run_btn, batch_md, out_data_file]
-	).success(
-	    fn = mapping_stage1,
-	    inputs = [app_info_json, channel_info_json], 
-		outputs = [app_info_json, channel_info_json, desc_md]
-	).success(
-	    fn = mapping_result,
-	    inputs = [app_info_json, channel_info_json],
-	    outputs = [app_info_json, channel_info_json, map_btn, desc_md, tpl_img, mapped_img, next_btn, run_btn]
-	)
-	
-	
-	# +========================================================================================+
-	# |                                 manage step transition                                 |
-	# +========================================================================================+
-	def init_next_step(app_info, channel_info, selected_radio, selected_chkbox):
-		stage1_info = app_info["stage1"]
+			
 		
-		# stage1-1 -> stage1-2
-		if stage1_info["state"] == "step2-initializing":
-			#print('step1 -> step2')
-			md = """
-			---
-			### Step2: Forwarding Unmatched Channels
-			Select one of your unmatched channels to forward its data to the empty template channel 
-			currently indicated in red.
-			"""
+		# ========================================step1=========================================
+		elif stage1_info["state"] == "step1-finished":
+			in_num = len(channel_info["inputOrder"])
+			matched_num = 30 - len(stage1_info["missingTemplates"])
 			
-			# initialize the progress indication label for step2
-			stage1_info.update({
-				"state" : "step2-selecting",
-				"fillingCount" : 1,
-				"totalFillingNum" : len(stage1_info["missingTemplates"])
-			})
-			name = stage1_info["missingTemplates"][0]
-			label = "{} (1/{})".format(name, stage1_info["totalFillingNum"])
+			# step1 to step4
+			# the in_channels has all the 30 tpl_channels (in_num>=30)
+			if matched_num == 30:
+				#print('step1 -> step4')
+				md = """
+				---
+				### Mapping result
+				(...)
+				"""
+				
+				# finalize and save the mapping results
+				filename = filepath+"mapping_result.json"
+				stage1_info, stage2_info, channel_info = app_utils.mapping_result(
+																		stage1_info, stage2_info, channel_info, filename)
+				#gr.Info('The mapping process has been finished.')
+				stage1_info["state"] = "finished"
+				app_info["stage1"] = stage1_info
+				app_info["stage2"] = stage2_info
+				yield {app_info_json : app_info,
+						channel_info_json : channel_info,
+						desc_md : gr.Markdown(md),
+						tpl_img : gr.Image(visible=False),
+						mapped_img : gr.Image(visible=False),
+						next_btn : gr.Button(visible=False),
+						out_json_file : gr.File(filename, visible=True),
+						run_btn : gr.Button(interactive=True)}
 			
-			app_info["stage1"] = stage1_info
-			# determine which button to display
-			if stage1_info["totalFillingNum"] == 1:
-				return {app_info_json : app_info,
-					    channel_info_json : channel_info,
-					    desc_md : gr.Markdown(md),
-					    tpl_img : gr.Image(visible=False),
-					    mapped_img : gr.Image(visible=False),
-					    radio_group : gr.Radio(choices=stage1_info["unassignedInputs"], value=[], label=label, visible=True),
-					    clear_btn : gr.Button(visible=True),
-					    next_btn : gr.Button("Next step")}
-			else:
-			    return {app_info_json : app_info,
-					    channel_info_json : channel_info,
-					    desc_md : gr.Markdown(md),
-					    tpl_img : gr.Image(visible=False),
-					    mapped_img : gr.Image(visible=False),
-					    radio_group : gr.Radio(choices=stage1_info["unassignedInputs"], value=[], label=label, visible=True),
-					    clear_btn : gr.Button(visible=True),
-					    step2_btn : gr.Button(visible=True),
-					    next_btn : gr.Button(visible=False)}
-		
-		# stage1-1 -> stage1-3
-		elif stage1_info["state"] == "step3-initializing":
-			#print('step1 -> step3')
-			md = """
-			---
-			### Step3: Filling Remaining Template Channels
-			To run the model successfully, we need to ensure that all 30 template channels are filled. 
-			In this step, you are required to select one of the methods provided below to fill the 
-			remaining empty template channels.
-			"""
-			return {desc_md : gr.Markdown(md),
-					tpl_img : gr.Image(visible=False),
-					mapped_img : gr.Image(visible=False),
-					in_fillmode : gr.Dropdown(visible=True),
-					fillmode_btn : gr.Button(visible=True),
-					next_btn : gr.Button(visible=False)}
+			# step1 to step2
+			# matched_num < 30, and there're still some unmatched in_channels
+			elif in_num > matched_num:
+				#print('step1 -> step2')
+				md = """
+				---
+				### Step2: Forwarding Unmatched Channels
+				Select one of your unmatched channels to forward its data to the empty template channel 
+				currently indicated in red.
+				"""
+				
+				# initialize the progress indication label for step2
+				stage1_info.update({
+					"fillingCount" : 1,
+					"totalFillingNum" : len(stage1_info["missingTemplates"])
+				})
+				name = stage1_info["missingTemplates"][0]
+				label = "{} (1/{})".format(name, stage1_info["totalFillingNum"])
+				
+				stage1_info["state"] = "step2-selecting"
+				app_info["stage1"] = stage1_info
+				# determine which button to display
+				if stage1_info["totalFillingNum"] == 1:
+					yield {app_info_json : app_info,
+							desc_md : gr.Markdown(md),
+							tpl_img : gr.Image(visible=False),
+							mapped_img : gr.Image(visible=False),
+							radio_group : gr.Radio(choices=stage1_info["unassignedInputs"], value=[], label=label, visible=True),
+							clear_btn : gr.Button(visible=True)}
+				else:
+					yield {app_info_json : app_info,
+							desc_md : gr.Markdown(md),
+							tpl_img : gr.Image(visible=False),
+							mapped_img : gr.Image(visible=False),
+							radio_group : gr.Radio(choices=stage1_info["unassignedInputs"], value=[], label=label, visible=True),
+							clear_btn : gr.Button(visible=True),
+							step2_btn : gr.Button(visible=True),
+							next_btn : gr.Button(visible=False)}
+			
+			# step1 to step3-1
+			# in_num < 30, but all of them can match to some tpl_channels
+			elif in_num == matched_num:
+				#print('step1 -> step3-1')
+				md = """
+				---
+				### Step3: Filling Remaining Template Channels
+				To run the model successfully, we need to ensure that all 30 template channels are filled. 
+				In this step, you are required to select one of the methods provided below to fill the 
+				remaining empty template channels.
+				"""
+				
+				stage1_info["state"] = "step3-select-method"
+				app_info["stage1"] = stage1_info
+				yield {app_info_json : app_info,
+						desc_md : gr.Markdown(md),
+						tpl_img : gr.Image(visible=False),
+						mapped_img : gr.Image(visible=False),
+						in_fillmode : gr.Dropdown(visible=True),
+						fillmode_btn : gr.Button(visible=True),
+						next_btn : gr.Button(visible=False)}
 		
-		# stage1-2 -> stage1-3 or stage2
+		# ========================================step2=========================================
 		elif stage1_info["state"] == "step2-selecting":
 		
-			# ----------------------store information before the button click----------------------
-			
+			# --------------------store information before the button click---------------------
 			# check if the user has selected an in_channel to forward to the previous target tpl_channel
 			if selected_radio != []:
 				prev_target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
@@ -634,40 +551,51 @@ with gr.Blocks() as demo:
 				
 				# store the index of the in_channel
 				selected_idx = channel_info["inputDict"][selected_radio]["index"]
-				stage1_info["newOrder"][prev_target_idx] = [selected_idx]
+				stage1_info["mappingData"][0]["newOrder"][prev_target_idx] = [selected_idx]
+				stage1_info["mappingData"][0]["fillFlags"][prev_target_idx] = False
 				# mark the in_channel as assigned and tpl_channel as matched
 				channel_info["templateDict"][prev_target_name]["matched"] = True
 				channel_info["inputDict"][selected_radio]["assigned"] = True				
-				print(prev_target_name, '<-', selected_radio)
-			
-			# ------------------------update information for the next step-------------------------
+				#print(prev_target_name, '<-', selected_radio)
 			
+			# -----------------------update information for the next step-----------------------
 			# update the list of unassignedInputs to exclude the selected in_channel of the previous round
-			stage1_info["unassignedInputs"] = [channel for channel in channel_info["inputOrder"]
-												if channel_info["inputDict"][channel]["assigned"]==False]
+			stage1_info["unassignedInputs"] = app_utils.get_unassigned_inputs(channel_info["inputOrder"],
+																			  channel_info["inputDict"])
 			# update the list of missingTemplates to exclude those filled in step2
-			stage1_info["missingTemplates"] = [channel for channel in channel_info["templateOrder"]
-											if channel_info["templateDict"][channel]["matched"]==False]
+			stage1_info["missingTemplates"] = app_utils.get_empty_templates(channel_info["templateOrder"],
+																			channel_info["templateDict"])
 			
-			# if all the unmatched tpl_channels were filled by in_channels
-			# -> stage2
+			# -----------------------------determine the next step------------------------------
+			# step2 to step4
+			# all the unmatched tpl_channels were filled by in_channels
 			if len(stage1_info["missingTemplates"]) == 0:
-				#print('step2 -> stage2')
-				stage1_info["state"] = "finished"
-				gr.Info('The mapping process has been finished.')
+				#print('step2 -> step4')
+				md = """
+				---
+				### Mapping result
+				(...)
+				"""
 				
+				# finalize and save the mapping results
+				filename = filepath+"mapping_result.json"
+				stage1_info, stage2_info, channel_info = app_utils.mapping_result(
+																		stage1_info, stage2_info, channel_info, filename)
+				#gr.Info('The mapping process has been finished.')
+				stage1_info["state"] = "finished"
 				app_info["stage1"] = stage1_info
-				return {app_info_json : app_info,
+				app_info["stage2"] = stage2_info
+				yield {app_info_json : app_info,
 						channel_info_json : channel_info,
-						desc_md : gr.Markdown(visible=False),
+						desc_md : gr.Markdown(md),
 						radio_group : gr.Radio(visible=False),
+						out_json_file : gr.File(filename, visible=True),
 						clear_btn : gr.Button(visible=False),
 						next_btn : gr.Button(visible=False),
 						run_btn : gr.Button(interactive=True)}
-						
-			# -> stage1-3
+			# step2 to step3-1
 			else:
-				#print('step2 -> step3')
+				#print('step2 -> step3-1')
 				md = """
 				---
 				### Step3: Filling Remaining Template Channels
@@ -676,8 +604,9 @@ with gr.Blocks() as demo:
 				remaining empty template channels.
 				"""
 				
+				stage1_info["state"] = "step3-select-method"
 				app_info["stage1"] = stage1_info
-				return {app_info_json : app_info,
+				yield {app_info_json : app_info,
 						channel_info_json : channel_info,
 						desc_md : gr.Markdown(md),
 						radio_group : gr.Radio(visible=False),
@@ -686,14 +615,86 @@ with gr.Blocks() as demo:
 						clear_btn : gr.Button(visible=False),
 						next_btn : gr.Button(visible=False)}
 		
-		# stage1-3 -> stage2
-		elif stage1_info["state"] == "step3-selecting":
-			#print('step3 -> stage2')
-			stage1_info["state"] = "finished"
-			gr.Info('The mapping process has been finished.')
+		# =======================================step3-1========================================
+		elif stage1_info["state"] == "step3-select-method":
 			
-			# ----------------------store information before the button click----------------------
+			# step3-1 to step4
+			if fillmode == "zero":
+				#print('step3-1 -> step4')
+				md = """
+				---
+				### Mapping result
+				(...)
+				"""
+				
+				# finalize and save the mapping results
+				filename = filepath+"mapping_result.json"
+				stage1_info, stage2_info, channel_info = app_utils.mapping_result(
+																		stage1_info, stage2_info, channel_info, filename)
+				#gr.Info('The mapping process has been finished.')
+				stage1_info["state"] = "finished"
+				app_info["stage1"] = stage1_info
+				app_info["stage2"] = stage2_info
+				yield {app_info_json : app_info,
+						channel_info_json : channel_info,
+						desc_md : gr.Markdown(md),
+						in_fillmode : gr.Dropdown(visible=False),
+						fillmode_btn : gr.Button(visible=False),
+						out_json_file : gr.File(filename, visible=True),
+						run_btn : gr.Button(interactive=True)}
+			# step3-1 to step3-2
+			elif fillmode == "mean":
+				#print('step3-1 -> step3-2')
+				md = """
+				---
+				### Step3: Fill the remaining template channels
+				The current empty template channel, indicated in red, will be filled with the average 
+				value of the data from the selected channels. (By default, the 4 nearest channels are pre-selected.)
+				"""
+				
+				# find the 4 nearest in_channels for each unmatched tpl_channels
+				stage1_info["mappingData"][0]["newOrder"] = app_utils.find_neighbors(
+																		channel_info,
+																		stage1_info["missingTemplates"],
+																		stage1_info["mappingData"][0]["newOrder"])
+				
+				# initialize the progress indication label
+				stage1_info.update({
+					"fillingCount" : 1,
+					"totalFillingNum" : len(stage1_info["missingTemplates"])
+				})
+				target_name = stage1_info["missingTemplates"][0]
+				target_idx = channel_info["templateDict"][target_name]["index"]
+				chkbox_value = stage1_info["mappingData"][0]["newOrder"][target_idx]
+				chkbox_value = [channel_info["inputOrder"][i] for i in chkbox_value]
+				chkbox_label = "{} (1/{})".format(target_name, stage1_info["totalFillingNum"])
+				
+				stage1_info["state"] = "step3-2-selecting"
+				app_info["stage1"] = stage1_info
+				# determine which button to display
+				if stage1_info["totalFillingNum"] == 1:
+					yield {app_info_json : app_info,
+					        desc_md : gr.Markdown(md),
+							in_fillmode : gr.Dropdown(visible=False),
+				    		fillmode_btn : gr.Button(visible=False),
+				    		chkbox_group : gr.CheckboxGroup(choices=channel_info["inputOrder"],
+															value=chkbox_value, label=chkbox_label, visible=True),
+				            next_btn : gr.Button(visible=True)}
+				else:
+					yield {app_info_json : app_info,
+					        desc_md : gr.Markdown(md),
+							in_fillmode : gr.Dropdown(visible=False),
+				    		fillmode_btn : gr.Button(visible=False),
+				    		chkbox_group : gr.CheckboxGroup(choices=channel_info["inputOrder"],
+															value=chkbox_value, label=chkbox_label, visible=True),
+				            step3_btn : gr.Button(visible=True)}
+		
+		# =======================================step3-2========================================
+		# step3-2 to step4
+		elif stage1_info["state"] == "step3-2-selecting":
+			#print('step3-2 -> step4')
 			
+			# --------------------store information before the button click---------------------
 			# if the user didn't uncheck all in_channel checkboxes
 			if selected_chkbox != []:
 				prev_target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
@@ -701,22 +702,36 @@ with gr.Blocks() as demo:
 				
 				# store the indices of the in_channels
 				selected_indices = [channel_info["inputDict"][channel]["index"] for channel in selected_chkbox]
-				stage1_info["newOrder"][prev_target_idx] = selected_indices
-				#print(f'{prev_target_name}({prev_target_idx}): {selected_chkbox}')
-			# -------------------------------------------------------------------------------------
+				stage1_info["mappingData"][0]["newOrder"][prev_target_idx] = selected_indices
+				#print(f'{prev_target_name}({prev_target_idx}): {selected_indices}')
+			# ----------------------------------------------------------------------------------
+			md = """
+			---
+			### Mapping result
+			(...)
+			"""
 			
+			# finalize and save the mapping results
+			filename = filepath+"mapping_result.json"
+			stage1_info, stage2_info, channel_info = app_utils.mapping_result(
+																	stage1_info, stage2_info, channel_info, filename)
+			#gr.Info('The mapping process has been finished.')
+			stage1_info["state"] = "finished"
 			app_info["stage1"] = stage1_info
-			return {app_info_json : app_info,
-					desc_md : gr.Markdown(visible=False),
+			app_info["stage2"] = stage2_info
+			yield {app_info_json : app_info,
+					channel_info_json : channel_info,
+					desc_md : gr.Markdown(md),
 					chkbox_group : gr.CheckboxGroup(visible=False),
 					next_btn : gr.Button(visible=False),
+					out_json_file : gr.File(filename, visible=True),
 					run_btn : gr.Button(interactive=True)}
 	
 	next_btn.click(
 		fn = init_next_step,
-		inputs = [app_info_json, channel_info_json, radio_group, chkbox_group],
-		outputs = [app_info_json, channel_info_json, desc_md, tpl_img, mapped_img, radio_group,
-		            in_fillmode, chkbox_group, fillmode_btn, clear_btn, step2_btn, next_btn, run_btn]
+		inputs = [app_info_json, channel_info_json, in_fillmode, radio_group, chkbox_group],
+		outputs = [app_info_json, channel_info_json, desc_md, tpl_img, mapped_img, radio_group, clear_btn, step2_btn,
+		            in_fillmode, fillmode_btn, chkbox_group, step3_btn, out_json_file, next_btn, run_btn]
 	).success(
 	    fn = None,
 	    js = init_js,
@@ -726,9 +741,25 @@ with gr.Blocks() as demo:
 	
 	
 	# +========================================================================================+
-	# |                                        stage1-2                                        |
+	# |                                      Stage1-step0                                      |
 	# +========================================================================================+
-	# ....
+	map_btn.click(
+	    fn = reset_all,
+	    inputs = [in_data_file, in_loc_file, in_samplerate],
+	    outputs = [app_info_json, channel_info_json, map_btn, desc_md, next_btn, tpl_img, mapped_img,
+	    			radio_group, clear_btn, step2_btn, in_fillmode, fillmode_btn, chkbox_group, step3_btn, out_json_file,
+	                run_btn, batch_md, out_data_file]
+	).success(
+	    fn = init_next_step,
+	    inputs = [app_info_json, channel_info_json, in_fillmode, radio_group, chkbox_group], 
+		outputs = [app_info_json, channel_info_json, map_btn, desc_md, tpl_img, mapped_img, next_btn]
+	)
+	
+	
+	# +========================================================================================+
+	# |                                      Stage1-step2                                      |
+	# +========================================================================================+
+	# ...
 	@radio_group.select(inputs = app_info_json, outputs = [step2_btn, next_btn])
 	def determine_button(app_info):
 		stage1_info = app_info["stage1"]
@@ -751,8 +782,7 @@ with gr.Blocks() as demo:
 	def update_radio(app_info, channel_info, selected):
 		stage1_info = app_info["stage1"]
 		
-		# ----------------------store information before the button click----------------------
-		
+		# ----------------------store information before the button click-----------------------
 		# check if the user has selected an in_channel to forward to the previous target tpl_channel
 		if selected != []:
 			prev_target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
@@ -760,18 +790,18 @@ with gr.Blocks() as demo:
 			
 			# store the index of the selected in_channel
 			selected_idx = channel_info["inputDict"][selected]["index"]
-			stage1_info["newOrder"][prev_target_idx] = [selected_idx]
+			stage1_info["mappingData"][0]["newOrder"][prev_target_idx] = [selected_idx]
+			stage1_info["mappingData"][0]["fillFlags"][prev_target_idx] = False
 			# mark the in_channel as assigned and tpl_channel as matched
 			channel_info["templateDict"][prev_target_name]["matched"] = True
 			channel_info["inputDict"][selected]["assigned"] = True
-			print(prev_target_name, '<-', selected)
+			#print(prev_target_name, '<-', selected)
 		
-		# ------------------------update information for the new round-------------------------
+		# ------------------------update information for the new round--------------------------
 		stage1_info["fillingCount"] += 1
 		
 		# update the list of unassignedInputs to exclude the selected in_channel of the previous round
-		stage1_info["unassignedInputs"] = [channel for channel in channel_info["inputOrder"]
-												if channel_info["inputDict"][channel]["assigned"]==False]
+		stage1_info["unassignedInputs"] = app_utils.get_unassigned_inputs(channel_info["inputOrder"], channel_info["inputDict"])
 		# update the progress indication label
 		target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
 		radio_label = "{} ({}/{})".format(target_name, stage1_info["fillingCount"], stage1_info["totalFillingNum"])
@@ -784,7 +814,7 @@ with gr.Blocks() as demo:
 					radio_group : gr.Radio(choices=stage1_info["unassignedInputs"],
 					                        value=[], label=radio_label),
 					step2_btn : gr.Button(visible=False),
-					next_btn : gr.Button("Next step", visible=True)}
+					next_btn : gr.Button(visible=True)}
 		else:
 			return {app_info_json : app_info,
 					channel_info_json : channel_info,
@@ -804,71 +834,12 @@ with gr.Blocks() as demo:
 	
 	
 	# +========================================================================================+
-	# |                                        stage1-3                                        |
-	# +========================================================================================+
-	def fill_value(app_info, channel_info, fillmode):
-		stage1_info = app_info["stage1"]
-		
-		if fillmode == "zero":
-			stage1_info["state"] = "finished"
-			gr.Info('The mapping process has been finished.')
-			
-			app_info["stage1"] = stage1_info
-			return {app_info_json : app_info,
-					desc_md : gr.Markdown(visible=False),
-					in_fillmode : gr.Dropdown(visible=False),
-					fillmode_btn : gr.Button(visible=False),
-					run_btn : gr.Button(interactive=True)}
-		
-		elif fillmode == "mean":
-			md = """
-			---
-			### Step3: Fill the remaining template channels
-			The current empty template channel, indicated in red, will be filled with the average 
-			value of the data from the selected channels. (By default, the 4 nearest channels are pre-selected.)
-			"""
-			
-			# find the 4-NN in_channels for each of the unmatched tpl_channels
-			new_idx = find_neighbors(channel_info, stage1_info["missingTemplates"], stage1_info["newOrder"])
-			
-			stage1_info.update({
-				"state" : "step3-selecting",
-				"newOrder" : new_idx,
-				"fillingCount" : 1,
-				"totalFillingNum" : len(stage1_info["missingTemplates"])
-			})
-			
-			# initialize the progress indication label
-			target_name = stage1_info["missingTemplates"][0]
-			target_idx = channel_info["templateDict"][target_name]["index"]
-			chkbox_value = stage1_info["newOrder"][target_idx]
-			chkbox_value = [channel_info["inputOrder"][i] for i in chkbox_value]
-			chkbox_label = "{} (1/{})".format(target_name, stage1_info["totalFillingNum"])
-			
-			app_info["stage1"] = stage1_info
-			# determine which button to display
-			if stage1_info["totalFillingNum"] == 1:
-			    return {app_info_json : app_info,
-			            desc_md : gr.Markdown(md),
-					    in_fillmode : gr.Dropdown(visible=False),
-		        		fillmode_btn : gr.Button(visible=False),
-		        		chkbox_group : gr.CheckboxGroup(choices=channel_info["inputOrder"],
-													    value=chkbox_value, label=chkbox_label, visible=True),
-		                next_btn : gr.Button(visible=True)}
-			else:
-			    return {app_info_json : app_info,
-			            desc_md : gr.Markdown(md),
-					    in_fillmode : gr.Dropdown(visible=False),
-		        		fillmode_btn : gr.Button(visible=False),
-		        		chkbox_group : gr.CheckboxGroup(choices=channel_info["inputOrder"],
-													    value=chkbox_value, label=chkbox_label, visible=True),
-		                step3_btn : gr.Button(visible=True)}
-	
+	# |                                      Stage1-step3                                      |
+	# +========================================================================================+	
 	def update_chkbox(app_info, channel_info, selected):
 		stage1_info = app_info["stage1"]
 		
-		# ----------------------store information before the button click----------------------
-		
+		# ----------------------store information before the button click-----------------------
 		# if the user didn't uncheck all in_channel checkboxes
 		if selected != []:
 			prev_target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
@@ -876,16 +847,16 @@ with gr.Blocks() as demo:
 			
 			# store the indices of the selected in_channels
 			selected_indices = [channel_info["inputDict"][channel]["index"] for channel in selected]
-			stage1_info["newOrder"][prev_target_idx] = selected_indices
-			#print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
+			stage1_info["mappingData"][0]["newOrder"][prev_target_idx] = selected_indices
+			#print(f'{prev_target_name}({prev_target_idx}): {selected_indices}')
 		
-		# ------------------------update information for the new round-------------------------
+		# ------------------------update information for the new round--------------------------
 		stage1_info["fillingCount"] += 1
 		
 		# update the progress indication label
 		target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
 		target_idx = channel_info["templateDict"][target_name]["index"]
-		chkbox_value = stage1_info["newOrder"][target_idx]
+		chkbox_value = stage1_info["mappingData"][0]["newOrder"][target_idx]
 		chkbox_value = [channel_info["inputOrder"][i] for i in chkbox_value]
 		chkbox_label = "{} ({}/{})".format(target_name, stage1_info["fillingCount"], stage1_info["totalFillingNum"])
 		
@@ -895,15 +866,16 @@ with gr.Blocks() as demo:
 			return {app_info_json : app_info,
 					chkbox_group : gr.CheckboxGroup(value=chkbox_value, label=chkbox_label),
 					step3_btn : gr.Button(visible=False),
-					next_btn : gr.Button("Submit", visible=True)}
+					next_btn : gr.Button(visible=True)}
 		else:
 			return {app_info_json : app_info,
 					chkbox_group : gr.CheckboxGroup(value=chkbox_value, label=chkbox_label)}
 	
 	fillmode_btn.click(
-	    fn = fill_value,
-	    inputs = [app_info_json, channel_info_json, in_fillmode],
-		outputs = [app_info_json, desc_md, in_fillmode, fillmode_btn, chkbox_group, step3_btn, next_btn, run_btn]
+	    fn = init_next_step,
+		inputs = [app_info_json, channel_info_json, in_fillmode, radio_group, chkbox_group],
+		outputs = [app_info_json, channel_info_json, desc_md, in_fillmode, fillmode_btn, chkbox_group, step3_btn,
+					out_json_file, next_btn, run_btn]
 	).success(
 	    fn = None,
 	    js = init_js,
@@ -924,59 +896,45 @@ with gr.Blocks() as demo:
 	
 	
     # +========================================================================================+
-	# |                                  stage2: decode data                                   |
+	# |                                  Stage2: decode data                                   |
 	# +========================================================================================+
-	def reset_run(app_info, channel_info, modelname):
+	def reset_run(app_info, modelname):
 		stage1_info = app_info["stage1"]
 		stage2_info = app_info["stage2"]
 	    
-	    # delete the previous folder of stage2
-		filepath = stage2_info["filepath"]
+	    # delete the previous folder of Stage2
+		filepath = stage2_info["filePath"]
 		utils.dataDelete(filepath)
-		# establish a new folder for stage2
-		new_filepath = app_info["rootFilepath"]+"stage2_"+str(random.randint(1,10000))+"/"
+		# establish a new folder for Stage2
+		new_filepath = app_info["rootPath"]+"stage2_"+str(random.randint(1,10000))+"/"
 		os.mkdir(new_filepath)
 		# generate the output filename
-		filename = stage1_info["filenames"]["input_data"]
+		filename = stage1_info["fileNames"]["input_data"]
 		filename = os.path.basename(str(filename))
 		new_filename = os.path.splitext(filename)[0]+'_'+modelname+'.csv'
 		
-	    # reset inputChannel.assigned back to the state after stage1
-		for channel in stage1_info["unassignedInputs"]:
-			channel_info["inputDict"][channel]["assigned"] = False
-		# calculate how many times the model needs to be run
-		unassigned_num = len(stage1_info["unassignedInputs"])
-		batch_num = math.ceil(unassigned_num/30) + 1
-		
-		app_info.update({
-			#"currentStage" : "stage2",
-			"stage2" : {
-				"filepath" : new_filepath,
-				"filenames" : {
-					"output_data" : new_filepath + new_filename
-				},
-				#"state" : "initializing",
-				"totalBatchNum" : batch_num,
-				"newOrder" : [[]]*30,
-				"unassignedInputs" : stage1_info["unassignedInputs"]
+		stage2_info.update({
+			"filePath" : new_filepath,
+			"fileNames" : {
+				"output_data" : new_filepath + new_filename
 			}
 		})
+		app_info["stage2"] = stage2_info
 		return {app_info_json : app_info,
-				channel_info_json : channel_info,
 				#run_btn : gr.Button(interactive=False),
 				batch_md : gr.Markdown(visible=False),
 				out_data_file : gr.File(visible=False)}
 	
-	def run_model(app_info, channel_info, modelname):
+	def run_model(app_info, modelname):
 		stage1_info = app_info["stage1"]
 		stage2_info = app_info["stage2"]
 		
-		filepath = stage2_info["filepath"]
+		filepath = stage2_info["filePath"]
 		samplerate = app_info["sampleRate"]
-		filename = stage1_info["filenames"]["input_data"]
-		new_filename = stage2_info["filenames"]["output_data"]
+		filename = stage1_info["fileNames"]["input_data"]
+		new_filename = stage2_info["fileNames"]["output_data"]
 		
-		# set a flag to record whether the user has clicked the map_btn or run_btn while running the model
+		# flag to indicate if the process has been interrupted by the user
 		break_flag = False
 		
 		# run the model multiple times until all in_channels are reconstructed
@@ -988,7 +946,7 @@ with gr.Blocks() as demo:
 				#utils.dataDelete(filepath+"temp_data/")
 				#os.mkdir(filepath+"temp_data/")
 			except FileNotFoundError:
-				#print('break1!!')
+				print('break1!!')
 				break_flag = True
 				break
 			except OSError as e:
@@ -998,33 +956,32 @@ with gr.Blocks() as demo:
 			md = "Running model({}/{})...".format(i+1, stage2_info["totalBatchNum"])
 			yield {batch_md : gr.Markdown(md, visible=True)}
 			
-			# ....
-			if i == 0:
-			    new_idx = stage1_info["newOrder"]
-			else:
-			    # if this is not the first time running the model, the in_channels that have
-			    # not been reconstructed yet will be optimally mapped to the template.
-				stage2_info, channel_info = mapping_stage2(stage2_info, channel_info)
-				new_idx = stage2_info["newOrder"]
-			print('unassigned num:', len(stage2_info["unassignedInputs"]))
-			
+			# get the mapped index order and the filled status for each tpl_channels
+			new_idx = stage1_info["mappingData"][i]["newOrder"]
+			fill_flags = stage1_info["mappingData"][i]["fillFlags"]
 			# ----------------------------------------------------------------------
 			try:
 				# step1: Reorder input data
-				reorder_input_data(new_idx, filename, filepath+"temp_data/mapped.csv")
+				data_shape = app_utils.reorder_data(new_idx, fill_flags, filename, filepath+"temp_data/mapped.csv")
+				if modelname == "(mapped data)":
+					new_filename = filepath+"temp_data/mapped.csv"
+					break
 				# step2: Data preprocessing
 				total_file_num = utils.preprocessing(filepath+"temp_data/", "mapped.csv", samplerate)
 				# step3: Signal reconstruction
 				utils.reconstruct(modelname, total_file_num, filepath+"temp_data/", "denoised.csv", samplerate)
+				if modelname == "(denoised data)":
+					new_filename = filepath+"temp_data/denoised.csv"
+					break
 				# step4: Restore original order
-				restore_original_order(channel_info, i, new_idx, filepath+"temp_data/denoised.csv", new_filename)
+				app_utils.restore_order(i, data_shape, new_idx, fill_flags, filepath+"temp_data/denoised.csv", new_filename)
+				break
 			except FileNotFoundError:
-				#print('break2!!')
+				print('break2!!')
 				break_flag = True
 				break
 			# ----------------------------------------------------------------------
 			utils.dataDelete(filepath+"temp_data/")
-			app_info["stage2"] = stage2_info
 		
 		if break_flag == True:
 			yield {batch_md : gr.Markdown(visible=False)}
@@ -1035,12 +992,12 @@ with gr.Blocks() as demo:
 	
 	run_btn.click(
 		fn = reset_run,
-		inputs = [app_info_json, channel_info_json, in_modelname],
-		outputs = [app_info_json, channel_info_json, run_btn, batch_md, out_data_file]
+		inputs = [app_info_json, in_modelname],
+		outputs = [app_info_json, run_btn, batch_md, out_data_file]
 		
 	).success(
 		fn = run_model,
-		inputs = [app_info_json, channel_info_json, in_modelname],
+		inputs = [app_info_json, in_modelname],
 		outputs = [run_btn, batch_md, out_data_file]
 	)
 	

channel_mapping.py → app_utils.py

@@ -1,144 +1,208 @@
 import utils
-import time
 import os
+import time
+import math
+import json
 import numpy as np
-import gradio as gr
-
+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 reorder_input_data(old_idx, filename, new_filename):
-	old_data = utils.read_train_data(filename) # original raw data
-	new_data = np.zeros((30, old_data.shape[1])) # to store reordered raw data
-	print('old.shape, new.shape: ', old_data.shape, new_data.shape)
-	print('new index order:', old_idx)
-	
-	zero_arr = np.zeros((1, old_data.shape[1]))
-	old_data = np.concatenate((old_data, zero_arr), axis=0)
-	
-	for i in range(30):
-		idx_set = old_idx[i]
-		#print("channel_{}'s index set: {}".format(i, idx_set))
-		
-		if idx_set == []:
+def reorder_data(idx_order, fill_flags, filename, new_filename):
+	# read the input data
+	raw_data = utils.read_train_data(filename)
+	new_data = np.zeros((30, raw_data.shape[1]))
+	
+	zero_arr = np.zeros((1, raw_data.shape[1]))
+	for i, (idx_set, flag) in enumerate(zip(idx_order, fill_flags)):
+		if flag == False:
+			new_data[i, :] = raw_data[idx_set[0], :]
+		elif idx_set == []:
 			new_data[i, :] = zero_arr
 		else:
-			tmp_data = [old_data[j, :] for j in idx_set]
+			tmp_data = [raw_data[j, :] for j in idx_set]
 			new_data[i, :] = np.mean(tmp_data, axis=0)
 	
+	#print(raw_data.shape, new_data.shape)
 	utils.save_data(new_data, new_filename)
-	return
+	return raw_data.shape
 
-def restore_original_order(channel_info, cnt, old_idx, filename, new_filename):
-	old_data = utils.read_train_data(filename) # denoised data
-	tpl_order = channel_info["templateOrder"]
-	in_order = channel_info["inputOrder"]
-	
-	if cnt == 0:
-		new_data = np.zeros((len(in_order), old_data.shape[1]))
+def restore_order(batch_cnt, raw_data_shape, idx_order, fill_flags, filename, new_filename):
+	# read the denoised data
+	d_data = utils.read_train_data(filename)
+	if batch_cnt == 0:
+		new_data = np.zeros((raw_data_shape[0], d_data.shape[1]))
 	else:
 		new_data = utils.read_train_data(new_filename)
 	
-	for i, channel in enumerate(tpl_order):
-		idx_set = old_idx[i]
-		
-		# ignore if this channel was filled with fillmode ('mean' or 'zero')
-		if len(idx_set)==1 and channel_info["templateDict"][channel]["matched"]==True:
-			new_data[idx_set[0], :] = old_data[i, :]
+	for i, (idx_set, flag) in enumerate(zip(idx_order, fill_flags)):
+		# ignore if this channel was filled using "fillmode"
+		if flag == False:
+			new_data[idx_set[0], :] = d_data[i, :]
 	
-	print('old.shape, new.shape: ', old_data.shape, new_data.shape)
+	#print(d_data.shape, new_data.shape)
 	utils.save_data(new_data, new_filename)
 	return
 
+def get_matched(tpl_order, tpl_dict):
+    return [channel for channel in tpl_order if tpl_dict[channel]["matched"]==True]
+
+def get_empty_templates(tpl_order, tpl_dict):
+    return [channel for channel in tpl_order if tpl_dict[channel]["matched"]==False]
+
+def get_unassigned_inputs(in_order, in_dict):
+    return [channel for channel in in_order if in_dict[channel]["assigned"]==False]
+
 def read_montage_data(loc_file):
-		
 	tpl_montage = read_custom_montage("./template_chanlocs.loc")
 	in_montage = read_custom_montage(loc_file)
+	tpl_order = tpl_montage.ch_names
+	in_order = in_montage.ch_names
 	tpl_dict = {}
 	in_dict = {}
 	
-	for i in range(30):
-		channel = tpl_montage.ch_names[i]
-		tpl_montage.rename_channels({channel: str.upper(channel)}) # convert all channel names to uppercase
-		
-		channel = str.upper(channel)
-		tpl_dict[channel] = {
+	# convert all channel names to uppercase and store the channel information
+	for i, channel in enumerate(tpl_order):
+		up_channel = str.upper(channel)
+		tpl_montage.rename_channels({channel: up_channel})
+		tpl_dict[up_channel] = {
 			"index" : i,
-			"coord_3d" : tpl_montage.get_positions()['ch_pos'][channel],
+			"coord_3d" : tpl_montage.get_positions()['ch_pos'][up_channel],
 			"matched" : False
 		}
-	for i in range(len(in_montage.ch_names)):
-		channel = in_montage.ch_names[i]
-		in_montage.rename_channels({channel: str.upper(channel)}) # convert all channel names to uppercase
-		
-		channel = str.upper(channel)
-		in_dict[channel] = {
+	for i, channel in enumerate(in_order):
+		up_channel = str.upper(channel)
+		in_montage.rename_channels({channel: up_channel})
+		in_dict[up_channel] = {
 			"index" : i,
-			"coord_3d" : in_montage.get_positions()['ch_pos'][channel],
+			"coord_3d" : in_montage.get_positions()['ch_pos'][up_channel],
 			"assigned" : False
 		}
-	
 	return tpl_montage, in_montage, tpl_dict, in_dict
 
-def align_coords(channel_info, tpl_montage, in_montage):
-
+def save_figures(channel_info, tpl_montage, filename1, filename2):
+	tpl_order = channel_info["templateOrder"]
+	in_order = channel_info["inputOrder"]
 	tpl_dict = channel_info["templateDict"]
 	in_dict = channel_info["inputDict"]
+	
+	# get the 2D coordinates
+	tpl_x = [tpl_dict[channel]["coord_2d"][0] for channel in tpl_order]
+	tpl_y = [tpl_dict[channel]["coord_2d"][1] for channel in tpl_order]
+	in_x = [in_dict[channel]["coord_2d"][0] for channel in in_order]
+	in_y = [in_dict[channel]["coord_2d"][1] for channel in in_order]
+	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))
+	plt.close()
+	
+	# -------------------------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, channel in enumerate(in_order):
+		ax.text(in_coords[i,0]+0.003, in_coords[i,1], channel, color='black', fontsize=10.0, va='center')
+	# save input_montage
+	fig.savefig(filename1)
+	
+	# ---------------------------add indications-------------------------------
+	# plot unmatched input channels in red
+	indices = [in_dict[channel]["index"] for channel in in_order if in_dict[channel]["assigned"]==False]
+	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], in_order[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()
+	
+	for i, channel in enumerate(tpl_order):
+		css_left = (tpl_coords[i,0]-11)/6.4
+		css_bottom = (tpl_coords[i,1]-7)/6.4
+		tpl_dict[channel]["css_position"] = [str(round(css_left, 2))+"%", str(round(css_bottom, 2))+"%"]
+	for i, channel in enumerate(in_order):
+		css_left = (in_coords[i,0]-11)/6.4
+		css_bottom = (in_coords[i,1]-7)/6.4
+		in_dict[channel]["css_position"] = [str(round(css_left, 2))+"%", str(round(css_bottom, 2))+"%"]
+	
+	channel_info.update({
+		"templateDict" : tpl_dict,
+		"inputDict" : in_dict
+	})
+	return channel_info
+
+def align_coords(channel_info, tpl_montage, in_montage):
 	tpl_order = channel_info["templateOrder"]
 	in_order = channel_info["inputOrder"]
-	matched = [channel for channel in tpl_order if tpl_dict[channel]["matched"]==True]
-	
+	tpl_dict = channel_info["templateDict"]
+	in_dict = channel_info["inputDict"]
+	matched = get_matched(tpl_order, tpl_dict)
 	
-	# 2-D (to indicate the location of the missing template channel when fill_mode:'mean')
+	# 2D alignment (for visualization purposes)
 	fig = [tpl_montage.plot(), in_montage.plot()]
 	ax = [fig[0].axes[0], fig[1].axes[0]]
 	
-	# get the original coords
-	#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)
+	# extract the displayed 2D coordinates from the plots
 	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[channel]["index"]] for channel in matched])
 	matched_in = np.array([all_in[in_dict[channel]["index"]] for channel in matched])
 	
-	# transform the xy axis (input's -> template's)
+	# apply TPS to transform in_channels positions 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 to all input channels
+	# 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
 	
-	# save template's and input's 2d position
+	# store the 2D positions
 	for i, channel in enumerate(tpl_order):
 		tpl_dict[channel]["coord_2d"] = all_tpl[i]
 	for i, channel in enumerate(in_order):
 		in_dict[channel]["coord_2d"] = transformed_in[i].tolist()
 	
 	
-	# 3-d (to use KNN)
-	# get the original coords
+	# 3D alignment
 	all_tpl = np.array([tpl_dict[channel]["coord_3d"].tolist() for channel in tpl_order])
 	all_in = np.array([in_dict[channel]["coord_3d"].tolist() for channel in in_order])
 	matched_tpl = np.array([all_tpl[tpl_dict[channel]["index"]] for channel in matched])
 	matched_in = np.array([all_in[in_dict[channel]["index"]] for channel in matched])
 	
-	# transform the xyz axis (input's -> template's)
 	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')
 	
-	# apply to all input channels
 	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
 	
-	# update input's 3d position
+	# update in_channels' 3D positions
 	for i, channel in enumerate(in_order):
 		in_dict[channel]["coord_3d"] = transformed_in[i].tolist()
 	
@@ -149,133 +213,153 @@ def align_coords(channel_info, tpl_montage, in_montage):
 	return channel_info
 
 def find_neighbors(channel_info, missing_channels, new_idx):
+	in_order = channel_info["inputOrder"]
 	tpl_dict = channel_info["templateDict"]
 	in_dict = channel_info["inputDict"]
-	in_order = channel_info["inputOrder"]
 	
+	# get the 3D coordinates
 	all_in = [np.array(in_dict[channel]["coord_3d"]) for channel in in_order]
-	missing_tpl = [np.array(tpl_dict[channel]["coord_3d"]) for channel in missing_channels]
+	empty_tpl = [np.array(tpl_dict[channel]["coord_3d"]) for channel in missing_channels]
 	
 	# use KNN to choose k nearest channels
 	k = 4 if len(in_order)>4 else len(in_order)
 	knn = NearestNeighbors(n_neighbors=k, metric='euclidean')
 	knn.fit(all_in)
-	
 	for i, channel in enumerate(missing_channels):
-		distances, indices = knn.kneighbors(missing_tpl[i].reshape(1,-1))
-		#selected = [in_order[j] for j in indices[0]]
-		#print(channel, ':', selected)
-		
+		distances, indices = knn.kneighbors(empty_tpl[i].reshape(1,-1))
 		idx = tpl_dict[channel]["index"]
 		new_idx[idx] = indices[0].tolist()
     
 	return new_idx
 
-def mapping_stage1(app_info, channel_info):
-	yield app_info, channel_info, gr.Markdown("Mapping...", visible=True)
-	second1 = time.time()
-	
-	loc_file = app_info["stage1"]["filenames"]["input_loc"]
+def match_names(stage1_info, channel_info):
+    # read the location file
+	loc_file = stage1_info["fileNames"]["input_loc"]
 	tpl_montage, in_montage, tpl_dict, in_dict = read_montage_data(loc_file)
 	tpl_order = tpl_montage.ch_names
 	in_order = in_montage.ch_names
-	new_idx = [[]]*30 # store the indices of the in_channels in the order of tpl_channls
+	# list to store the indices of the in_channels in the order of tpl_channls
+	new_idx = [[]]*30
+	# flags to record if each tpl_channel's data is filled by "fillmode"
+	fill_flags = [True]*30
+	
 	alias_dict = {
 		'T3': 'T7',
 		'T4': 'T8',
 		'T5': 'P7',
 		'T6': 'P8'
 	}
-	
-	# match the names of input channels and template channels
 	for i, channel in enumerate(tpl_order):
 		if channel in alias_dict and alias_dict[channel] in in_dict:
-			tpl_montage.rename_channels({channel: alias_dict[channel]}) # rename the current tpl_channel
+			tpl_montage.rename_channels({channel: alias_dict[channel]})
 			tpl_dict[alias_dict[channel]] = tpl_dict.pop(channel)
 			channel = alias_dict[channel]
 			
 		if channel in in_dict:
 			new_idx[i] = [in_dict[channel]["index"]]
+			fill_flags[i] = False
 			tpl_dict[channel]["matched"] = True
 			in_dict[channel]["assigned"] = True
 	
 	# update the names
 	tpl_order = tpl_montage.ch_names
 	
+	stage1_info.update({
+	    "unassignedInputs" : get_unassigned_inputs(in_order, in_dict),
+	    "missingTemplates" : get_empty_templates(tpl_order, tpl_dict),
+	    "mappingData" : [
+		{
+			"newOrder" : new_idx,
+			"fillFlags" : fill_flags
+		}
+	    ]
+	})
 	channel_info.update({
+		"templateOrder" : tpl_order,
+		"inputOrder" : in_order,
 	    "templateDict" : tpl_dict,
-	    "inputDict" : in_dict,
-	    "templateOrder" : tpl_order,
-		"inputOrder" : in_order
-	})
-	app_info["stage1"].update({
-	    "newOrder" : new_idx,
-	    "unassignedInputs" : [channel for channel in in_order if in_dict[channel]["assigned"]==False],
-	    "missingTemplates" : [channel for channel in tpl_order if tpl_dict[channel]["matched"]==False]
+	    "inputDict" : in_dict
 	})
-	
-	# align input, template's coordinates
-	channel_info = align_coords(channel_info, tpl_montage, in_montage)
-	
-	second2 = time.time()
-	print('Mapping (stage1) finished in',second2 - second1,'s.')
-	yield app_info, channel_info, gr.Markdown("", visible=False)
+	return stage1_info, channel_info, tpl_montage, in_montage
 
-def mapping_stage2(stage2_info, channel_info):
-	second1 = time.time()
-
-	tpl_dict = channel_info["templateDict"]
-	in_dict = channel_info["inputDict"]
+def optimal_mapping(channel_info):
 	tpl_order = channel_info["templateOrder"]
 	in_order = channel_info["inputOrder"]
-	unassigned = stage2_info["unassignedInputs"]
-	
-	tpl_coords = np.array([tpl_dict[channel]["coord_3d"] for channel in tpl_order])
-	unassigned_coords = np.array([in_dict[channel]["coord_3d"] for channel in unassigned])
-	
+	tpl_dict = channel_info["templateDict"]
+	in_dict = channel_info["inputDict"]
+	unassigned = get_unassigned_inputs(in_order, in_dict)
 	# reset all tpl.matched to False
 	for channel in tpl_dict:
 		tpl_dict[channel]["matched"] = False
 	
-	# initialize the cost matrix
+	# get the 3D coordinates
+	all_tpl = np.array([tpl_dict[channel]["coord_3d"] for channel in tpl_order])
+	unassigned_in = np.array([in_dict[channel]["coord_3d"] for channel in unassigned])
+	
+	# initialize the cost matrix for the Hungarian algorithm
 	if len(unassigned) < 30:
 		cost_matrix = np.full((30, 30), 1e6) # add dummy channels to ensure num_col >= num_row
 	else:
 		cost_matrix = np.zeros((30, len(unassigned)))
+	# fill the cost matrix with Euclidean distances between tpl_channels and unassigned in_channels
 	for i in range(30):
 		for j in range(len(unassigned)):
-			cost_matrix[i][j] = np.linalg.norm((tpl_coords[i]-unassigned_coords[j])*1000) # Euclidean distance
-			#print(cost_matrix[i][j], tpl_coords[i] - unassigned_coords[j])
+			cost_matrix[i][j] = np.linalg.norm((all_tpl[i]-unassigned_in[j])*1000)
 	
-	# use Hungarian Algorithm to find the minimum sum of distance of (input's coord to template's coord)...?
+	# apply the Hungarian algorithm to optimally assign each in_channel to a tpl_channel
+	# by minimizing the total distances between their positions.
 	row_idx, col_idx = linear_sum_assignment(cost_matrix)
 	
+	# store the mapping result
 	new_idx = [[]]*30
+	fill_flags = [True]*30
 	for i in range(30):
 		if col_idx[i] < len(unassigned): # filter out dummy channels
 			tpl_channel = tpl_order[row_idx[i]]
 			in_channel = unassigned[col_idx[i]]
+			
+			new_idx[row_idx[i]] = [in_dict[in_channel]["index"]]
+			fill_flags[row_idx[i]] = False
 			tpl_dict[tpl_channel]["matched"] = True
 			in_dict[in_channel]["assigned"] = True
-			new_idx[row_idx[i]] = [in_dict[in_channel]["index"]]
-			
-			print(f'{tpl_order[row_idx[i]]}({row_idx[i]}) <- {unassigned[col_idx[i]]}({col_idx[i]})')
+			#print(f'{tpl_channel}({row_idx[i]}) <- {in_channel}({col_idx[i]})')
 	
-	# fill the missing_channels
-	missing_channels = [channel for channel in tpl_order if tpl_dict[channel]["matched"]==False]
+	# fill the remaining empty tpl_channels
+	missing_channels = get_empty_templates(tpl_order, tpl_dict)
 	if missing_channels != []:
 	    new_idx = find_neighbors(channel_info, missing_channels, new_idx)
 	
-	stage2_info.update({
-	    "newOrder" : new_idx,
-	    "unassignedInputs" : [channel for channel in in_order if in_dict[channel]["assigned"]==False]
-	})
+	mapping_data = {
+		"newOrder" : new_idx,
+		"fillFlags" : fill_flags
+	}
 	channel_info.update({
 	    "templateDict" : tpl_dict,
 	    "inputDict" : in_dict
 	})
+	return mapping_data, channel_info
+
+def mapping_result(stage1_info, stage2_info, channel_info, filename):
+	# 1. calculate how many times the model needs to be run
+	unassigned_num = len(stage1_info["unassignedInputs"])
+	batch_num = math.ceil(unassigned_num/30) + 1
+	
+	# 2. map the remaining in_channels
+	for i in range(1, batch_num):
+		# optimally select 30 in_channels to map to the tpl_channels based on proximity
+		new_mapping_data, channel_info = optimal_mapping(channel_info)
+		stage1_info["mappingData"] += [new_mapping_data]
+	
+	# 3. save the mapping result
+	new_dict = {
+		#"templateOrder" : channel_info["templateOrder"],
+		#"inputOrder" : channel_info["inputOrder"],
+		"batchNum" : batch_num,
+		"mappingData" : stage1_info["mappingData"]
+	}
+	with open(filename, 'w') as jsonfile:
+		jsonfile.write(json.dumps(new_dict))
 	
-	second2 = time.time()
-	print("The mapping process has been finished in", second2 - second1, "s.")
-	return stage2_info, channel_info
+	stage2_info["totalBatchNum"] = batch_num
+	return stage1_info, stage2_info, channel_info