update

app.py

@@ -6,41 +6,44 @@ import random
 
 readme = """
 
-## Introduction
-This tool is designed to denoise your EEG data using our pre-trained models: **ART**, **IC-U-Net**, **IC-U-Net++**, and **IC-U-Net-Attn**. The models was trained using the EEG signals of 30 channels, including: ``Fp1, Fp2, F7, F3, Fz, F4, F8, FT7, FC3, FCz, FC4, FT8, T7, C3, Cz, C4, T8, TP7, CP3, CPz, CP4, TP8, P7, P3, Pz, P4, P8, O1, Oz, O2``.So it is essential to map your channels to our template channels before using the models. This guide will walk you through the steps of mapping your channels, ensuring proper alignment with our template, and preparing your data for optimal processing.
+This tool serves two main purposes:
+1. **Channel Mapping**: Align your EEG channels with our template channels to ensure compatibility with our models.
+2. **EEG Data Denoising**: Use our pre-trained models—**ART**, **IC-U-Net**, **IC-U-Net++**, and **IC-U-Net-Attn**—to denoise your EEG data.
+
+## File Requirements and Preparation
+- **Channel locations**: If you don't have the channel location file, we recommend you to download the standard montage <a href="">here</a>. If the channels in those files don't match yours, you can use **EEGLAB** to adjust them to your required montage.
+- **Raw data**: The data need to be a two-dimentional array (channel, timepoint).
+- **Channel requirements**: Your data must include some channels that correspond to our template channels, including: ``Fp1, Fp2, F7, F3, Fz, F4, F8, FT7, FC3, FCz, FC4, FT8, T7, C3, Cz, C4, T8, TP7, CP3, CPz, CP4, TP8, P7, P3, Pz, P4, P8, O1, Oz, O2``. At least some of them need to be present for successful mapping.
+- **Channel removal**: Before uploading your files, please remove any reference, ECG, EOG, EMG... channels.
 
 ## 1. Channel Mapping
-The following steps will guide you through the process of mapping your EEG channels to our template channels, to ensure compatibility before denoising.
+The following steps will guide you through the process of mapping your EEG channels to our template channels.
 
 ### Step1: Initial Matching and Rescaling
 After clicking on ``Mapping`` button, we will first match your channels to our template channels by their names. Using the matched channels as reference points, we will apply Thin Plate Spline (TPS) transformation to rescale your montage to align with our template's scale. The template montage and your rescaled montage will be displayed side by side for comparison. Channels that do not have a match in our template will be **highlighted in red**.
-- If your data includes all the 30 template channels, you can proceed directly to **run the models**.
+- If your data includes all the 30 template channels, you will be directed to **Mapping Results**.
 - If your data doesn't include all the 30 template channels and you have some channels that do not match the template, you will be directed to **Step2**.
 - If all your channels are included in our template but you have fewer than 30 channels, you will be directed to **Step3**.
 
 ### Step2: Forwarding Unmatched Channels
 In this step, you will handle the channels that didn't have a direct match with our template, by manually assigning them to the template channels that are still empty, ensuring the most efficient use of your data.  
 Your unmatched channels, previously highlighted in red, will be shown on your montage with a radio button displayed above each. You can choose to forward the data from these unmatched channels to the empty template channels. The interface will display each empty template channel in sequence, allowing you to select which of your unmatched channels to forward.
-- If all empty template channels are filled by your selections, you can proceed to **run the models**.
+- If all empty template channels are filled by your selections, you will be directed to **Mapping Results**.
 - If there are still empty template channels remaining, you will be directed to **Step3**.
 
 ### 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 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.
-
+Choose the 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 will be directed to **Mapping Results**.
 
-### Mapping Result
+### Mapping Results
+After completing the previous steps, your channels will be aligned with the template channels required by our models. In case there are still some channels that haven't been mapped, we will automatically batch and optimally assign them to the template.  
+Once the channel mapping process is completed, a **JSON file** containing the mapping results will be generated. This file is necessary only if you plan to run the models using the source code; otherwise, you can ignore it.
 
 ## 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.
-1. **Initial model run**: Using the results from the previous mapping steps, the model will first run on the processed data. After this initial run, the reconstructed signals will be restored to their original positions in your data.
-2. **Subsequent processing**: If your data contains more than 30 channels, after the initial model run, we will automatically handle any remaining channels that were not initially processed. These channels will be mapped to the template channels using the Hungarian Algorithm, which optimally assigns each input channel to a template channel based on proximity. The reconstructed signals will then be placed back into their original positions in your data. This process continues iteratively until all input channels are reconstructed. If the total number of your channels is not a multiple of 30, the final iteration may result in some template channels being empty. In this case, **Mean** filling method from Step3 will be used to fill the values of these remaining empty template channels.
-
+After clicking the ``Run`` button, we will process your EEG data based on the mapping results. If necessary, your data will be devided into batches and run the models on each batch sequentially, ensuring that all channels are properly processed.
 """
 
 icunet = """
@@ -222,75 +225,75 @@ with gr.Blocks() as demo:
 	with gr.Row():
 		gr.Markdown(
 			"""
-			<p style="text-align: center;">(...)</p>
+			<p style="text-align: center;">...</p>
 			"""
 		)
 	with gr.Row():
 
 		with gr.Column():
-			gr.Markdown("# 1.Channel Mapping")
-			# ------------------------input--------------------------
-			gr.Markdown("""
-			- The data need to be a two-dimentional array (channel, timepoint).
-			- If you don't have the channel location file, we recommend you to download 
-			the standard montage <a href="">here</a>. If the channels in those files 
-			don't match yours, you can use **EEGLAB** to modify them to your needed montage.
-			""")
-			with gr.Row():
-				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", interactive=False, scale=1)
-			
-			# ------------------------mapping------------------------
-			desc_md = gr.Markdown(visible=False)
-			# 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)
-			# step2 : forward unmatched input channels to empty template channels
-			radio_group = gr.Radio(elem_id="radio-group", visible=False)
-			# step3 : fill the remaining template channels
-			with gr.Row():
-				in_fillmode = gr.Dropdown(choices=["mean", "zero"],
-									value="mean",
-									label="Filling method",
-									visible=False,
-									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)
-				step2_btn = gr.Button("Next", visible=False)
-				step3_btn = gr.Button("Next", visible=False)
-				next_btn = gr.Button("Next step", visible=False)
-			# -------------------------------------------------------
+			with gr.Row(variant='panel'):
+				with gr.Column():
+					gr.Markdown("# 1.Channel Mapping")
+					# ------------------------input--------------------------
+					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)
+					
+					# ------------------------mapping------------------------
+					desc_md = gr.Markdown(visible=False)
+					# step1 : initial matching and 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)
+					# step2 : forward unmatched input channels to empty template channels
+					radio_group = gr.Radio(elem_id="radio-group", visible=False)
+					# step3 : fill the remaining template channels
+					with gr.Row():
+						in_fillmode = gr.Dropdown(choices=["mean", "zero"],
+											value="mean",
+											label="Filling method",
+											visible=False,
+											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)
+					res_md = gr.Markdown("""
+					(Download this file if you plan to run the models using the source code.)
+					""", visible=False)
+					
+					with gr.Row():
+						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)
+					# -------------------------------------------------------
 
 		with gr.Column():
-			gr.Markdown("# 2.Decode Data")
-			# ------------------------input--------------------------
-			with gr.Row():
-				in_modelname = gr.Dropdown(choices=[
-									("ART", "EEGART"),
-									("IC-U-Net", "ICUNet"),
-									("IC-U-Net++", "UNetpp"),
-									("IC-U-Net-Attn", "AttUnet"),
-									"(mapped data)",
-									"(denoised data)"],
-									value="EEGART",
-									label="Model",
-									scale=2)
-				run_btn = gr.Button(scale=1, interactive=False)
-			
-			# ------------------------output-------------------------
-			batch_md = gr.Markdown(visible=False)
-			out_data_file = gr.File(label="Denoised data", visible=False)
-			# -------------------------------------------------------	
+			with gr.Row(variant='panel'):
+				with gr.Column():
+					gr.Markdown("# 2.Decode Data")
+					# ------------------------input--------------------------
+					in_data_file = gr.File(label="Raw data (.csv)", file_types=[".csv"])
+					with gr.Row():
+						in_modelname = gr.Dropdown(choices=[
+											("ART", "EEGART"),
+											("IC-U-Net", "ICUNet"),
+											("IC-U-Net++", "UNetpp"),
+											("IC-U-Net-Attn", "AttUnet")],
+											#"(mapped data)",
+											#"(denoised data)"],
+											value="EEGART",
+											label="Model",
+											scale=2)
+						run_btn = gr.Button(interactive=False, scale=1)
+					
+					# ------------------------output-------------------------
+					batch_md = gr.Markdown(visible=False)
+					out_data_file = gr.File(label="Denoised data", visible=False)
+					# -------------------------------------------------------
 				
 	with gr.Row():
 		with gr.Tab("ART"):
@@ -304,31 +307,20 @@ with gr.Blocks() as demo:
 		with gr.Tab("README"):
 			gr.Markdown(readme)
     
-	#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],
-	                    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)
-	
+	#demo.load(js=js)	
 	
 	# +========================================================================================+
 	# |                                Stage1: channel mapping                                 |
 	# +========================================================================================+
-	def reset_all(in_data, in_loc, samplerate):
+	def reset_all(in_loc, samplerate):
 		# establish a new folder for the current session
-		rootpath = os.path.dirname(str(in_data))
+		rootpath = os.path.dirname(str(in_loc))
 		try:
 			os.mkdir(rootpath+"/session_data/")
 		except OSError as e:
 			utils.dataDelete(rootpath+"/session_data/")
 			os.mkdir(rootpath+"/session_data/")
 			print(e)
-		# establish new folders for stage1 and stage2
 		os.mkdir(rootpath+"/session_data/stage1/")
 		os.mkdir(rootpath+"/session_data/stage2/")
 		
@@ -340,7 +332,6 @@ with gr.Blocks() as demo:
 			"stage1" : {
 				"filePath" : rootpath+"/session_data/stage1/",
 				"fileNames" : {
-					"input_data" : in_data,
 					"input_loc" : in_loc,
 					"input_montage" : "",
 		    		"mapped_montage" : ""
@@ -361,12 +352,12 @@ with gr.Blocks() as demo:
 			"stage2" : {
 				"filePath" : rootpath+"/session_data/stage2/",
 				"fileNames" : {
+					"input_data" : "",
 					"output_data" : ""
 				},
 				"totalBatchNum" : None
 			}
 		}
-		# reset layout
 		return {app_info_json : app_info,
 				channel_info_json : channel_info,
 				# --------------------Stage1-------------------------
@@ -383,7 +374,9 @@ with gr.Blocks() as demo:
 				chkbox_group : gr.CheckboxGroup(choices=[], value=[], label="", visible=False),
 				step3_btn : gr.Button(visible=False),
 				out_json_file : gr.File(value=None, visible=False),
+				res_md : gr.Markdown(visible=False),
 				# --------------------Stage2-------------------------
+				in_data_file : gr.File(value=None),
 				run_btn : gr.Button(interactive=False),
 				batch_md : gr.Markdown(visible=False),
 				out_data_file : gr.File(visible=False)}
@@ -400,17 +393,13 @@ with gr.Blocks() as demo:
 		# ========================================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)}
+		    
+			# match the names of in_channels and tpl_channels
 			stage1_info, channel_info, tpl_montage, in_montage = app_utils.match_names(stage1_info, channel_info)
-			
-			# 2. rescale coordinates
-			yield {desc_md : gr.Markdown("Rescaling...")}
+			# rescale coordinates
 			channel_info = app_utils.align_coords(channel_info, tpl_montage, in_montage)
-		    
-			# 3. generate and save figures of the montages
+			# 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)
@@ -418,24 +407,19 @@ with gr.Blocks() as demo:
 				"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.
 				"""
 			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.
 				"""
-			
 			stage1_info["state"] = "step1-finished"
 			app_info["stage1"] = stage1_info
 			yield {app_info_json : app_info,
@@ -455,13 +439,6 @@ with gr.Blocks() as demo:
 			# 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(
@@ -472,24 +449,22 @@ with gr.Blocks() as demo:
 				app_info["stage2"] = stage2_info
 				yield {app_info_json : app_info,
 						channel_info_json : channel_info,
-						desc_md : gr.Markdown(md),
+						desc_md : gr.Markdown(visible=False),
 						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),
+						res_md : gr.Markdown(visible=True),
 						run_btn : gr.Button(interactive=True)}
 			
 			# 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,
@@ -521,15 +496,12 @@ with gr.Blocks() as demo:
 			# 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,
@@ -570,13 +542,6 @@ with gr.Blocks() as demo:
 			# step2 to step4
 			# all the unmatched tpl_channels were filled by in_channels
 			if len(stage1_info["missingTemplates"]) == 0:
-				#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(
@@ -587,23 +552,21 @@ with gr.Blocks() as demo:
 				app_info["stage2"] = stage2_info
 				yield {app_info_json : app_info,
 						channel_info_json : channel_info,
-						desc_md : gr.Markdown(md),
+						desc_md : gr.Markdown(visible=False),
 						radio_group : gr.Radio(visible=False),
 						out_json_file : gr.File(filename, visible=True),
+						res_md : gr.Markdown(visible=True),
 						clear_btn : gr.Button(visible=False),
 						next_btn : gr.Button(visible=False),
 						run_btn : gr.Button(interactive=True)}
 			# step2 to step3-1
 			else:
-				#print('step2 -> 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,
@@ -620,13 +583,6 @@ with gr.Blocks() as demo:
 			
 			# 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(
@@ -637,27 +593,24 @@ with gr.Blocks() as demo:
 				app_info["stage2"] = stage2_info
 				yield {app_info_json : app_info,
 						channel_info_json : channel_info,
-						desc_md : gr.Markdown(md),
+						desc_md : gr.Markdown(visible=False),
 						in_fillmode : gr.Dropdown(visible=False),
 						fillmode_btn : gr.Button(visible=False),
 						out_json_file : gr.File(filename, visible=True),
+						res_md : gr.Markdown(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,
@@ -692,7 +645,6 @@ with gr.Blocks() as demo:
 		# =======================================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
@@ -705,12 +657,6 @@ with gr.Blocks() as demo:
 				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(
@@ -721,17 +667,18 @@ with gr.Blocks() as demo:
 			app_info["stage2"] = stage2_info
 			yield {app_info_json : app_info,
 					channel_info_json : channel_info,
-					desc_md : gr.Markdown(md),
+					desc_md : gr.Markdown(visible=False),
 					chkbox_group : gr.CheckboxGroup(visible=False),
 					next_btn : gr.Button(visible=False),
 					out_json_file : gr.File(filename, visible=True),
+					res_md : gr.Markdown(visible=True),
 					run_btn : gr.Button(interactive=True)}
 	
 	next_btn.click(
 		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, 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]
+		            in_fillmode, fillmode_btn, chkbox_group, step3_btn, out_json_file, res_md, next_btn, run_btn]
 	).success(
 	    fn = None,
 	    js = init_js,
@@ -745,10 +692,10 @@ with gr.Blocks() as demo:
 	# +========================================================================================+
 	map_btn.click(
 	    fn = reset_all,
-	    inputs = [in_data_file, in_loc_file, in_samplerate],
+	    inputs = [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]
+	    			radio_group, clear_btn, step2_btn, in_fillmode, fillmode_btn, chkbox_group, step3_btn,
+	    			out_json_file, res_md, in_data_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], 
@@ -781,7 +728,6 @@ with gr.Blocks() as demo:
 	
 	def update_radio(app_info, channel_info, selected):
 		stage1_info = app_info["stage1"]
-		
 		# ----------------------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 != []:
@@ -838,7 +784,6 @@ with gr.Blocks() as demo:
 	# +========================================================================================+	
 	def update_chkbox(app_info, channel_info, selected):
 		stage1_info = app_info["stage1"]
-		
 		# ----------------------store information before the button click-----------------------
 		# if the user didn't uncheck all in_channel checkboxes
 		if selected != []:
@@ -875,7 +820,7 @@ with gr.Blocks() as demo:
 	    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]
+					out_json_file, res_md, next_btn, run_btn]
 	).success(
 	    fn = None,
 	    js = init_js,
@@ -898,7 +843,7 @@ with gr.Blocks() as demo:
     # +========================================================================================+
 	# |                                  Stage2: decode data                                   |
 	# +========================================================================================+
-	def reset_run(app_info, modelname):
+	def reset_run(app_info, in_data, modelname):
 		stage1_info = app_info["stage1"]
 		stage2_info = app_info["stage2"]
 	    
@@ -909,13 +854,13 @@ with gr.Blocks() as demo:
 		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 = os.path.basename(str(filename))
+		filename = os.path.basename(str(in_data))
 		new_filename = os.path.splitext(filename)[0]+'_'+modelname+'.csv'
 		
 		stage2_info.update({
 			"filePath" : new_filepath,
 			"fileNames" : {
+				"input_data" : in_data,
 				"output_data" : new_filepath + new_filename
 			}
 		})
@@ -931,13 +876,10 @@ with gr.Blocks() as demo:
 		
 		filepath = stage2_info["filePath"]
 		samplerate = app_info["sampleRate"]
-		filename = stage1_info["fileNames"]["input_data"]
+		filename = stage2_info["fileNames"]["input_data"]
 		new_filename = stage2_info["fileNames"]["output_data"]
 		
-		# 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
+		break_flag = False # flag to indicate if the process has been interrupted by the user
 		for i in range(stage2_info["totalBatchNum"]):
 			# establish a temp folder
 			try:
@@ -992,7 +934,7 @@ with gr.Blocks() as demo:
 	
 	run_btn.click(
 		fn = reset_run,
-		inputs = [app_info_json, in_modelname],
+		inputs = [app_info_json, in_data_file, in_modelname],
 		outputs = [app_info_json, run_btn, batch_md, out_data_file]
 		
 	).success(