update app.py

app.py

@@ -2,60 +2,12 @@ import gradio as gr
 import numpy as np
 import os
 import utils
-import channel_mapping
+from channel_mapping import mapping, reorder_data
 
-def test(model_name, fill_mode, csv_file, chanlocs_file):
-    tmp_path = os.path.dirname(str(csv_file))
-    
-    input_name = os.path.basename(str(csv_file))
-    output_name = os.path.splitext(input_name)[0]+'_'+model_name+'.csv'
-    
-    # Channel mapping
-    #input_loc_name = os.path.basename(str(chanlocs_file))
-    fill_mode = "adjacent" if fill_mode == "adjacent channel's value" else fill_mode
-    channel_mapping.mapping(csv_file, chanlocs_file, fill_mode)
-    
-    # step1: Data preprocessing
-    total_file_num = utils.preprocessing(tmp_path, 'mapped.csv', 256)
-    
-    # step2: Signal reconstruction
-    utils.reconstruct(model_name, total_file_num, tmp_path, output_name)
-    
-    return tmp_path+'/'+output_name
+import mne
+from mne.channels import read_custom_montage
 
-
-with gr.Blocks() as app:
-    with gr.Row():
-        gr.Markdown(
-        """
-# Introduction
-(...)
-        """
-        )
-    with gr.Row():
-        with gr.Column():
-            input_model_name = gr.Dropdown(choices=["ICUNet", "UNetpp", "AttUnet", "EEGART"],
-                                value="ICUNet",
-                                label="Model")
-            input_fill_mode = gr.Dropdown(choices=["zero", "adjacent channel's value"],
-                                value="zero",
-                                label="Imputation")
-            input_data = gr.File(label="Raw data (.csv)", file_types=[".csv"])
-            input_loc = gr.File(label="Channel locations (.loc, .locs)", file_types=[".loc", "locs"])
-            btn = gr.Button()
-        with gr.Column():
-            output_data = gr.File(label="Denoised data")
-            # put template_loc file
-            gr.Markdown(
-            """
-            (Missing channels)  
-            (New channel locations)
-            """
-            )
-    with gr.Row():
-        with gr.Tab("README"):
-            gr.Markdown(
-            """
+readme = """
 # Quickstart
 
 ### Raw data
@@ -86,28 +38,292 @@ Therefore, you need to
 <b><font color=#FF0000>remove these channels</font></b>
 after you download the denoised data.
 
-            """
-            # ### Denoised data: Once the reconstructing process finished, the denoised data will be downloadable here.
-            # ### New channel locations: The template channel locations is downloadable here.
-            )
-        with gr.Tab("IC-U-Net"):
-            gr.Markdown(
-            """
+"""
+# ### Denoised data: Once the reconstructing process finished, the denoised data will be downloadable here.
+# ### New channel locations: The template channel locations is downloadable here.
+
+icunet = """
 # IC-U-Net
 ### Abstract
-Electroencephalography (EEG) signals are often contaminated with artifacts. It is imperative to develop a practical and reliable artifact removal method to prevent the misinterpretation of neural signals and the underperformance of brain–computer interfaces. Based on the U-Net architecture, we developed a new artifact removal model, IC-U-Net, for removing pervasive EEG artifacts and reconstructing brain signals. IC-U-Net was trained using mixtures of brain and non-brain components decomposed by independent component analysis. It uses an ensemble of loss functions to model complex signal fluctuations in EEG recordings. The effectiveness of the proposed method in recovering brain activities and removing various artifacts (e.g., eye blinks/movements, muscle activities, and line/channel noise) was demonstrated in a simulation study and four real-world EEG experiments. IC-U-Net can reconstruct a multi-channel EEG signal and is applicable to most artifact types, offering a promising end-to-end solution for automatically removing artifacts from EEG recordings. It also meets the increasing need to image natural brain dynamics in a mobile setting. 
+Electroencephalography (EEG) signals are often contaminated with artifacts. It is imperative to develop a practical and reliable artifact removal method to prevent the misinterpretation of neural signals and the underperformance of brain–computer interfaces. Based on the U-Net architecture, we developed a new artifact removal model, IC-U-Net, for removing pervasive EEG artifacts and reconstructing brain signals. IC-U-Net was trained using mixtures of brain and non-brain components decomposed by independent component analysis. It uses an ensemble of loss functions to model complex signal fluctuations in EEG recordings. The effectiveness of the proposed method in recovering brain activities and removing various artifacts (e.g., eye blinks/movements, muscle activities, and line/channel noise) was demonstrated in a simulation study and four real-world EEG experiments. IC-U-Net can reconstruct a multi-channel EEG signal and is applicable to most artifact types, offering a promising end-to-end solution for automatically removing artifacts from EEG recordings. It also meets the increasing need to image natural brain dynamics in a mobile setting.
+"""
 
+chk_html = """
+<form name="test" id="chs-form">
+    <input type="checkbox" />
+</form>
+"""
+
+chk_script = """
+
+let channels = document.getElementById("chs-checkbox");
+
+// init generate checkboxgroup
+let obj = document.getElementById("map-result").value; // emmm......
+let channels = obj.channels;
+let num = channels.length;
+
+for(i=0; i<num; i++){
+	document.getElementById("gen-checkbox").
+		innerHTML += '<input type="checkbox" class="channels" name="channel" value=channels[i].name />'
+}
+
+// check if mapping just finished
+const result = document.getElementById("map-result")
+result.addEventListener("change", function() {
+	const res_obj = this.value;
+	
+	if(res_obj.fill_mode=="mean" && res_obj.missing_channels.length!=0 && !res_obj.start){
+		gen_chkbox(res_obj.missing_channels);
+		res_obj.start = True;
+	}
+})
+
+function gen_chkbox(channels){
+	let num = channels.length;
+	let chs_form = document.getElementById("chs-form");
+	
+	chs_form.innerHTML = "";
+	for(i=0; i<num; i++){
+		chs_form.innerHTML += '<input type="checkbox" class="channels" name="channel" value=channels[i].name />'
+	}
+}
+
+
+"""
+
+with gr.Blocks() as demo:
+
+	state_json = gr.JSON(elem_id="state", visible=False)
+
+	with gr.Row():
+		gr.Markdown(
+		"""
+# Introduction
+(...)
+		"""
+		)
+	with gr.Row():
+		with gr.Column():
+			gr.Markdown(
+			"""
+# 1.
             """
-            )
-        with gr.Tab("IC-U-Net++"):
-            gr.Markdown(
-            """
-# Test
-            """
-            )
-        
-    inputs = [input_model_name, input_fill_mode, input_data, input_loc]
-    btn.click(fn=test, inputs=inputs, outputs=output_data)
+			)
+			with gr.Row():
+				input_raw_data = gr.File(label="Raw data (.csv)", file_types=[".csv"])
+				input_raw_loc = gr.File(label="Channel locations (.loc, .locs)", file_types=[".loc", "locs"])
+			with gr.Row():
+				input_fill_mode = gr.Dropdown(choices=["zero",
+									("adjacent channel's value", "adjacent"),
+									("mean(manual)", "mean")],
+									value="zero",
+									label="Imputation",
+									scale=2)
+				map_btn = gr.Button("Mapping", scale=1)
+				channels_json = gr.JSON(visible=False)
+			with gr.Row():
+				tpl_montage = gr.Image("./template_montage.png", label="Template montage", visible=False)
+				map_montage = gr.Image(label="Mapping result", visible=False)
+				input_montage = gr.Image(label="Input montage", visible=False)
+			with gr.Accordion(visible=False) as accordion: #???
+				with gr.Column():
+					#chk_block = gr.HTML(chk_html)
+					#input_montage = gr.Image(label="Input montage")
+					chs_chkbox = gr.CheckboxGroup(elem_classes="chs-chkbox", label="")
+					next_btn = gr.Button("Submmit/Next", interactive=False)
+			tpl_loc_file = gr.File("./template_chanlocs.loc", label="Template channel locations file")
+		with gr.Column():
+			gr.Markdown(
+			"""
+# 2.
+			"""
+			)
+			with gr.Row():
+				input_model_name = gr.Dropdown(choices=["ICUNet", "UNetpp", "AttUnet", "EEGART"],
+									value="ICUNet",
+									label="Model",
+									scale=2)
+				run_btn = gr.Button(scale=1, interactive=False)
+			ouput_denoised_data = gr.File(label="Denoised data")
+			
+			#with gr.Row():
+			#	mapResult_montage_img = gr.HTML()
+			#	tpl_montage_img = gr.HTML()
+			gr.Markdown(
+			"""
+			(Missing channels)
+			"""
+			)
+				
+	with gr.Row():
+		with gr.Tab("EEGART"):
+			gr.Markdown()
+		with gr.Tab("IC-U-Net"):
+			gr.Markdown(icunet)
+		with gr.Tab("IC-U-Net++"):
+			gr.Markdown()
+		with gr.Tab("IC-U-Net-Att"):
+			gr.Markdown()
+		with gr.Tab("README"):
+			gr.Markdown(readme)
+    
+	#demo.load(js=chk_script)
+	
+	def reset_layout(raw_data):
+		# establish temp folders
+		filepath = os.path.dirname(str(raw_data))
+		
+		#filename = filepath+"/temp_data/img/input_montage.png"
+		#if os.path.exists(filename):
+		#	os.remove(filename)
+		#	print("Previous input image has been deleted!")
+		
+		try:
+			os.mkdir(filepath+"/temp_data/")
+		except OSError as e:
+			utils.dataDelete(filepath+"/temp_data/")
+			os.mkdir(filepath+"/temp_data/")
+			print(e)
+		try:
+			os.mkdir(filepath+"/temp_data/img/")
+		except OSError as e:
+			utils.dataDelete(filepath+"/temp_data/img/")
+			os.mkdir(filepath+"/temp_data/img/")
+			print(e)
+		
+		return {state_json : {},
+				accordion : gr.Accordion(visible=False),
+				chs_chkbox : gr.CheckboxGroup(choices=[], value=[], label=""), # choices, value ???
+				next_btn : gr.Button(interactive=False),
+				run_btn : gr.Button(interactive=False),
+				tpl_montage : gr.Image(visible=False),
+				map_montage : gr.Image(value=None, visible=False),
+				input_montage : gr.Image(value=None, visible=False),}
+    	
+	def mapping_result(channels_obj, raw_data, fill_mode):
+		state_obj = channels_obj.copy()
+		filepath = os.path.dirname(str(raw_data))
+		if fill_mode=="mean" and channels_obj["missingIndex"]!=[]:
+			state_obj.update({
+				"state" : "initializing",
+				"currentIndex" : 0,
+				"maxIndex" : len(channels_obj["missingIndex"])-1
+			})
+			#print("Missing channels:", state_obj["missingIndex"])
+			return {state_json : state_obj,
+					accordion : gr.Accordion(visible=True)}
+		else:
+			state_obj.update({
+				"state" : "finished",
+				"currentIndex" : -1,
+				"maxIndex" : -1
+			})
+			reorder_data(raw_data, channels_obj["newOrder"], fill_mode)
+			return {state_json : state_obj,
+					run_btn : gr.Button(interactive=True)}
+
+
+	def generate_chkbox(state_obj):
+		if state_obj["state"] == "finished":
+			return {state_json : {}}
+		if state_obj["state"] == "initializing":
+			in_channels = [channel for channel in state_obj["inputByName"]]
+			state_obj["state"] = "selecting"
+			# and img....
+			
+			first_idx = state_obj["missingIndex"][0]
+			first_name = state_obj["templateByIndex"][first_idx]
+			chkbox_label = first_name+' (1/'+str(state_obj["maxIndex"]+1)+')'
+			return {state_json : state_obj,
+					chs_chkbox : gr.CheckboxGroup(choices=in_channels, label=chkbox_label),
+					next_btn : gr.Button(interactive=True)}
+	
+	def show_montage(state_obj, raw_data, raw_loc):
+		filepath = os.path.dirname(str(raw_data))
+		
+		raw_montage = read_custom_montage(raw_loc)
+		if state_obj["state"] == "initializing":
+			filename = filepath+"/temp_data/img/input_montage.png"
+			input_fig = raw_montage.plot(show=False)
+			input_fig.savefig(filename)
+			return {tpl_montage : gr.Image(visible=True),
+					input_montage : gr.Image(value=filename, visible=True)}
+		
+		elif state_obj["state"] == "finished":
+			# didn't find any way to remove channels from DigMontage...
+			# tmp
+			filename = filepath+"/temp_data/img/input_montage.png"
+			input_fig = raw_montage.plot(show=False)
+			input_fig.savefig(filename)
+			return {input_montage : gr.Image(visible=False),
+					tpl_montage : gr.Image(visible=True),
+					map_montage : gr.Image(value=filename, visible=True)} # value=, 
+					
+		elif state_obj["state"] == "selecting":
+			# need to update input_montage here
+			return {input_montage : gr.Image()}
+
+	
+	map_btn.click(reset_layout, input_raw_data,
+					[state_json, accordion, chs_chkbox, next_btn, run_btn, tpl_montage, map_montage, input_montage]).success(
+				  mapping, [input_raw_data, input_raw_loc, input_fill_mode], channels_json).success(
+				  mapping_result, [channels_json, input_raw_data, input_fill_mode], [state_json, accordion, run_btn]).success(
+				  show_montage, [state_json, input_raw_data, input_raw_loc], [input_montage, tpl_montage, map_montage]).success(
+				  generate_chkbox, state_json, [state_json, chs_chkbox, next_btn])
+    
+
+	def check_next(state_obj, selected, raw_data, fill_mode):
+		if state_obj["state"] == "selecting":
+			if state_obj["currentIndex"] <= state_obj["maxIndex"]:
+				
+				# info before clicking on next_btn
+				prev_target_idx = state_obj["missingIndex"][state_obj["currentIndex"]]
+				prev_target_name = state_obj["templateByIndex"][prev_target_idx]
+				
+				selected_idx = [state_obj["inputByName"][channel]["index"] for channel in selected]
+				state_obj["newOrder"][prev_target_idx] = selected_idx
+				
+				print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
+				
+				# current info
+				state_obj["currentIndex"] += 1
+				if state_obj["currentIndex"] <= state_obj["maxIndex"]:
+					target_idx = state_obj["missingIndex"][state_obj["currentIndex"]]
+					target_name = state_obj["templateByIndex"][target_idx]
+					chkbox_label = target_name+' ('+str(state_obj["currentIndex"]+1)+'/'+str(state_obj["maxIndex"]+1)+')'
+					
+					#return {chk_block : chk_html, state_json : state_obj}
+					return {state_json : state_obj,
+							chs_chkbox : gr.CheckboxGroup(value=[], label=chkbox_label)}
+				else:
+					state_obj["state"] = "finished"
+					reorder_data(raw_data, state_obj["newOrder"], fill_mode)
+					
+					return {state_json : state_obj,
+							accordion : gr.Accordion(visible=False),
+							run_btn : gr.Button(interactive=True)}
+	
+	next_btn.click(check_next, [state_json, chs_chkbox, input_raw_data, input_fill_mode],
+	                            [state_json, accordion, chs_chkbox, run_btn]).success(
+	               show_montage, [state_json, input_raw_data, input_raw_loc], [input_montage, tpl_montage, map_montage])
+    
+    
+	@run_btn.click(inputs=[input_model_name, input_raw_data], outputs=ouput_denoised_data)
+	def run_model(model_name, raw_file):    	
+		file_path = os.path.dirname(str(raw_file))
+    	
+		input_name = os.path.basename(str(raw_file))
+		output_name = os.path.splitext(input_name)[0]+'_'+model_name+'.csv'
+    	
+    	# step1: Data preprocessing
+		total_file_num = utils.preprocessing(file_path+'/temp_data', 'mapped.csv', 256)
+		
+		# step2: Signal reconstruction
+		utils.reconstruct(model_name, total_file_num, file_path+'/temp_data/', output_name)
+		
+		return file_path+'/'+output_name
+    
 
 if __name__ == "__main__":
-    app.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)