update comment

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 
+import time
 import os
 import random
 import math
@@ -9,7 +10,7 @@ import mne
 from mne.channels import read_custom_montage
 
 import utils
-from channel_mapping import mapping_stage1, mapping_stage2, reorder_to_template, reorder_to_origin, find_neighbors
+from channel_mapping import mapping_stage1, mapping_stage2, reorder_input_data, restore_original_order, find_neighbors
 
 
 quickstart = """
@@ -20,7 +21,8 @@ quickstart = """
 
 ## Channel locations
 Upload your data's channel locations in `.loc` format, which can be obtained using **EEGLAB**.  
->If you cannot obtain it, we recommend you to download the standard montage <a href="">here</a>. If the channels in those files doesn't match yours, you can use **EEGLAB** to modify them to your needed montage.
+**Note:**
+If you cannot obtain it, we recommend you to download the standard montage <a href="">here</a>. If the channels in those files doesn't match yours, you can use **EEGLAB** to modify them to your needed montage.
 
 ## Mapping
 (...)
@@ -55,32 +57,33 @@ Electroencephalography (EEG) signals are often contaminated with artifacts. It i
 """
 
 init_js = """
-(app_state, channel_info) => {
-	app_state = JSON.parse(JSON.stringify(app_state));
+(app_info, channel_info) => {
+	app_info = JSON.parse(JSON.stringify(app_info));
 	channel_info = JSON.parse(JSON.stringify(channel_info));
+	stage1_info = app_info.stage1
 	
 	let selector, classname, attribute;
 	let channel, left, bottom;
 	
-	if(app_state.stage1State == "step2-selecting"){
+	if(stage1_info.state == "step2-selecting"){
 	    selector = "#radio-group > div:nth-of-type(2)";
 	    //classname = "radio";
 	    attribute = "value";
-	}else if(app_state.stage1State == "step3-selecting"){
+	}else if(stage1_info.state == "step3-selecting"){
 	    selector = "#chkbox-group > div:nth-of-type(2)";
 	     //classname = "chkbox";
 	    attribute = "name";
 	}else return;
 	
 	
-	// add figure of the mapping result
+	// add figure of the input montage
 	document.querySelector(selector).style.cssText = `
 		position: relative;
 		width: 100%;
 		aspect-ratio: 1;
 		//width: 560px;
 		//height: 560px;
-		background: url("file=${app_state.filenames.raw_montage}");
+		background: url("file=${stage1_info.filenames.input_montage}");
 		background-size: contain;
 
 	`;
@@ -99,13 +102,13 @@ init_js = """
 	
 	
 	// add indication for the missing channels
-	channel = app_state.missingTemplates[0];
+	channel = stage1_info.missingTemplates[0];
 	left = channel_info.templateDict[channel].css_position[0];
 	bottom = channel_info.templateDict[channel].css_position[1];
 	
 	let dot_rule = `
 		${selector}::before {
-			content: '';
+			content: "";
 			position: absolute;
 			background-color: red;
 			width: 10px;
@@ -145,14 +148,15 @@ init_js = """
 """
 
 update_js = """
-(app_state, channel_info) => {
-	app_state = JSON.parse(JSON.stringify(app_state));
+(app_info, channel_info) => {
+	app_info = JSON.parse(JSON.stringify(app_info));
 	channel_info = JSON.parse(JSON.stringify(channel_info));
+	stage1_info = app_info.stage1
 	
 	let selector;
 	let channel, left, bottom;
 	
-	if(app_state.stage1State == "step2-selecting"){
+	if(stage1_info.state == "step2-selecting"){
 	    selector = "#radio-group > div:nth-of-type(2)";
 	    
 	    // update the radios
@@ -166,12 +170,12 @@ update_js = """
 			item.className = "";
 			item.querySelector(":scope > span").innerText = "";
 		});
-	}else if(app_state.stage1State == "step3-selecting"){
+	}else if(stage1_info.state == "step3-selecting"){
 	    selector = "#chkbox-group > div:nth-of-type(2)";
 	}else return;
 	
 	// update indication
-	channel = app_state.missingTemplates[app_state["fillingCount"]-1];
+	channel = stage1_info.missingTemplates[stage1_info["fillingCount"]-1];
 	left = channel_info.templateDict[channel].css_position[0];
 	bottom = channel_info.templateDict[channel].css_position[1];
 	
@@ -202,7 +206,7 @@ update_js = """
 		}
 	`;
 	
-	// check if indicator already exist
+	// update the rules
 	const styleSheet = document.styleSheets[0];
 	for(let i=0; i<styleSheet.cssRules.length; i++){
 		let tmp = styleSheet.cssRules[i].selectorText;
@@ -219,7 +223,7 @@ update_js = """
 
 with gr.Blocks() as demo:
 
-	app_state_json = gr.JSON(visible=False)
+	app_info_json = gr.JSON(visible=False)
 	channel_info_json = gr.JSON(visible=False)
 
 	with gr.Row():
@@ -234,27 +238,25 @@ with gr.Blocks() as demo:
 			gr.Markdown("# 1.Channel Mapping")
 			# ------------------------input--------------------------
 			with gr.Row():
-				in_raw_data = gr.File(label="Raw data (.csv)", file_types=[".csv"])
-				in_raw_loc = gr.File(label="Channel locations (.loc, .locs)", file_types=[".loc", "locs"])
+				in_data_file = gr.File(label="Raw data (.csv)", file_types=[".csv"])
+				in_loc_file = gr.File(label="Channel locations (.loc, .locs, .xyz, .sfp, .txt)",
+									file_types=[".loc", "locs", ".xyz", ".sfp", ".txt"])
 			with gr.Row():
 				in_samplerate = gr.Textbox(label="Sampling rate (Hz)", scale=2)
-				map_btn = gr.Button("Mapping", scale=1)
+				map_btn = gr.Button("Mapping", interactive=False, scale=1)
 			
 			# ------------------------mapping------------------------
 			# description for stage1-123
-			desc_md = gr.Markdown("### Step1: Mapping result", visible=False) # """??? # test
-			
+			desc_md = gr.Markdown(visible=False)
 			# stage1-1 : mapping result
 			with gr.Row():
-				tpl_montage = gr.Image("./template_montage.png", label="Template montage", visible=False)
+				tpl_montage = gr.Image("./template_montage.png", label="Template channels", visible=False)
 				mapped_montage = gr.Image(label="Input channels", visible=False)
-			
 			# stage1-2 : assign 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
 			with gr.Row():
-				in_fill_mode = gr.Dropdown(choices=["mean", "zero"],
+				in_fillmode = gr.Dropdown(choices=["mean", "zero"],
 									value="mean",
 									label="Filling method",
 									visible=False,
@@ -263,7 +265,7 @@ with gr.Blocks() as demo:
 			chkbox_group = gr.CheckboxGroup(elem_id="chkbox-group", visible=False)
 			
 			with gr.Row():
-				clear_btn = gr.Button("Clear", visible=False) #, interactive=False
+				clear_btn = gr.Button("Clear", visible=False)
 				step2_btn = gr.Button("Next", visible=False)
 				step3_btn = gr.Button("Next", visible=False)
 				next_btn = gr.Button("Next step", visible=False)
@@ -273,13 +275,11 @@ with gr.Blocks() as demo:
 			gr.Markdown("# 2.Decode Data")
 			# ------------------------input--------------------------
 			with gr.Row():
-				in_model_name = gr.Dropdown(choices=[
+				in_modelname = gr.Dropdown(choices=[
 									("ART", "EEGART"),
 									("IC-U-Net", "ICUNet"),
 									("IC-U-Net++", "UNetpp"),
-									("IC-U-Net-Attn", "AttUnet"),
-									"(mapped data)",
-									"(denoised data)"],
+									("IC-U-Net-Attn", "AttUnet")],
 									value="EEGART",
 									label="Model",
 									scale=2)
@@ -287,8 +287,7 @@ with gr.Blocks() as demo:
 			
 			# ------------------------output-------------------------
 			batch_md = gr.Markdown(visible=False)
-			out_denoised_data = gr.File(label="Denoised data", visible=False)
-			
+			out_data_file = gr.File(label="Denoised data", visible=False)
 			# -------------------------------------------------------	
 				
 	with gr.Row():
@@ -302,59 +301,91 @@ with gr.Blocks() as demo:
 			gr.Markdown()
 		with gr.Tab("QuickStart"):
 			gr.Markdown(quickstart)
-    
-	#demo.load(js=tmp_js)
-	
-	# -------------------------stage1: channel mapping-------------------------------
-	def reset_all(raw_data, raw_loc, samplerate):
-		# verify that all required inputs have been provided
-		if raw_data == None or raw_loc == None:
-			gr.Warning('Please upload both the raw data and the channel location files.')
-			return
-		if samplerate == "":
-			gr.Warning('Please enter the sampling rate.')
-			return
-	    
-		# establish temp folder
-		filepath = os.path.dirname(str(raw_data))
+	
+	
+	# 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],
+	                    inputs = [in_data_file, in_loc_file, in_samplerate], outputs = map_btn)
+	def check_input(in_data, in_loc, samplerate):
+	    if in_data!=None and in_loc!=None and samplerate!="":
+	        return gr.Button(interactive=True)
+	    else:
+	        return gr.Button(interactive=False)
+	
+	
+	# +========================================================================================+
+	# |                                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))
 		try:
-			os.mkdir(filepath+"/temp_data/")
+			os.mkdir(filepath+"/session_data/")
 		except OSError as e:
-			utils.dataDelete(filepath+"/temp_data/")
-			os.mkdir(filepath+"/temp_data/")
-			#print(e)
+			utils.dataDelete(filepath+"/session_data/")
+			os.mkdir(filepath+"/session_data/")
+			print(e)
+		# establish new folders for stage1 and stage2
+		os.mkdir(filepath+"/session_data/stage1/")
+		os.mkdir(filepath+"/session_data/stage2/")
 		
-		# initialize channel_info, app_state
+		# initialize channel_info, app_info
 		channel_info = {}
-		app_state = {
-			"filepath": filepath+"/temp_data/",
-			"filenames": {},
-			"sampleRate": int(samplerate),
-			"stage1State" : "step1"
+		app_info = {
+			"rootFilepath" : filepath+"/session_data/",
+			"sampleRate" : int(samplerate),
+			#"currentStage" : "stage1",
+			"stage1" : {
+				"filepath" : filepath+"/session_data/stage1/",
+				"filenames" : {
+					"input_data" : in_data,
+					"input_loc" : in_loc,
+					"input_montage" : "",
+		    		"mapped_montage" : ""
+				},
+				"state" : None,
+				"fillingCount" : None,
+				"totalFillingNum" : None,
+				"newOrder" : None,
+	    		"unassignedInputs" : None,
+	    		"missingTemplates" : None
+			},
+			"stage2" : {
+				"filepath" : filepath+"/session_data/stage2/",
+				"filenames" : {
+					"output_data" : ""
+				},
+				#"state" : None,
+				"totalBatchNum" : None,
+				"newOrder" : None,
+	    		"unassignedInputs" : None
+			}
 		}
-		
 		# reset layout
-		return {app_state_json : app_state,
+		return {app_info_json : app_info,
 				channel_info_json : channel_info,
 				# ------------------stage1-----------------------
-				desc_md : gr.Markdown("### Step1: Mapping result", visible=False),
+				map_btn : gr.Button(interactive=False),
+				desc_md : gr.Markdown(visible=False),
 				tpl_montage : gr.Image(visible=False),
 				mapped_montage : gr.Image(value=None, visible=False),
 				radio_group : gr.Radio(choices=[], value=[], label="", visible=False),
-				in_fill_mode : gr.Dropdown(value="mean", 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),
 				step3_btn : gr.Button(visible=False),
-				next_btn : gr.Button(visible=False),
+				next_btn : gr.Button("Next step", visible=False),
 				# ------------------stage2-----------------------
 				run_btn : gr.Button(interactive=False),
 				batch_md : gr.Markdown(visible=False),
-				out_denoised_data : gr.File(visible=False)}
+				out_data_file : gr.File(visible=False)}
 	
 	
-	# ---------------------------stage1-1-------------------------------
+	# +========================================================================================+
+	# |                                        stage1-1                                        |
+	# +========================================================================================+
 	def save_figures(channel_info, filename1, filename2):
 		
 		template_montage = read_custom_montage("./template_chanlocs.loc")
@@ -363,6 +394,14 @@ with gr.Blocks() as demo:
 		template_order = channel_info["templateOrder"]
 		input_order = channel_info["inputOrder"]
 		
+		# get template and input's 2d coords
+		tpl_x = [template_dict[channel]["coord_2d"][0] for channel in template_order]
+		tpl_y = [template_dict[channel]["coord_2d"][1] for channel in template_order]
+		in_x = [input_dict[channel]["coord_2d"][0] for channel in input_order]
+		in_y = [input_dict[channel]["coord_2d"][1] for channel in input_order]
+		tpl_coords = np.vstack((tpl_x, tpl_y)).T
+		in_coords = np.vstack((in_x, in_y)).T
+		
 		# get template's head figure
 		tpl_fig = template_montage.plot()
 		tpl_ax = tpl_fig.axes[0]
@@ -373,14 +412,6 @@ with gr.Blocks() as demo:
 			head_lines.append((x,y))
 		plt.close()
 		
-		# get template's and input's 2d coords
-		tpl_x = [template_dict[channel]["coord_2d"][0] for channel in template_order]
-		tpl_y = [template_dict[channel]["coord_2d"][1] for channel in template_order]
-		in_x = [input_dict[channel]["coord_2d"][0] for channel in input_order]
-		in_y = [input_dict[channel]["coord_2d"][1] for channel in input_order]
-		tpl_coords = np.vstack((tpl_x, tpl_y)).T
-		in_coords = np.vstack((in_x, in_y)).T
-		
 		# -------------------------plot input montage------------------------------
 		fig = plt.figure(figsize=(6.4,6.4), dpi=100)
 		ax = fig.add_subplot(111)
@@ -391,12 +422,11 @@ with gr.Blocks() as demo:
 		# plot template's head
 		for x, y in head_lines:
 			ax.plot(x, y, color='black', linewidth=1.0)
-		# plot input channels
+		# plot input channels on it
 		ax.scatter(in_coords[:,0], in_coords[:,1], s=35, color='black')
 		for i, channel in enumerate(input_order):
 			ax.text(in_coords[i,0]+0.003, in_coords[i,1], channel, color='black', fontsize=10.0, va='center')
-		
-		# save raw_montage
+		# save input_montage
 		fig.savefig(filename1)
 		
 		# ---------------------------add indications-------------------------------
@@ -406,15 +436,14 @@ with gr.Blocks() as demo:
 		ax.scatter(in_coords[indices,0], in_coords[indices,1], s=35, color='red')
 		for i in indices:
 			ax.text(in_coords[i,0]+0.003, in_coords[i,1], input_order[i], color='red', fontsize=10.0, va='center')
-		
 		# save mapped_montage
 		fig.savefig(filename2)
-		plt.close()
 		
 		# -------------------------------------------------------------------------
-		# save the template and input channels' display position (in px).
+		# 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(template_order):
 			css_left = (tpl_coords[i,0]-11)/6.4
@@ -432,156 +461,189 @@ with gr.Blocks() as demo:
 		})
 		return channel_info
 	
-	def mapping_result(app_state, channel_info):
-		filepath = app_state["filepath"]
-		filename1 = filepath+"raw_montage_"+str(random.randint(1,10000))+".png"
+	def mapping_result(app_info, channel_info):
+		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)
-		
-		app_state["filenames"].update({
-		    "raw_montage" : filename1,
+		stage1_info["filenames"].update({
+		    "input_montage" : filename1,
 		    "mapped_montage" : filename2
 		})
 		
-		# ------------------determine the next step-----------------------
+		# -------------------------determine the next step--------------------------
 		
-		in_num = len(channel_info["inputOrder"])
-		matched_num = 30 - len(app_state["missingTemplates"])
+		input_num = len(channel_info["inputOrder"])
+		matched_num = 30 - len(stage1_info["missingTemplates"])
 		
-		# if the input channels(>=30) has all the 30 template channels
+		# if the in_channels has all the 30 tpl_channels (input_num>=30)
 		# -> stage2
 		if matched_num == 30:
-			app_state["stage1State"] = "finished"
+			stage1_info["state"] = "finished"
 			gr.Info('The mapping process has been finished.')
+			if input_num > 30:
+			    md = """
+			    ### Mapping result
+			    (...in red...)
+			    """
+			else:
+			    md = """
+			    ### Mapping result
+			    (...)
+			    """
 			
-			return {app_state_json : app_state,
+			app_info["stage1"] = stage1_info
+			return {app_info_json : app_info,
 					channel_info_json : channel_info,
-					desc_md : gr.Markdown("### Mapping result", visible=True),
+					map_btn : gr.Button(interactive=True),
+					desc_md : gr.Markdown(md, visible=True),
 					tpl_montage : gr.Image(visible=True),
 					mapped_montage : gr.Image(value=filename2, visible=True),
 					run_btn : gr.Button(interactive=True)}
 		
-		# if matched channels < 30, and there're still some unmatched input channels
-		# -> assign these input channels to nearby unmatched/empty template channels
-		if in_num > matched_num:
-			app_state["stage1State"] = "step2-initializing"
-		
-		# if input channels < 30, but all of them can match to some template channels
-		# -> directly use fill_mode to fill the remaining channels
-		elif in_num == matched_num:
-			app_state["stage1State"] = "step3-initializing"
-		
-		return {app_state_json : app_state,
-				channel_info_json : channel_info,
-				desc_md : gr.Markdown("### Step1: Mapping result", visible=True),
-				tpl_montage : gr.Image(visible=True),
-				mapped_montage : gr.Image(value=filename2, visible=True),
-				next_btn : gr.Button("Next step", visible=True)}
+		else:
+			# if matched_num < 30, and there're still some unmatched in_channels
+			# -> assign these in_channels to nearby unmatched tpl_channels
+			if input_num > matched_num:
+				stage1_info["state"] = "step2-initializing"
+				md = """
+				### Step1: Mapping result
+				(...in red...)
+				"""
+			
+			# if input_num < 30, but all of them can match to some tpl_channels
+			# -> directly use fillmode to fill the remaining tpl_channels
+			elif input_num == matched_num:
+				stage1_info["state"] = "step3-initializing"
+				md = """
+				### Step1: Mapping result
+				(...)
+				"""
+			
+			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_montage : gr.Image(visible=True),
+					mapped_montage : gr.Image(value=filename2, visible=True),
+					next_btn : gr.Button(visible=True)}
 	
-	map_btn.click(
+	start_stage1 = map_btn.click(
 	    fn = reset_all,
-	    inputs = [in_raw_data, in_raw_loc, in_samplerate],
-	    outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio_group,
-	                in_fill_mode, chkbox_group, fillmode_btn, clear_btn, step2_btn, step3_btn, next_btn,
-	                run_btn, batch_md, out_denoised_data]
+	    inputs = [in_data_file, in_loc_file, in_samplerate],
+	    outputs = [app_info_json, channel_info_json, map_btn, desc_md, tpl_montage, mapped_montage, 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_state_json, channel_info_json, in_raw_loc], 
-		outputs = [app_state_json, channel_info_json, desc_md]
-	    
+	    inputs = [app_info_json, channel_info_json], 
+		outputs = [app_info_json, channel_info_json, desc_md]
 	).success(
 	    fn = mapping_result,
-	    inputs = [app_state_json, channel_info_json],
-	    outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, next_btn, run_btn]
+	    inputs = [app_info_json, channel_info_json],
+	    outputs = [app_info_json, channel_info_json, map_btn, desc_md, tpl_montage, mapped_montage, next_btn, run_btn]
 	)
 	
-
-	def init_next_step(app_state, channel_info, selected_radio, selected_chkbox):
+	
+	# +========================================================================================+
+	# |                                 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 app_state["stage1State"] == "step2-initializing":
-			print('step1 -> step2')
-			app_state["missingTemplates"] = [channel for channel in channel_info["templateOrder"]
-											if channel_info["templateDict"][channel]["matched"]==False]
-			app_state.update({
-				"stage1State" : "step2-selecting",
+		if stage1_info["state"] == "step2-initializing":
+			#print('step1 -> step2')
+			md = """
+			### Step2: Assign unmatched input channels
+			(...)
+			"""
+			
+			# initialize the progress indication label for step2
+			stage1_info.update({
+				"state" : "step2-selecting",
 				"fillingCount" : 1,
-				"totalFillingNum" : len(app_state["missingTemplates"])
+				"totalFillingNum" : len(stage1_info["missingTemplates"])
 			})
+			name = stage1_info["missingTemplates"][0]
+			label = "{} (1/{})".format(name, stage1_info["totalFillingNum"])
 			
-			name = app_state["missingTemplates"][0]
-			label = name+' (1/'+str(app_state["totalFillingNum"])+')'
-			
-			if len(app_state["stage1UnassignedInputs"])==1 or app_state["totalFillingNum"]==1:
-				return {app_state_json : app_state,
+			app_info["stage1"] = stage1_info
+			# determine which button to display
+			if len(stage1_info["unassignedInputs"])==1 or stage1_info["totalFillingNum"]==1:
+				return {app_info_json : app_info,
 					    channel_info_json : channel_info,
-					    desc_md : gr.Markdown("### Step2: Assign unmatched input channels"),
+					    desc_md : gr.Markdown(md),
 					    tpl_montage : gr.Image(visible=False),
 					    mapped_montage : gr.Image(visible=False),
-					    radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
+					    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_state_json : app_state,
+			    return {app_info_json : app_info,
 					    channel_info_json : channel_info,
-					    desc_md : gr.Markdown("### Step2: Assign unmatched input channels"),
+					    desc_md : gr.Markdown(md),
 					    tpl_montage : gr.Image(visible=False),
 					    mapped_montage : gr.Image(visible=False),
-					    radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"], value=[], label=label, visible=True),
+					    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 app_state["stage1State"] == "step3-initializing":
-			print('step1 -> step3')
-			app_state["missingTemplates"] = [channel for channel in channel_info["templateOrder"]
-											if channel_info["templateDict"][channel]["matched"]==False]
-			app_state.update({
-				"stage1State" : "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: Fill the remaining template channels"),
+		elif stage1_info["state"] == "step3-initializing":
+			#print('step1 -> step3')
+			md = """
+			### Step3: Fill the remaining template channels
+			(...)
+			"""
+			return {desc_md : gr.Markdown(md),
 					tpl_montage : gr.Image(visible=False),
 					mapped_montage : gr.Image(visible=False),
-					in_fill_mode : gr.Dropdown(visible=True),
+					in_fillmode : gr.Dropdown(visible=True),
 					fillmode_btn : gr.Button(visible=True),
 					next_btn : gr.Button(visible=False)}
 		
 		# stage1-2 -> stage1-3 or stage2
-		elif app_state["stage1State"] == "step2-selecting":
+		elif stage1_info["state"] == "step2-selecting":
 		
-			# save info before clicking on next_btn
-			prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
-			prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
-			if selected_radio == []:
-				app_state["stage1NewOrder"][prev_target_idx] = []
-			else:
-				selected_idx = channel_info["inputDict"][selected_radio]["index"]
-				app_state["stage1NewOrder"][prev_target_idx] = [selected_idx]
+			# ----------------------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]
+				prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
+				
+				# store the index of the in_channel
+				selected_idx = channel_info["inputDict"][selected_radio]["index"]
+				stage1_info["newOrder"][prev_target_idx] = [selected_idx]
+				# 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)
+				print(prev_target_name, '<-', selected_radio)
 			
-			app_state.update({
-				"stage1UnassignedInputs" : [channel for channel in channel_info["inputOrder"]
-									if channel_info["inputDict"][channel]["assigned"]==False],
-				"missingTemplates" : [channel for channel in channel_info["templateOrder"]
-									if channel_info["templateDict"][channel]["matched"]==False]
-			})
+			# ------------------------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]
+			# 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]
 			
-			# if all the unmatched template channels were filled by input channels
+			# if all the unmatched tpl_channels were filled by in_channels
 			# -> stage2
-			if len(app_state["missingTemplates"]) == 0:
-				print('step2 -> stage2')
+			if len(stage1_info["missingTemplates"]) == 0:
+				#print('step2 -> stage2')
+				stage1_info["state"] = "finished"
 				gr.Info('The mapping process has been finished.')
-				app_state["stage1State"] = "finished"
 				
-				return {app_state_json : app_state,
+				app_info["stage1"] = stage1_info
+				return {app_info_json : app_info,
 						channel_info_json : channel_info,
 						desc_md : gr.Markdown(visible=False),
 						radio_group : gr.Radio(visible=False),
@@ -591,42 +653,43 @@ with gr.Blocks() as demo:
 						
 			# -> stage1-3
 			else:
-				print('step2 -> step3')
-				app_state.update({
-					"stage1State" : "step3-initializing",
-					"fillingCount" : 1,
-					"totalFillingNum" : len(app_state["missingTemplates"])
-				})
-				return {app_state_json : app_state,
+				#print('step2 -> step3')
+				md = """
+				### Step3: Fill the remaining template channels
+				(...)
+				"""
+				
+				app_info["stage1"] = stage1_info
+				return {app_info_json : app_info,
 						channel_info_json : channel_info,
-						desc_md : gr.Markdown("### Step3: Fill the remaining template channels"),
+						desc_md : gr.Markdown(md),
 						radio_group : gr.Radio(visible=False),
-						in_fill_mode : gr.Dropdown(visible=True),
+						in_fillmode : gr.Dropdown(visible=True),
 						fillmode_btn : gr.Button(visible=True),
 						clear_btn : gr.Button(visible=False),
 						next_btn : gr.Button(visible=False)}
 		
 		# stage1-3 -> stage2
-		elif app_state["stage1State"] == "step3-selecting":
-			
-			# save info before clicking on next_btn
-			prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
-			prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
-			if selected_chkbox == []:
-				app_state["stage1NewOrder"][prev_target_idx] = []
-			else:
-				selected_idx = [channel_info["inputDict"][channel]["index"] for channel in selected_chkbox]
-				app_state["stage1NewOrder"][prev_target_idx] = selected_idx				
-			#print(f'{prev_target_name}({prev_target_idx}): {selected_chkbox}')
-			
+		elif stage1_info["state"] == "step3-selecting":
+			#print('step3 -> stage2')
+			stage1_info["state"] = "finished"
 			gr.Info('The mapping process has been finished.')
-			app_state["stage1State"] = "finished"
-			print('step3 -> stage2')
 			
-			app_state["missingTemplates"] = [channel for channel in channel_info["templateOrder"]
-											if channel_info["templateDict"][channel]["matched"]==False]
+			# ----------------------store information before the button click----------------------
+			
+			# check if the user has not unchecked all in_channel checkboxes
+			if selected_chkbox != []:
+				prev_target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
+				prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
+				
+				# 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}')
+			# -------------------------------------------------------------------------------------
 			
-			return {app_state_json : app_state,
+			app_info["stage1"] = stage1_info
+			return {app_info_json : app_info,
 					desc_md : gr.Markdown(visible=False),
 					chkbox_group : gr.CheckboxGroup(visible=False),
 					next_btn : gr.Button(visible=False),
@@ -634,62 +697,71 @@ with gr.Blocks() as demo:
 	
 	next_btn.click(
 		fn = init_next_step,
-		inputs = [app_state_json, channel_info_json, radio_group, chkbox_group],
-		outputs = [app_state_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio_group,
-		            in_fill_mode, chkbox_group, fillmode_btn, clear_btn, step2_btn, next_btn, run_btn]
+		inputs = [app_info_json, channel_info_json, radio_group, chkbox_group],
+		outputs = [app_info_json, channel_info_json, desc_md, tpl_montage, mapped_montage, radio_group,
+		            in_fillmode, chkbox_group, fillmode_btn, clear_btn, step2_btn, next_btn, run_btn]
 	).success(
 	    fn = None,
 	    js = init_js,
-	    inputs = [app_state_json, channel_info_json],
+	    inputs = [app_info_json, channel_info_json],
 		outputs = []
 	)
 	
-	# ---------------------------stage1-2-------------------------------
-	def update_radio(app_state, channel_info, selected):
+	
+	# +========================================================================================+
+	# |                                        stage1-2                                        |
+	# +========================================================================================+
+	def update_radio(app_info, channel_info, selected):
+		stage1_info = app_info["stage1"]
 		
-		# save info before clicking on next_btn
-		prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
-		prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
-		if selected == []:
-			app_state["stage1NewOrder"][prev_target_idx] = []
-		else:
-			selected_idx = channel_info["inputDict"][selected]["index"]
-			app_state["stage1NewOrder"][prev_target_idx] = [selected_idx]
+		# ----------------------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]
+			prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
+			
+			# store the index of the selected in_channel
+			selected_idx = channel_info["inputDict"][selected]["index"]
+			stage1_info["newOrder"][prev_target_idx] = [selected_idx]
+			# 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-------------------------
+		stage1_info["fillingCount"] += 1
 		
-		# update the current round
-		app_state["fillingCount"] += 1
-		app_state["stage1UnassignedInputs"] = [channel for channel in channel_info["inputOrder"]
+		# 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]
+		# 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"])
 		
-		target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
-		radio_label = target_name+' ('+str(app_state["fillingCount"])+'/'+str(app_state["totalFillingNum"])+')'
-				
-		if len(app_state["stage1UnassignedInputs"])==1 or app_state["fillingCount"]==app_state["totalFillingNum"]:
-			return {app_state_json : app_state,
+		app_info["stage1"] = stage1_info
+		# determine which button to display
+		if len(stage1_info["unassignedInputs"])==1 or stage1_info["fillingCount"]==stage1_info["totalFillingNum"]:
+			return {app_info_json : app_info,
 					channel_info_json : channel_info,
-					radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"],
+					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)}
 		else:
-			return {app_state_json : app_state,
+			return {app_info_json : app_info,
 					channel_info_json : channel_info,
-					radio_group : gr.Radio(choices=app_state["stage1UnassignedInputs"],
+					radio_group : gr.Radio(choices=stage1_info["unassignedInputs"],
 					                        value=[], label=radio_label)}
 	
 	step2_btn.click(
 	    fn = update_radio,
-	    inputs = [app_state_json, channel_info_json, radio_group],
-	    outputs = [app_state_json, channel_info_json, radio_group, step2_btn, next_btn]
-	    
+	    inputs = [app_info_json, channel_info_json, radio_group],
+	    outputs = [app_info_json, channel_info_json, radio_group, step2_btn, next_btn]
 	).success(
 	    fn = None,
 	    js = update_js,
-	    inputs = [app_state_json, channel_info_json],
+	    inputs = [app_info_json, channel_info_json],
 		outputs = []
 	)
 	
@@ -700,183 +772,261 @@ with gr.Blocks() as demo:
 	)
 	
 	
-	# ---------------------------stage1-3-------------------------------
-	def fill_value(app_state, channel_info, fill_mode):
+	# +========================================================================================+
+	# |                                        stage1-3                                        |
+	# +========================================================================================+
+	def fill_value(app_info, channel_info, fillmode):
+		stage1_info = app_info["stage1"]
 		
-		if fill_mode == 'zero':
-			app_state["stage1State"] = "finished"
+		if fillmode == "zero":
+			stage1_info["state"] = "finished"
 			gr.Info('The mapping process has been finished.')
 			
-			return {app_state_json : app_state,
+			app_info["stage1"] = stage1_info
+			return {app_info_json : app_info,
 					desc_md : gr.Markdown(visible=False),
-					in_fill_mode : gr.Dropdown(visible=False),
+					in_fillmode : gr.Dropdown(visible=False),
 					fillmode_btn : gr.Button(visible=False),
 					run_btn : gr.Button(interactive=True)}
 		
-		elif fill_mode == 'mean':
-			app_state["stage1State"] = "step3-selecting"
-			app_state = find_neighbors(app_state, channel_info)
+		elif fillmode == "mean":
+			md = """
+			### Step3: Fill the remaining template channels
+			(...)
+			"""
 			
-			# init stage1-3-selecting
-			target_name = app_state["missingTemplates"][0]
-			target_idx = channel_info["templateDict"][target_name]["index"]
+			# 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"])
+			})
 			
-			chkbox_value = app_state["stage1NewOrder"][target_idx]
+			# initialize the progress indicator 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 = target_name+' (1/'+str(app_state["totalFillingNum"])+')'
+			chkbox_label = "{} (1/{})".format(target_name, stage1_info["totalFillingNum"])
 			
-			if app_state["totalFillingNum"] == 1:
-			    return {app_state_json : app_state,
-					    in_fill_mode : gr.Dropdown(visible=False),
+			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_state_json : app_state,
-					    in_fill_mode : gr.Dropdown(visible=False),
+			    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)}
 	
-	def update_chkbox(app_state, channel_info, selected):
-	
-		# save info before clicking on next_btn
-		prev_target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
-		prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
-		if selected == []:
-			app_state["stage1NewOrder"][prev_target_idx] = []
-		else:
-			selected_idx = [channel_info["inputDict"][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))
+	def update_chkbox(app_info, channel_info, selected):
+		stage1_info = app_info["stage1"]
 		
-		# update the current round
-		app_state["fillingCount"] += 1
+		# ----------------------store information before the button click----------------------
 		
-		target_name = app_state["missingTemplates"][app_state["fillingCount"]-1]
-		target_idx = channel_info["templateDict"][target_name]["index"]
+		# check if the user has not unchecked all in_channel checkboxes
+		if selected != []:
+			prev_target_name = stage1_info["missingTemplates"][stage1_info["fillingCount"]-1]
+			prev_target_idx = channel_info["templateDict"][prev_target_name]["index"]
+			
+			# 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))
+		
+		# ------------------------update information for the new round-------------------------
+		stage1_info["fillingCount"] += 1
 		
-		chkbox_value = app_state["stage1NewOrder"][target_idx]
+		# 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 = [channel_info["inputOrder"][i] for i in chkbox_value]
-		chkbox_label = target_name+' ('+str(app_state["fillingCount"])+'/'+str(app_state["totalFillingNum"])+')'
+		chkbox_label = "{} ({}/{})".format(target_name, stage1_info["fillingCount"], stage1_info["totalFillingNum"])
 		
-		if app_state["fillingCount"] == app_state["totalFillingNum"]:
-			return {app_state_json : app_state,
+		app_info["stage1"] = stage1_info
+		# determine which button to display
+		if stage1_info["fillingCount"] == stage1_info["totalFillingNum"]:
+			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)}
 		else:
-			return {app_state_json : app_state,
+			return {app_info_json : app_info,
 					chkbox_group : gr.CheckboxGroup(value=chkbox_value, label=chkbox_label)}
 	
 	fillmode_btn.click(
 	    fn = fill_value,
-	    inputs = [app_state_json, channel_info_json, in_fill_mode],
-		outputs = [app_state_json, desc_md, in_fill_mode, fillmode_btn, chkbox_group, step3_btn, next_btn, run_btn]
+	    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]
 	).success(
 	    fn = None,
 	    js = init_js,
-	    inputs = [app_state_json, channel_info_json],
+	    inputs = [app_info_json, channel_info_json],
 		outputs = []
 	)
 	
 	step3_btn.click(
 	    fn = update_chkbox,
-	    inputs = [app_state_json, channel_info_json, chkbox_group],
-	    outputs = [app_state_json, chkbox_group, step3_btn, next_btn]
-	    
+	    inputs = [app_info_json, channel_info_json, chkbox_group],
+	    outputs = [app_info_json, chkbox_group, step3_btn, next_btn]
 	).success(
 	    fn = None,
 	    js = update_js,
-	    inputs = [app_state_json, channel_info_json],
+	    inputs = [app_info_json, channel_info_json],
 		outputs = []
 	)
 	
-    # -------------------------stage2: decode data-------------------------------
-	def delete_file(filename):
-		try:
-			os.remove(filename)
-		except OSError as e:
-			print(e)
 	
-	def reset_run(app_state, channel_info, raw_data, model_name):
+    # +========================================================================================+
+	# |                                  stage2: decode data                                   |
+	# +========================================================================================+
+	def reset_run(app_info, channel_info, modelname):
+		stage1_info = app_info["stage1"]
+		stage2_info = app_info["stage2"]
 	    
-	    # reset in.assigned back to the state after stage1
-		for channel in app_state["stage1UnassignedInputs"]:
-			channel_info["inputDict"][channel]["assigned"] = False
-		
-		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'
+	    # delete the previous folder of stage2 if it exists
+		filepath = stage2_info["filepath"]
+		utils.dataDelete(filepath)
+		# establish a new folder for stage2
+		new_filepath = app_info["rootFilepath"]+"stage2_"+str(random.randint(1,10000))+"/"
+		os.mkdir(new_filepath)
+		# generate the output filename
+		filename = stage1_info["filenames"]["input_data"]
+		filename = os.path.basename(str(filename))
+		new_filename = os.path.splitext(filename)[0]+'_'+modelname+'.csv'
 		
-		in_num = len(channel_info["inputOrder"])
-		assigned_num = len(app_state["stage1UnassignedInputs"])
-		batch_num = math.ceil((in_num-assigned_num)/30) + 1
+	    # 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_state["filenames"]["denoised"] = filepath + output_name
-		app_state.update({
-			"runningState" : "stage1",
-			"batchCount" : 1,
-			"totalBatchNum" : batch_num,
-			"stage2UnassignedInputs" : app_state["stage1UnassignedInputs"],
-			"stage2NewOrder" : [[]]*30,
+		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"]
+			}
 		})
-		return {app_state_json : app_state,
+		return {app_info_json : app_info,
 				channel_info_json : channel_info,
-				run_btn : gr.Button(interactive=False),
+				#run_btn : gr.Button(interactive=False),
 				batch_md : gr.Markdown(visible=False),
-				out_denoised_data : gr.File(visible=False)}
+				out_data_file : gr.File(visible=False)}
 	
-	def run_model(app_state, channel_info, raw_data, model_name):
-		filepath = app_state["filepath"]
-		samplerate = app_state["sampleRate"]
-		new_filename = app_state["filenames"]["denoised"]
+	def run_model(app_info, channel_info, modelname):
+		stage1_info = app_info["stage1"]
+		stage2_info = app_info["stage2"]
 		
-		while app_state["runningState"] != "finished":
-			md = 'Running model('+str(app_state["batchCount"])+'/'+str(app_state["totalBatchNum"])+')...'
-			yield {batch_md : gr.Markdown(md, visible=True)}
-			
-			if app_state["batchCount"] > 1:
-				app_state, channel_info = mapping_stage2(app_state, channel_info)
-				if app_state["runningState"] == "finished":
-					#yield {batch_md : gr.Markdown("error", visible=True)}
-					break
-			
-			reorder_to_template(app_state, raw_data)
-			# step1: Data preprocessing
-			total_file_num = utils.preprocessing(filepath, 'mapped.csv', samplerate)
-			# step2: Signal reconstruction
-			utils.reconstruct(model_name, total_file_num, filepath, 'denoised.csv', samplerate)
-			reorder_to_origin(app_state, channel_info, new_filename)
+		filepath = stage2_info["filepath"]
+		samplerate = app_info["sampleRate"]
+		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
+		break_flag = False
+		
+		# run the model multiple times until all in_channels are reconstructed
+		for i in range(stage2_info["totalBatchNum"]):
+			# establish a temp folder
+			try:
+				os.mkdir(filepath+"temp_data/")
+			#except FileExistsError:
+				#utils.dataDelete(filepath+"temp_data/")
+				#os.mkdir(filepath+"temp_data/")
+			except FileNotFoundError:
+				#print('break1!!')
+				break_flag = True
+				break
+			except OSError as e:
+				print(e)
 			
-			#if model_name == "(mapped data)":
-				#return {out_denoised_data : filepath + 'mapped.csv'}
-			#elif model_name == "(denoised data)":
-				#return {out_denoised_data : filepath + 'denoised.csv'}
+			# update the running status
+			md = "Running model({}/{})...".format(i+1, stage2_info["totalBatchNum"])
+			yield {batch_md : gr.Markdown(md, visible=True)}
 			
-			delete_file(filepath+'mapped.csv')
-			delete_file(filepath+'denoised.csv')
-			app_state["batchCount"] += 1
+			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"]))
 			
-		yield {run_btn : gr.Button(interactive=True),
-				batch_md : gr.Markdown(visible=False),
-				out_denoised_data : gr.File(new_filename, visible=True)}
+			# ----------------------------------------------------------------------
+			try:
+				# step1: Reorder input data
+				reorder_input_data(new_idx, filename, filepath+"temp_data/mapped.csv")
+				# 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)
+				# step4: Restore original order
+				restore_original_order(channel_info, i, new_idx, filepath+"temp_data/denoised.csv", new_filename)
+			except FileNotFoundError:
+				#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)}
+		else:
+			yield {#run_btn : gr.Button(interactive=True),
+					batch_md : gr.Markdown(visible=False),
+					out_data_file : gr.File(new_filename, visible=True)}
 	
 	run_btn.click(
 		fn = reset_run,
-		inputs = [app_state_json, channel_info_json, in_raw_data, in_model_name],
-		outputs = [app_state_json, channel_info_json, run_btn, batch_md, out_denoised_data]
+		inputs = [app_info_json, channel_info_json, in_modelname],
+		outputs = [app_info_json, channel_info_json, run_btn, batch_md, out_data_file]
 		
 	).success(
 		fn = run_model,
-		inputs = [app_state_json, channel_info_json, in_raw_data, in_model_name],
-		outputs = [run_btn, batch_md, out_denoised_data]
+		inputs = [app_info_json, channel_info_json, in_modelname],
+		outputs = [run_btn, batch_md, out_data_file]
 	)
 	
 if __name__ == "__main__":
     demo.launch()
+
+
+"""
+--------
+	|----(inputname).csv
+	|----session_data
+		|----stage1
+			|----input_montage.png
+			|----mapped_montage.png
+		|----stage2_(...)
+			|----temp_data
+				|----mapped.csv
+				|----denoised.csv
+				|----temp2
+					|...
+			|----(outputname).csv
+"""
+

channel_mapping.py

@@ -10,13 +10,10 @@ from scipy.interpolate import Rbf
 from scipy.optimize import linear_sum_assignment
 from sklearn.neighbors import NearestNeighbors
 
-def reorder_to_template(app_state, filename):
-	old_idx = app_state["stage1NewOrder"] if app_state["runningState"]=="stage1" else app_state["stage2NewOrder"]
+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])) # reordered raw data
-	new_filename = app_state["filepath"]+'mapped.csv'
-	#print('new order 1:', app_state["stage1NewOrder"])
-	#print('new order 2:', app_state["stage2NewOrder"])
+	new_data = np.zeros((30, old_data.shape[1])) # to store reordered raw data
+	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)
@@ -31,25 +28,25 @@ def reorder_to_template(app_state, filename):
 			tmp_data = [old_data[j, :] for j in idx_set]
 			new_data[i, :] = np.mean(tmp_data, axis=0)
 	
-	print('old.shape, new.shape: ', old_data.shape, new_data.shape)
+	old_shape = (old_data.shape[0]-1, old_data.shape[1])
+	print('old.shape, new.shape: ', old_shape, new_data.shape)
 	utils.save_data(new_data, new_filename)
 	return
 
-def reorder_to_origin(app_state, channel_info, new_filename):
-	filename = app_state["filepath"]+'denoised.csv'
-	old_idx = app_state["stage1NewOrder"] if app_state["runningState"]=="stage1" else app_state["stage2NewOrder"]
+def restore_original_order(channel_info, cnt, old_idx, filename, new_filename):
 	old_data = utils.read_train_data(filename) # denoised data
 	template_order = channel_info["templateOrder"]
+	input_order = channel_info["inputOrder"]
 	
-	if app_state["runningState"] == "stage1":
-		new_data = np.zeros((len(channel_info["inputOrder"]), old_data.shape[1]))
+	if cnt == 0:
+		new_data = np.zeros((len(input_order), old_data.shape[1]))
 	else:
 		new_data = utils.read_train_data(new_filename)
 	
 	for i, channel in enumerate(template_order):
 		idx_set = old_idx[i]
 		
-		# ignore if this channel was filled with fill_mode ('mean' or 'zero')
+		# 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, :]
 	
@@ -97,9 +94,9 @@ def align_coords(channel_info, template_montage, input_montage):
 	
 	
 	# --------------------------------2-D------------------------------------
-	# (for the indication of missing template channel's position when fill_mode:'mean')
+	# (for the indicate the location missing template channel's position when fill_mode:'mean')
 	
-	fig = [template_montage.plot(), input_montage.plot()]	
+	fig = [template_montage.plot(), input_montage.plot()]
 	ax = [fig[0].axes[0], fig[1].axes[0]]
 	
 	# get the original coords
@@ -154,48 +151,38 @@ def align_coords(channel_info, template_montage, input_montage):
 	})
 	return channel_info
 
-def find_neighbors(app_state, channel_info):
-	new_idx = app_state["stage1NewOrder"] if app_state["runningState"]=="stage1" else app_state["stage2NewOrder"]
+def find_neighbors(channel_info, missing_channels, new_idx):
 	template_dict = channel_info["templateDict"]
 	input_dict = channel_info["inputDict"]
-	template_order = channel_info["templateOrder"]
 	input_order = channel_info["inputOrder"]
-	missing_channels = app_state["missingTemplates"]
-	if missing_channels == []:
-		return app_state # change nothing
 	
-	
-	in_coords = [input_dict[channel]["coord_3d"] for channel in input_order]
-	in_coords = np.array([in_coords[i] for i in range(len(in_coords))])
+	all_in = [np.array(input_dict[channel]["coord_3d"]) for channel in input_order]
+	missing_tpl = [np.array(template_dict[channel]["coord_3d"]) for channel in missing_channels]
 	
 	# use KNN to choose k nearest channels
 	k = 4 if len(input_order)>4 else len(input_order)
 	knn = NearestNeighbors(n_neighbors=k, metric='euclidean')
-	knn.fit(in_coords)
+	knn.fit(all_in)
 	
-	for channel in missing_channels:
-		distances, indices = knn.kneighbors(np.array(template_dict[channel]["coord_3d"]).reshape(1,-1))
-		selected = [input_order[i] for i in indices[0]]
+	for i, channel in enumerate(missing_channels):
+		distances, indices = knn.kneighbors(missing_tpl[i].reshape(1,-1))
+		#selected = [input_order[j] for j in indices[0]]
 		#print(channel, ':', selected)
 		
 		idx = template_dict[channel]["index"]
 		new_idx[idx] = indices[0].tolist()
-	
-	if app_state["runningState"] == "stage1":
-		app_state["stage1NewOrder"] = new_idx
-	else:
-		app_state["stage2NewOrder"] = new_idx
     
-	return app_state
+	return new_idx
 
-def mapping_stage1(app_state, channel_info, loc_file):
-	yield app_state, channel_info, gr.Markdown("Mapping...", visible=True) # app_state, channel_info???
+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"]
 	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
+	new_idx = [[]]*30 # store the indices of the in_channels in the order of tpl_channls
 	alias_dict = {
 		'T3': 'T7',
 		'T4': 'T8',
@@ -203,21 +190,20 @@ def mapping_stage1(app_state, channel_info, loc_file):
 		'T6': 'P8'
 	}
 	
-	# match the names of input channels -> template channels
+	# match the names of input channels and template channels
 	for i, channel in enumerate(template_order):
 		if channel in alias_dict and alias_dict[channel] in input_dict:
-			template_montage.rename_channels({channel: alias_dict[channel]})
+			template_montage.rename_channels({channel: alias_dict[channel]}) # rename the current tpl_channel
 			template_dict[alias_dict[channel]] = template_dict.pop(channel)
 			channel = alias_dict[channel]
-		
+			
 		if channel in input_dict:
 			new_idx[i] = [input_dict[channel]["index"]]
 			template_dict[channel]["matched"] = True
 			input_dict[channel]["assigned"] = True
 	
-	# update names
+	# update the names
 	template_order = template_montage.ch_names
-	input_order = input_montage.ch_names
 	
 	channel_info.update({
 	    "templateDict" : template_dict,
@@ -225,10 +211,9 @@ def mapping_stage1(app_state, channel_info, loc_file):
 	    "templateOrder" : template_order,
 		"inputOrder" : input_order
 	})
-	app_state.update({
-	    "stage1NewOrder" : new_idx,
-	    "runningState" : "stage1",
-	    "stage1UnassignedInputs" : [channel for channel in input_order if input_dict[channel]["assigned"]==False],
+	app_info["stage1"].update({
+	    "newOrder" : new_idx,
+	    "unassignedInputs" : [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]
 	})
 	
@@ -237,19 +222,16 @@ def mapping_stage1(app_state, channel_info, loc_file):
 	
 	second2 = time.time()
 	print('Mapping (stage1) finished in',second2 - second1,'s.')
-	yield app_state, channel_info, gr.Markdown("", visible=False)
+	yield app_info, channel_info, gr.Markdown("", visible=False)
 
-def mapping_stage2(app_state, channel_info):
+def mapping_stage2(stage2_info, channel_info):
 	second1 = time.time()
 
 	template_dict = channel_info["templateDict"]
 	input_dict = channel_info["inputDict"]
 	template_order = channel_info["templateOrder"]
 	input_order = channel_info["inputOrder"]
-	unassigned = app_state["stage2UnassignedInputs"]
-	if unassigned == []:
-		app_state["runningState"] = "finished"
-		return app_state, channel_info
+	unassigned = stage2_info["unassignedInputs"]
 	
 	tpl_coords = np.array([template_dict[channel]["coord_3d"] for channel in template_order])
 	unassigned_coords = np.array([input_dict[channel]["coord_3d"] for channel in unassigned])
@@ -274,31 +256,29 @@ def mapping_stage2(app_state, channel_info):
 	new_idx = [[]]*30
 	for i in range(30):
 		if col_idx[i] < len(unassigned): # filter out dummy channels
-			print(f'({row_idx[i]}, {col_idx[i]})')
-			
 			tpl_channel = template_order[row_idx[i]]
 			in_channel = unassigned[col_idx[i]]
 			template_dict[tpl_channel]["matched"] = True
 			input_dict[in_channel]["assigned"] = True
 			new_idx[row_idx[i]] = [input_dict[in_channel]["index"]]
 			
-			print(template_order[row_idx[i]], '<-', unassigned[col_idx[i]])
+			print(f'{template_order[row_idx[i]]}({row_idx[i]}) <- {unassigned[col_idx[i]]}({col_idx[i]})')
 	
+	# fill the missing_channels
+	missing_channels = [channel for channel in template_order if template_dict[channel]["matched"]==False]
+	if missing_channels != []:
+	    new_idx = find_neighbors(channel_info, missing_channels, new_idx)
+	
+	stage2_info.update({
+	    "newOrder" : new_idx,
+	    "unassignedInputs" : [channel for channel in input_order if input_dict[channel]["assigned"]==False]
+	})
 	channel_info.update({
 	    "templateDict" : template_dict,
 	    "inputDict" : input_dict
 	})
-	app_state.update({
-	    "stage2NewOrder" : new_idx,
-	    "runningState" : "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 missing_channels
-	app_state = find_neighbors(app_state, channel_info)
 	
 	second2 = time.time()
-	print(f'Mapping (stage2-{app_state["batchCount"]-1}) finished in {second2 - second1}s.')
-	return app_state, channel_info
+	print("The mapping process has been finished in", second2 - second1, "s.")
+	return stage2_info, channel_info
 

utils.py

@@ -98,7 +98,7 @@ def cut_data(filepath, raw_data):
     total = int(len(raw_data[0]) / 1024)
     for i in range(total):
         table = raw_data[:, i * 1024:(i + 1) * 1024]
-        filename = filepath + '/temp2/' + str(i) + '.csv'
+        filename = filepath + 'temp2/' + str(i) + '.csv'
         with open(filename, 'w', newline='') as csvfile:
             writer = csv.writer(csvfile)
             writer.writerows(table)
@@ -213,10 +213,10 @@ def decode_data(data, std_num, mode=5):
 def preprocessing(filepath, filename, samplerate):
     # establish temp folder
     try:
-        os.mkdir(filepath+"/temp2/")
+        os.mkdir(filepath+"temp2/")
     except OSError as e:
-        dataDelete(filepath+"/temp2/")
-        os.mkdir(filepath+"/temp2/")
+        dataDelete(filepath+"temp2/")
+        os.mkdir(filepath+"temp2/")
         print(e)
     
     # read data
@@ -239,7 +239,7 @@ def reconstruct(model_name, total, filepath, outputfile, samplerate):
     # -------------------decode_data---------------------------
     second1 = time.time()
     for i in range(total):
-        file_name = filepath + '/temp2/{}.csv'.format(str(i))
+        file_name = filepath + 'temp2/{}.csv'.format(str(i))
         data_noise = read_train_data(file_name)
         
         std = np.std(data_noise)
@@ -251,17 +251,18 @@ def reconstruct(model_name, total, filepath, outputfile, samplerate):
         d_data = decode_data(data_noise, std, model_name)
         d_data = d_data[0]
 
-        outputname = filepath + '/temp2/output{}.csv'.format(str(i))
+        outputname = filepath + 'temp2/output{}.csv'.format(str(i))
         save_data(d_data, outputname)
 
     # --------------------glue_data----------------------------
-    signal = glue_data(filepath+"/temp2/", total, filepath+outputfile)
+    signal = glue_data(filepath+"temp2/", total, filepath+outputfile)
     #print(signal.shape)
     # -------------------delete_data---------------------------
-    dataDelete(filepath+"/temp2/")
+    dataDelete(filepath+"temp2/")
     # --------------------resample-----------------------------
     signal = resample_(signal, 256, samplerate) # 256Hz -> original sampling rate
     #print(signal.shape)
+    # --------------------save_data----------------------------
     save_data(signal, filepath+outputfile)
     second2 = time.time()