update

app.py

@@ -5,30 +5,29 @@ import os
 import random
 
 readme = """
-
-This tool serves two main purposes:
+## Introduction
+This tool is designed to assist you in two key tasks:
 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.
+- **Channel requirements**: Your data must include some channels that correspond to our template channels, which include: ``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. Additionally, please remove any reference, ECG, EOG, EMG... channels before uploading your files.
 
 ## 1. Channel Mapping
 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 will be directed to **Mapping Results**.
+- If your data includes all the 30 template channels, you will see the **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 will be directed to **Mapping Results**.
+- If all empty template channels are filled by your selections, you will see the **Mapping Results**.
 - If there are still empty template channels remaining, you will be directed to **Step3**.
 
 ### Step3: Filling Remaining Template Channels
@@ -36,7 +35,7 @@ To run the models successfully, we need to ensure that all 30 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.  
-Once all template channels are filled, you will be directed to **Mapping Results**.
+Once all template channels are filled, you will see the **Mapping Results**.
 
 ### 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.  
@@ -58,7 +57,7 @@ init_js = """
 	channel_info = JSON.parse(JSON.stringify(channel_info));
 	stage1_info = app_info.stage1
 	
-	let selector, classname, attribute;
+	let selector, attribute; //, classname;
 	let channel, left, bottom;
 	
 	if(stage1_info.state == "step2-selecting"){
@@ -81,7 +80,6 @@ init_js = """
 		//height: 560px;
 		background: url("file=${stage1_info.fileNames.input_montage}");
 		background-size: contain;
-
 	`;
 	
 	// move the radios/checkboxes
@@ -96,7 +94,6 @@ init_js = """
 		item.querySelector(":scope > span").innerText = "";
 	});
 	
-	
 	// add indication for the missing channels
 	channel = stage1_info.missingTemplates[0];
 	left = channel_info.templateDict[channel].css_position[0];
@@ -170,7 +167,7 @@ update_js = """
 	    selector = "#chkbox-group > div:nth-of-type(2)";
 	}else return;
 	
-	// update indication
+	// update the indication
 	channel = stage1_info.missingTemplates[stage1_info["fillingCount"]-1];
 	left = channel_info.templateDict[channel].css_position[0];
 	bottom = channel_info.templateDict[channel].css_position[1];
@@ -216,7 +213,6 @@ update_js = """
 }
 """
 
-
 with gr.Blocks() as demo:
 
 	app_info_json = gr.JSON(visible=False)
@@ -231,51 +227,50 @@ with gr.Blocks() as demo:
 	with gr.Row():
 
 		with gr.Column():
+			# ------------------------input--------------------------
 			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"],
+			# ------------------------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)
-					# -------------------------------------------------------
+				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():
+			# ------------------------input--------------------------
 			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=[
@@ -283,17 +278,15 @@ with gr.Blocks() as demo:
 											("IC-U-Net", "ICUNet"),
 											("IC-U-Net++", "UNetpp"),
 											("IC-U-Net-Attn", "AttUnet")],
-											#"(mapped data)",
-											#"(denoised data)"],
+											#"(mapped 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)
-					# -------------------------------------------------------
+			# ------------------------output-------------------------
+			batch_md = gr.Markdown(visible=False)
+			out_data_file = gr.File(label="Denoised data", visible=False)
+			# -------------------------------------------------------
 				
 	with gr.Row():
 		with gr.Tab("ART"):
@@ -306,8 +299,6 @@ with gr.Blocks() as demo:
 			gr.Markdown()
 		with gr.Tab("README"):
 			gr.Markdown(readme)
-    
-	#demo.load(js=js)	
 	
 	# +========================================================================================+
 	# |                                Stage1: channel mapping                                 |
@@ -320,7 +311,7 @@ with gr.Blocks() as demo:
 		except OSError as e:
 			utils.dataDelete(rootpath+"/session_data/")
 			os.mkdir(rootpath+"/session_data/")
-			print(e)
+			#print(e)
 		os.mkdir(rootpath+"/session_data/stage1/")
 		os.mkdir(rootpath+"/session_data/stage2/")
 		
@@ -429,7 +420,6 @@ with gr.Blocks() as demo:
 					tpl_img : gr.Image(visible=True),
 					mapped_img : gr.Image(value=filename2, visible=True),
 					next_btn : gr.Button(visible=True)}
-			
 		
 		# ========================================step1=========================================
 		elif stage1_info["state"] == "step1-finished":
@@ -706,16 +696,16 @@ with gr.Blocks() as demo:
 	# +========================================================================================+
 	# |                                      Stage1-step2                                      |
 	# +========================================================================================+
-	# ...
+	# determine which button to display based on the current state
 	@radio_group.select(inputs = app_info_json, outputs = [step2_btn, next_btn])
 	def determine_button(app_info):
 		stage1_info = app_info["stage1"]
-		if len(stage1_info["unassignedInputs"])==1 and stage1_info["fillingCount"]<stage1_info["totalFillingNum"]:
+		if len(stage1_info["unassignedInputs"]) == 1:
 			return {step2_btn : gr.Button(visible=False),
 					next_btn : gr.Button(visible=True)}
 		else:
 			return {step2_btn : gr.Button()} # change nothing
-			
+	# clear the selected value and reset the buttons
 	@clear_btn.click(inputs = app_info_json, outputs = [radio_group, step2_btn, next_btn])
 	def clear_value(app_info):
 		stage1_info = app_info["stage1"]
@@ -905,19 +895,12 @@ with gr.Blocks() as demo:
 			try:
 				# step1: Reorder input data
 				data_shape = app_utils.reorder_data(new_idx, fill_flags, filename, filepath+"temp_data/mapped.csv")
-				if modelname == "(mapped data)":
-					new_filename = filepath+"temp_data/mapped.csv"
-					break
 				# step2: Data preprocessing
 				total_file_num = utils.preprocessing(filepath+"temp_data/", "mapped.csv", samplerate)
 				# step3: Signal reconstruction
 				utils.reconstruct(modelname, total_file_num, filepath+"temp_data/", "denoised.csv", samplerate)
-				if modelname == "(denoised data)":
-					new_filename = filepath+"temp_data/denoised.csv"
-					break
 				# step4: Restore original order
 				app_utils.restore_order(i, data_shape, new_idx, fill_flags, filepath+"temp_data/denoised.csv", new_filename)
-				break
 			except FileNotFoundError:
 				print('break2!!')
 				break_flag = True

app_utils.py

@@ -14,6 +14,7 @@ from sklearn.neighbors import NearestNeighbors
 def reorder_data(idx_order, fill_flags, filename, new_filename):
 	# read the input data
 	raw_data = utils.read_train_data(filename)
+	#print(raw_data.shape)
 	new_data = np.zeros((30, raw_data.shape[1]))
 	
 	zero_arr = np.zeros((1, raw_data.shape[1]))
@@ -26,7 +27,6 @@ def reorder_data(idx_order, fill_flags, filename, new_filename):
 			tmp_data = [raw_data[j, :] for j in idx_set]
 			new_data[i, :] = np.mean(tmp_data, axis=0)
 	
-	#print(raw_data.shape, new_data.shape)
 	utils.save_data(new_data, new_filename)
 	return raw_data.shape
 
@@ -35,6 +35,7 @@ def restore_order(batch_cnt, raw_data_shape, idx_order, fill_flags, filename, ne
 	d_data = utils.read_train_data(filename)
 	if batch_cnt == 0:
 		new_data = np.zeros((raw_data_shape[0], d_data.shape[1]))
+		#print(new_data.shape)
 	else:
 		new_data = utils.read_train_data(new_filename)
 	
@@ -43,7 +44,6 @@ def restore_order(batch_cnt, raw_data_shape, idx_order, fill_flags, filename, ne
 		if flag == False:
 			new_data[idx_set[0], :] = d_data[i, :]
 	
-	#print(d_data.shape, new_data.shape)
 	utils.save_data(new_data, new_filename)
 	return
 
@@ -238,10 +238,8 @@ def match_names(stage1_info, channel_info):
 	tpl_montage, in_montage, tpl_dict, in_dict = read_montage_data(loc_file)
 	tpl_order = tpl_montage.ch_names
 	in_order = in_montage.ch_names
-	# list to store the indices of the in_channels in the order of tpl_channls
-	new_idx = [[]]*30
-	# flags to record if each tpl_channel's data is filled by "fillmode"
-	fill_flags = [True]*30
+	new_idx = [[]]*30 # store the indices of the in_channels in the order of tpl_channls
+	fill_flags = [True]*30 # record if each tpl_channel's data is filled by "fillmode"
 	
 	alias_dict = {
 		'T3': 'T7',
@@ -340,17 +338,17 @@ def optimal_mapping(channel_info):
 	return mapping_data, channel_info
 
 def mapping_result(stage1_info, stage2_info, channel_info, filename):
-	# 1. calculate how many times the model needs to be run
+	# calculate how many times the model needs to be run
 	unassigned_num = len(stage1_info["unassignedInputs"])
 	batch_num = math.ceil(unassigned_num/30) + 1
 	
-	# 2. map the remaining in_channels
+	# map the remaining in_channels
 	for i in range(1, batch_num):
 		# optimally select 30 in_channels to map to the tpl_channels based on proximity
 		new_mapping_data, channel_info = optimal_mapping(channel_info)
 		stage1_info["mappingData"] += [new_mapping_data]
 	
-	# 3. save the mapping result
+	# save the mapping result
 	new_dict = {
 		#"templateOrder" : channel_info["templateOrder"],
 		#"inputOrder" : channel_info["inputOrder"],