update v2

app.py

@@ -2,8 +2,10 @@ import gradio as gr
 import numpy as np
 import os
 import random
+import math
+import json
 import utils
-from channel_mapping import mapping, reorder_data
+from channel_mapping import mapping_stage1, mapping_stage2, reorder_to_template, reorder_to_origin
 
 import mne
 from mne.channels import read_custom_montage
@@ -11,12 +13,9 @@ from mne.channels import read_custom_montage
 quickstart = """
 # Quickstart
 
-## 1. Channel mapping
-
 ### Raw data
 1. The data need to be a two-dimensional array (channel, timepoint).
-2. Make sure you have **resampled** your data to **256 Hz**.
-3. Upload your EEG data in `.csv` format.
+2. Upload your EEG data in `.csv` format.
 
 ### Channel locations
 Upload your data's channel locations in `.loc` format, which can be obtained using **EEGLAB**.  
@@ -27,29 +26,12 @@ The models was trained using the EEG signals of 30 channels, including: `Fp1, Fp
 We expect your input data to include these channels as well.  
 If your data doesn't contain all of the mentioned channels, there are 3 imputation ways you can choose from:
 
-<u>Manually</u>:  
-- **mean**: select the channels you wish to use for imputing the required one, and we will average their values. If you select nothing, zeros will be imputed. For example, you didn't have **FCZ** and you choose **FC1, FC2, FZ, CZ** to impute it(depending on the channels you have), we will compute the mean of these 4 channels and assign this new value to **FCZ**.
-
-<u>Automatically</u>:  
-Firstly, we will attempt to find neighboring channel to use as alternative. For instance, if the required channel is **FC3** but you only have **FC1**, we will use it as a replacement for **FC3**.  
-Then, depending on the **Imputation** way you chose, we will:
 - **zero**: fill the missing channels with zeros.
-- **adjacent**: fill the missing channels using neighboring channels which are located closer to the center. For example, if the required channel is **FC3** but you only have **F3, C3**, then we will choose **C3** as the imputing value for **FC3**.
->Note: The imputed channels **need to be removed** after the data being reconstructed.
+- **mean(auto)**: select 4 neareat channels for each missing channels, and we will average their values.
+- **mean(manual)**: select the channels you wish to use for imputing the required one, and we will average their values. If you select nothing, zeros will be imputed. For example, you didn't have **FCZ** and you choose **FC1, FC2, FZ, CZ** to impute it(depending on the channels you have), we will compute the mean of these 4 channels and assign this new value to **FCZ**.
 
 ### Mapping result
-Once the mapping process is finished, the **template montage** and the **input montage**(with the channels choosen by the mapping function displaying their names) will be shown.
-
-### Missing channels
-The channels displayed here are those for which the template didn't find suitable channels to use, and utilized **Imputation** to fill the missing values.  
-Therefore, you need to
-<span style="color:red">**remove these channels**</span>
-after you download the denoised data.
-
-### Template location file
-You need to use this as the **new location file** for the denoised data.
-
-## 2. Decode data
+Once the mapping process is finished, the **template montage** and the **input montage**(with the matched channels displaying their names) will be shown.
 
 ### Model
 Select the model you want to use.  
@@ -68,20 +50,43 @@ chkbox_js = """
 	state_json = JSON.parse(JSON.stringify(state_json));
 	if(state_json.state == "finished") return;
 	
-	document.querySelector("#chs-chkbox>div:nth-of-type(2)").style.cssText = `
+	// add figure of in_montage
+	document.querySelector("#chkbox-group> div:nth-of-type(2)").style.cssText = `
 		position: relative;
 		width: 560px;
 		height: 560px;
 		background: url("file=${state_json.files.raw_montage}");
 	`;
+	
+	// add indication for the missing channel
+	/*
+	let indicator = document.getElementById("indicator")
+	if(!indicator) document.querySelector("#chkbox-group> div:nth-of-type(2)").innerHTML += '<div id="indicator"></div>'
+	
+	let channel = state_json.missingChannelsIndex[0]
+	channel = state_json.templateByIndex[channel]
+	let left = state_json.templateByName[channel].css_position[0];
+	let bottom = state_json.templateByName[channel].css_position[1];
+	
+	document.getElementById("red-dot").style.cssText = `
+		position: absolute;
+		background-color: red;
+		width: 10px;
+		height: 10px;
+		border-radius: 50%;
+		left: ${left};
+		bottom: ${bottom};
+	`;
+	*/
 
-	let all_chkbox = document.querySelectorAll("#chs-chkbox> div:nth-of-type(2)> label");
+	// move the checkboxes
+	let all_chkbox = document.querySelectorAll("#chkbox-group> div:nth-of-type(2)> label");
 	all_chkbox = Array.apply(null, all_chkbox);
 	
 	all_chkbox.forEach((item, index) => {
-		let channel = state_json.inputByIndex[index];
-		let left = state_json.inputByName[channel].css_position[0];
-		let bottom = state_json.inputByName[channel].css_position[1];
+		channel = state_json.inputByIndex[index];
+		left = state_json.inputByName[channel].css_position[0];
+		bottom = state_json.inputByName[channel].css_position[1];
 		//console.log(`left: ${left}, bottom: ${bottom}`);
 		
 		item.style.cssText = `
@@ -90,42 +95,73 @@ chkbox_js = """
 			bottom: ${bottom};
 		`;
 		item.className = "";
-		item.querySelector("span").innerText = "";
-	});
+		item.querySelector(":scope> span").innerText = "";
+	});	
+}
+"""
+
+dot_js = """
+(state_json) => {
+	state_json = JSON.parse(JSON.stringify(state_json));
+	if(state_json.state == "finished") return;
+	
+	let channel = state_json.missingChannelsIndex[state_json["fillingCount"]]
+	channel = state_json.templateByIndex[channel]
+	let left = state_json.templateByName[channel].css_position[0];
+	let bottom = state_json.templateByName[channel].css_position[1];
 	
+	document.getElementById("indicator").style.cssText = `
+		position: absolute;
+		background-color: red;
+		width: 10px;
+		height: 10px;
+		border-radius: 50%;
+		left: ${left};
+		bottom: ${bottom};
+	`;
 }
 """
 
 
 with gr.Blocks() as demo:
 
-	state_json = gr.JSON(elem_id="state", visible=False)
+	state_json = gr.JSON(visible=False)
 
 	with gr.Row():
 		gr.Markdown(
 			"""
-			
+			<p style="text-align: center;">(...)</p>
 			"""
 		)
 	with gr.Row():
+
 		with gr.Column():
 			gr.Markdown(
 				"""
 				# 1.Channel Mapping
 				"""
 			)
+
+			# upload files, chose imputation way (???
 			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"])
-			with gr.Row():
-				in_fill_mode = gr.Dropdown(choices=["zero",
-									("adjacent channel", "adjacent"),
-									("mean (manually select channels)", "mean")],
-									value="zero",
-									label="Imputation",
-									scale=2)
-				map_btn = gr.Button("Mapping", scale=1)
-				channels_json = gr.JSON(visible=False)
+				with gr.Column(min_width=100):
+					in_sample_rate = gr.Textbox(label="Sampling rate (Hz)")
+					in_fill_mode = gr.Dropdown(choices=[
+										#("adjacent channel", "adjacent"),
+										("mean (auto)", "mean_auto"),
+										("mean (manual)", "mean_manual"),
+										("",""),
+										"zero"],
+										value="mean_auto",
+										label="Imputation")
+					map_btn = gr.Button("Mapping")
+			
+			chkbox_group = gr.CheckboxGroup(elem_id="chkbox-group", label="", visible=False)
+			next_btn = gr.Button("Next", interactive=False, visible=False)
+			
+			# mapping result
 			res_md = gr.Markdown(
 						"""
 						### Mapping result:
@@ -134,11 +170,8 @@ with gr.Blocks() as demo:
 					)
 			with gr.Row():
 				tpl_montage = gr.Image("./template_montage.png", label="Template montage", visible=False)
-				map_montage = gr.Image(label="Choosen channels", visible=False)
-			chs_chkbox = gr.CheckboxGroup(elem_id="chs-chkbox", label="", visible=False)
-			next_btn = gr.Button("Next", interactive=False, visible=False)
-			miss_txtbox = gr.Textbox(label="Missing channels", visible=False)
-			tpl_loc_file = gr.File("./template_chanlocs.loc", show_label=False, visible=False)
+				map_montage = gr.Image(label="Matched channels", visible=False)
+
 		with gr.Column():
 			gr.Markdown(
 				"""
@@ -146,29 +179,36 @@ with gr.Blocks() as demo:
 				"""
 			)
 			with gr.Row():
-				in_model_name = gr.Dropdown(choices=["ICUNet", "UNetpp", "AttUnet", "EEGART", "(mapped data)"],
-									value="ICUNet",
+				in_model_name = 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)
+			batch_md = gr.Markdown(visible=False)
 			out_denoised_data = gr.File(label="Denoised data")
 	
 				
 	with gr.Row():
-		with gr.Tab("EEGART"):
+		with gr.Tab("ART"):
 			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"):
+		with gr.Tab("IC-U-Net-Attn"):
 			gr.Markdown()
 		with gr.Tab("QuickStart"):
 			gr.Markdown(quickstart)
     
 	#demo.load(js=js)
 	
-	def reset_layout(raw_data):
+	def reset_layout(raw_data, samplerate):
 		# establish temp folder
 		filepath = os.path.dirname(str(raw_data))
 		try:
@@ -177,51 +217,53 @@ with gr.Blocks() as demo:
 			utils.dataDelete(filepath+"/temp_data/")
 			os.mkdir(filepath+"/temp_data/")
 			#print(e)
-		state_obj = {
+		
+		data = utils.read_train_data(raw_data)
+		state = {
 			"filepath": filepath+"/temp_data/",
-			"files": {}
+			"files": {},
+			"sampleRate": int(samplerate),
+			"dataShape" : data.shape
 		}
-		return {state_json : state_obj,
-				chs_chkbox : gr.CheckboxGroup(choices=[], value=[], label="", visible=False), # choices, value ???
+		return {state_json : state,
+				chkbox_group : gr.CheckboxGroup(choices=[], value=[], label="", visible=False),
 				next_btn : gr.Button("Next", interactive=False, visible=False),
 				run_btn : gr.Button(interactive=False),
 				tpl_montage : gr.Image(visible=False),
 				map_montage : gr.Image(value=None, visible=False),
-				miss_txtbox : gr.Textbox(visible=False),
 				res_md : gr.Markdown(visible=False),
-				tpl_loc_file : gr.File(visible=False)}
+				batch_md : gr.Markdown(visible=False)}
     	
-	def mapping_result(state_obj, channels_obj, raw_data, fill_mode):
-		state_obj.update(channels_obj)
+	def mapping_result(state, fill_mode):
+	
+		in_num = len(state["inputByName"])
+		matched_num = 30 - len(state["missingChannelsIndex"])
+		batch_num = math.ceil((in_num-matched_num)/30) + 1
+		state.update({
+			"runnigState" : "stage1",
+			"batchCount" : 1,
+			"totalBatchNum" : batch_num
+		})
 		
-		if fill_mode=="mean" and channels_obj["missingChannelsIndex"]!=[]:
-			state_obj.update({
+		if fill_mode=="mean_manual" and state["missingChannelsIndex"]!=[]:
+			state.update({
 				"state" : "initializing",
 				"fillingCount" : 0,
-				"totalFillingNum" : len(channels_obj["missingChannelsIndex"])-1
+				"totalFillingNum" : len(state["missingChannelsIndex"])-1
 			})
-			#print("Missing channels:", state_obj["missingChannelsIndex"])
-			return {state_json : state_obj,
+			#print("Missing channels:", state["missingChannelsIndex"])
+			return {state_json : state,
+					#chkbox_group : gr.CheckboxGroup(visible=True),
 					next_btn : gr.Button(visible=True)}
 		else:
-			reorder_data(raw_data, channels_obj["newOrder"], fill_mode, state_obj)
-		    
-			missing_channels = [state_obj["templateByIndex"][idx] for idx in state_obj["missingChannelsIndex"]]
-			missing_channels = ', '.join(missing_channels)
-		    
-			state_obj.update({
-				"state" : "finished",
-				#"fillingCount" : -1,
-				#"totalFillingNum" : -1
-			})
-			return {state_json : state_obj,
+			state["state"] = "finished"
+			
+			return {state_json : state,
 					res_md : gr.Markdown(visible=True),
-			        miss_txtbox : gr.Textbox(value=missing_channels, visible=True),
-			        tpl_loc_file : gr.File(visible=True),
 					run_btn : gr.Button(interactive=True)}
 	
-	def show_montage(state_obj, raw_loc):
-		filepath = state_obj["filepath"]
+	def show_montage(state, raw_loc):
+		filepath = state["filepath"]
 		raw_montage = read_custom_montage(raw_loc)
 		
 		# convert all channel names to uppercase
@@ -229,73 +271,63 @@ with gr.Blocks() as demo:
 			channel = raw_montage.ch_names[i]
 			raw_montage.rename_channels({channel: str.upper(channel)})
 		
-		if state_obj["state"] == "initializing":
+		if state["state"] == "initializing":
 			filename = filepath+"raw_montage_"+str(random.randint(1,10000))+".png"
-			state_obj["files"]["raw_montage"] = filename
+			state["files"]["raw_montage"] = filename
 			raw_fig = raw_montage.plot()
 			raw_fig.set_size_inches(5.6, 5.6)
 			raw_fig.savefig(filename, pad_inches=0)
 			
-			return {state_json : state_obj}#,
-					#tpl_montage : gr.Image(visible=True),
-					#in_montage : gr.Image(value=filename, visible=True),
-					#map_montage : gr.Image(visible=False)}
+			return {state_json : state}
 		
-		elif state_obj["state"] == "finished":
-			# didn't find any way to hide the dark points...
-			# tmp
+		elif state["state"] == "finished":
 			filename = filepath+"mapped_montage_"+str(random.randint(1,10000))+".png"
-			state_obj["files"]["map_montage"] = filename
+			state["files"]["map_montage"] = filename
 			
 			show_names= []
-			for channel in state_obj["inputByName"]:
-			    if state_obj["inputByName"][channel]["used"]:
-			        if channel=='CZ' and state_obj["CZImputed"]:
-			            continue
+			for channel in state["inputByName"]:
+			    if state["inputByName"][channel]["matched"]:
 			        show_names.append(channel)
 			mapped_fig = raw_montage.plot(show_names=show_names)
 			mapped_fig.set_size_inches(5.6, 5.6)
 			mapped_fig.savefig(filename, pad_inches=0)
 			
-			return {state_json : state_obj,
+			return {state_json : state,
 					tpl_montage : gr.Image(visible=True),
 					map_montage : gr.Image(value=filename, visible=True)}
 					
-		elif state_obj["state"] == "selecting":
-			# update in_montage here ?
-			#return {in_montage : gr.Image()}
-			return {state_json : state_obj}
+		else:
+			return {state_json : state} # change nothing
 	
-	def generate_chkbox(state_obj):
-		if state_obj["state"] == "initializing":
-			in_channels = [channel for channel in state_obj["inputByName"]]
-			state_obj["state"] = "selecting"
+	def generate_chkbox(state):
+		if state["state"] == "initializing":
+			in_channels = [channel for channel in state["inputByName"]]
+			state["state"] = "selecting"
 			
-			first_idx = state_obj["missingChannelsIndex"][0]
-			first_name = state_obj["templateByIndex"][first_idx]
-			chkbox_label = first_name+' (1/'+str(state_obj["totalFillingNum"]+1)+')'
-			return {state_json : state_obj,
-					chs_chkbox : gr.CheckboxGroup(choices=in_channels, label=chkbox_label, visible=True),
+			first_idx = state["missingChannelsIndex"][0]
+			first_name = state["templateByIndex"][first_idx]
+			chkbox_label = first_name+' (1/'+str(state["totalFillingNum"]+1)+')'
+			return {state_json : state,
+					chkbox_group : gr.CheckboxGroup(choices=in_channels, label=chkbox_label, visible=True),
 					next_btn : gr.Button(interactive=True)}
 		else:
-			return {state_json : state_obj}
+			return {state_json : state} # change nothing
 
 	
 	map_btn.click(
 	    fn = reset_layout,
-	    inputs = in_raw_data,
-		outputs = [state_json, chs_chkbox, next_btn, run_btn, tpl_montage, map_montage, miss_txtbox,
-					res_md, tpl_loc_file]
-		
+	    inputs = [in_raw_data, in_sample_rate],
+		outputs = [state_json, chkbox_group, next_btn, run_btn, tpl_montage, map_montage, res_md, batch_md]
+					
 	).success(
-		fn = mapping, 
-		inputs = [in_raw_data, in_raw_loc, in_fill_mode], 
-		outputs = channels_json
+		fn = mapping_stage1, 
+		inputs = [state_json, in_raw_data, in_raw_loc, in_fill_mode], 
+		outputs = state_json
 		
 	).success(
 		fn = mapping_result, 
-		inputs = [state_json, channels_json, in_raw_data, in_fill_mode], 
-		outputs = [state_json, chs_chkbox, next_btn, miss_txtbox, res_md, tpl_loc_file, run_btn]
+		inputs = [state_json, in_fill_mode], 
+		outputs = [state_json, next_btn, res_md, run_btn]
 		
 	).success(
 		fn = show_montage, 
@@ -305,7 +337,8 @@ with gr.Blocks() as demo:
 	).success(
 		fn = generate_chkbox, 
 		inputs = state_json, 
-		outputs = [state_json, chs_chkbox, next_btn]
+		outputs = [state_json, chkbox_group, next_btn]
+		
 	).success(
 		fn = None,
 		js = chkbox_js,
@@ -314,81 +347,85 @@ with gr.Blocks() as demo:
 	)
     
 
-	def check_next(state_obj, selected, raw_data, fill_mode):
-		if state_obj["state"] == "selecting":
+	def check_next(state, selected, raw_data, fill_mode):
+		#if state["state"] == "selecting":
 			
-			# save info before clicking on next_btn
-			prev_target_idx = state_obj["missingChannelsIndex"][state_obj["fillingCount"]]
-			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
-			
-			if len(selected)==1 and state_obj["inputByName"][selected[0]]["used"]==False:
-			    state_obj["inputByName"][selected[0]]["used"] = True
-			    state_obj["missingChannelsIndex"][state_obj["fillingCount"]] = -1
-			
-			print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
+		# save info before clicking on next_btn
+		prev_target_idx = state["missingChannelsIndex"][state["fillingCount"]]
+		prev_target_name = state["templateByIndex"][prev_target_idx]
+		
+		selected_idx = [state["inputByName"][channel]["index"] for channel in selected]
+		state["newOrder"][prev_target_idx] = selected_idx
+		
+		#if len(selected)==1 and state["inputByName"][selected[0]]["used"]==False:
+			#state["inputByName"][selected[0]]["used"] = True
+			#state["missingChannelsIndex"][state["fillingCount"]] = -1
+		
+		print('Selection for missing channel "{}"({}): {}'.format(prev_target_name, prev_target_idx, selected))
+		
+		# update next round
+		state["fillingCount"] += 1
+		if state["fillingCount"] <= state["totalFillingNum"]:
+			target_idx = state["missingChannelsIndex"][state["fillingCount"]]
+			target_name = state["templateByIndex"][target_idx]
+			chkbox_label = target_name+' ('+str(state["fillingCount"]+1)+'/'+str(state["totalFillingNum"]+1)+')'
+			btn_label = "Submit" if state["fillingCount"]==state["totalFillingNum"] else "Next"
 			
-			# update next round
-			state_obj["fillingCount"] += 1
-			if state_obj["fillingCount"] <= state_obj["totalFillingNum"]:
-				target_idx = state_obj["missingChannelsIndex"][state_obj["fillingCount"]]
-				target_name = state_obj["templateByIndex"][target_idx]
-				chkbox_label = target_name+' ('+str(state_obj["fillingCount"]+1)+'/'+str(state_obj["totalFillingNum"]+1)+')'
-				btn_label = "Submit" if state_obj["fillingCount"]==state_obj["totalFillingNum"] else "Next"
-				
-				return {state_json : state_obj,
-						chs_chkbox : gr.CheckboxGroup(value=[], label=chkbox_label),
-						next_btn : gr.Button(btn_label)}
-			else:
-				state_obj["state"] = "finished"
-				reorder_data(raw_data, state_obj["newOrder"], fill_mode, state_obj)
-				
-				missing_channels = []
-				for idx in state_obj["missingChannelsIndex"]:
-				    if idx != -1:
-				        missing_channels.append(state_obj["templateByIndex"][idx])
-				missing_channels = ', '.join(missing_channels)
-					
-				return {state_json : state_obj,
-						chs_chkbox : gr.CheckboxGroup(visible=False),
-						next_btn : gr.Button(visible=False),
-						res_md : gr.Markdown(visible=True),
-			        	miss_txtbox : gr.Textbox(value=missing_channels, visible=True),
-			        	tpl_loc_file : gr.File(visible=True),
-						run_btn : gr.Button(interactive=True)}
+			return {state_json : state,
+					chkbox_group : gr.CheckboxGroup(value=[], label=chkbox_label),
+					next_btn : gr.Button(btn_label)}
+		else:
+			state["state"] = "finished"			
+			return {state_json : state,
+					chkbox_group : gr.CheckboxGroup(visible=False),
+					next_btn : gr.Button(visible=False),
+					res_md : gr.Markdown(visible=True),
+					run_btn : gr.Button(interactive=True)}
 	
 	next_btn.click(
-	    fn = check_next, 
-	    inputs = [state_json, chs_chkbox, in_raw_data, in_fill_mode],
-	    outputs = [state_json, chs_chkbox, next_btn, run_btn, res_md, miss_txtbox, tpl_loc_file]
-	    
+	    fn = check_next,
+	    inputs = [state_json, chkbox_group, in_raw_data, in_fill_mode],
+	    outputs = [state_json, chkbox_group, next_btn, run_btn, res_md]
+	
 	).success(
 	    fn = show_montage, 
 	    inputs = [state_json, in_raw_loc], 
 	    outputs = [state_json, tpl_montage, map_montage]
 	)
-    
-    
-	@run_btn.click(inputs=[state_json, in_raw_data, in_model_name], outputs=out_denoised_data)
-	def run_model(state_obj, raw_file, model_name):    	
-		filepath = state_obj["filepath"]
-    	
-		input_name = os.path.basename(str(raw_file))
+	
+	@run_btn.click(inputs = [state_json, in_raw_data, in_model_name, in_fill_mode], outputs = out_denoised_data)
+	def run_model(state, raw_data, model_name, fill_mode):
+		filepath = state["filepath"]
+		samplerate = state["sampleRate"]
+		
+		#if batch > total_batch:
+			#return {batch_md : gr.Markdown("error", visible=True)}
+		
+		input_name = os.path.basename(str(raw_data))
 		output_name = os.path.splitext(input_name)[0]+'_'+model_name+'.csv'
 		
-		if model_name == "(mapped data)":
-		    return filepath + 'mapped.csv'
-    	
-    	# step1: Data preprocessing
-		total_file_num = utils.preprocessing(filepath, 'mapped.csv', 256)
+		while(state["runnigState"] != "finished"):
+			if state["batchCount"] > state["totalBatchNum"]:
+				break
+			if state["batchCount"] > 1:
+				state["runnigState"] = "stage2"
+				state = mapping_stage2(state, fill_mode)
+			state["batchCount"] += 1
+			
+			reorder_to_template(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(state, filepath+'denoised.csv', filepath+output_name)
+			
+			if model_name == "(mapped data)":
+				return {out_denoised_data : filepath + 'mapped.csv'}
+			elif model_name == "(denoised data)":
+				return {out_denoised_data : filepath + 'denoised.csv'}
 		
-		# step2: Signal reconstruction
-		utils.reconstruct(model_name, total_file_num, filepath, output_name)
+		return {out_denoised_data : filepath + output_name}
 		
-		return filepath + output_name
-    
-
+	
 if __name__ == "__main__":
     demo.launch()

channel_mapping.py

@@ -2,15 +2,19 @@ import utils
 import time
 import os
 import numpy as np
+import gradio as gr
 
 import mne
 from mne.channels import read_custom_montage
-
-def reorder_data(filename, old_idx, fill_mode, state_obj):
-	old_data = utils.read_train_data(filename)
-	new_data = np.zeros((30, old_data.shape[1]))
-	new_filename = state_obj["filepath"]+'mapped.csv'
-	#print('old data shape: ', old_data.shape)
+from scipy.interpolate import Rbf
+from scipy.optimize import linear_sum_assignment
+from sklearn.neighbors import NearestNeighbors
+
+def reorder_to_template(state, filename):
+	old_idx = state["newOrder"]
+	old_data = utils.read_train_data(filename) # original raw data
+	new_data = np.zeros((30, old_data.shape[1])) # reordered raw data
+	new_filename = state["filepath"]+'mapped.csv'
 	
 	zero_arr = np.zeros((1, old_data.shape[1]))
 	old_data = np.concatenate((old_data, zero_arr), axis=0)
@@ -24,286 +28,278 @@ def reorder_data(filename, old_idx, fill_mode, state_obj):
 		else:
 			tmp_data = [old_data[j, :] for j in curr_idx_set]
 			new_data[i, :] = np.mean(tmp_data, axis=0)
-
-	#print('new data shape: ', new_data.shape)
+	
+	print('old.shape, new.shape: ', old_data.shape, new_data.shape)
 	utils.save_data(new_data, new_filename)
 	return
 
+def reorder_to_origin(state, filename, new_filename):
+	old_idx = state["newOrder"]
+	old_data = utils.read_train_data(filename) # denoised data
+	template_order = state["templateByIndex"]
+	
+	if state["runnigState"] == "stage1":
+		new_data = np.zeros((len(state["inputByName"]), 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 doesn't exist
+		if len(idx_set)==1 and state["templateByName"][channel]["matched"]==True:
+			new_data[idx_set[0], :] = old_data[i, :]
+	
+	print('old.shape, new.shape: ', old_data.shape, new_data.shape)
+	utils.save_data(new_data, new_filename)
+	return
 
 class Channel:
 
-	def __init__(self, index, name=None, used=False, coord=None, css_position=None, topo_index=None, topo_position=None):
-
+	def __init__(self, index, name=None, matched=False, assigned=False, coord=None, css_position=None):
 		self.name = name
 		self.index = index
-		self.used = used
+		self.matched = matched
+		self.assigned = assigned # for input channels
 		self.coord = coord
 		self.css_position = css_position
-		self.topo_index = topo_index
-		self.topo_position = topo_position
-
-	def prefix(self):
-		ret = ''.join(filter(str.isalpha, self.name))
-		return ret[:len(ret) - 1] if ret[-1] == 'Z' else ret
 
-	def suffix(self):
-		return -1 if self.name[-1] == 'Z' else int(''.join(filter(str.isdigit, self.name)))
 
-
-def pack_data(new_idx, missing_channels, tpl_dict, in_dict, tpl_ordered_name, in_ordered_name):
-	
-	return {
-		"newOrder" : [([i] if i!=-1 else []) for i in new_idx],
-		"missingChannelsIndex" : missing_channels,
-		"templateByName" : {k : v.__dict__ for k,v in tpl_dict.items()},  # dict, {name:object}
-		"templateByIndex" : tpl_ordered_name, # list
-		"inputByName" : {k : v.__dict__ for k,v in in_dict.items()},
-		"inputByIndex" : in_ordered_name
-	}
-
-def mapping(data_file, loc_file, fill_mode):
-	second1 = time.time()
-	
-	data = utils.read_train_data(data_file)
-	
-	template_dict = {}
-	input_dict = {}
+def read_montage_data(loc_file):
+		
 	template_montage = read_custom_montage("./template_chanlocs.loc")
 	input_montage = read_custom_montage(loc_file)
+	template_dict = {}
+	input_dict = {}
 	
 	montages = [template_montage, input_montage]
 	dicts = [template_dict, input_dict]
 	num = [30, len(input_montage.ch_names)]
 	
 	for i in range(2):
-		fig = montages[i].plot()
-		fig.set_size_inches(5.6, 5.6)
-		ax = fig.axes[0]
-		ax.set_aspect('equal')
-		ax.figure.canvas.draw() #update the figure
-		coords = ax.collections[0].get_offsets().data
-		abs_coords = ax.transData.transform(coords)
-		#print("abs_coords)
 		for j in range(num[i]):
 			channel = montages[i].ch_names[j]
+			montages[i].rename_channels({channel: str.upper(channel)}) # convert all channel names to uppercase
 	        
-	        # convert all channel names to uppercase
-			montages[i].rename_channels({channel: str.upper(channel)})
-	        
-			css_left = (abs_coords[j][0]-11)/560
-			css_bottom = (abs_coords[j][1]-7)/560
 			channel = str.upper(channel)
-			dicts[i][channel] = Channel(index=j,
-										name=channel,
-										coord=montages[i].get_positions()['ch_pos'][channel],
-										css_position=[str(round(css_left*100, 2))+"%", str(round(css_bottom*100, 2))+"%"]
-								)
+			dicts[i][channel] = Channel(index=j, name=channel, coord=montages[i].get_positions()['ch_pos'][channel])
+	
+	return template_montage, input_montage, template_dict, input_dict
+
+def align_coords(state, template_montage, input_montage):
+
+	template_dict = state["templateByName"]
+	input_dict = state["inputByName"]
+	template_order = state["templateByIndex"]
+	input_order = state["inputByIndex"]
+	matched = [channel for channel in input_dict if input_dict[channel]["matched"]==True]
 	
+	# 2-d (fot the indication of missing template channel's position when fill_mode:'mean_manual')
+	fig = [template_montage.plot(), input_montage.plot()]
+	fig[0].set_size_inches(5.6, 5.6)
+	fig[1].set_size_inches(5.6, 5.6)
 	
-	new_idx = [-1]*30
-	missing_channels = []
-	exact_missing_channels = []
-	z_row_idx = data.shape[0]
+	ax = [fig[0].axes[0], fig[1].axes[0]]
+	ax[0].set_aspect('equal')
+	ax[1].set_aspect('equal')
+	ax[0].figure.canvas.draw() #update the figure
+	ax[1].figure.canvas.draw()
+	
+	# get the original coords
+	all_tpl = ax[0].transData.transform(ax[0].collections[0].get_offsets().data) # display coords (px)
+	all_in= ax[1].transData.transform(ax[1].collections[0].get_offsets().data)
+	matched_tpl = np.array([all_tpl[template_dict[channel]["index"]] for channel in matched])
+	matched_in = np.array([all_in[input_dict[channel]["index"]] for channel in matched])
+	
+	# transform the xy axis (template's -> input's)
+	rbf_x = Rbf(matched_tpl[:,0], matched_tpl[:,1], matched_in[:,0], function='thin_plate')
+	rbf_y = Rbf(matched_tpl[:,0], matched_tpl[:,1], matched_in[:,1], function='thin_plate')
+	
+	# apply to all template channels
+	transformed_tpl_x = rbf_x(all_tpl[:,0], all_tpl[:,1])
+	transformed_tpl_y = rbf_y(all_tpl[:,0], all_tpl[:,1])
+	#transformed_tpl = np.vstack((transformed_tpl_x, transformed_tpl_y)).T
+	
+	# update input, template's position
+	for i, channel in enumerate(template_order):
+		css_left = (transformed_tpl_x[i]-11)/560
+		css_bottom = (transformed_tpl_y[i]-7)/560
+		template_dict[channel]["css_position"] = [str(round(css_left*100, 2))+"%", str(round(css_bottom*100, 2))+"%"]
+	for i, channel in enumerate(input_order):
+		css_left = (all_in[i][0]-11)/560
+		css_bottom = (all_in[i][1]-7)/560
+		input_dict[channel]["css_position"] = [str(round(css_left*100, 2))+"%", str(round(css_bottom*100, 2))+"%"]
+	
+	
+	# 3-d (to use KNN)
+	# get the original coords
+	all_tpl = np.array([template_dict[channel]["coord"].tolist() for channel in template_order])
+	all_in = np.array([input_dict[channel]["coord"].tolist() for channel in input_order])
+	matched_tpl = np.array([all_tpl[template_dict[channel]["index"]] for channel in matched])
+	matched_in = np.array([all_in[input_dict[channel]["index"]] for channel in matched])
 	
+	# transform the xyz axis (input's -> template's)
+	rbf_x = Rbf(matched_in[:,0], matched_in[:,1], matched_in[:,2], matched_tpl[:,0], function='thin_plate')
+	rbf_y = Rbf(matched_in[:,0], matched_in[:,1], matched_in[:,2], matched_tpl[:,1], function='thin_plate')
+	rbf_z = Rbf(matched_in[:,0], matched_in[:,1], matched_in[:,2], matched_tpl[:,2], function='thin_plate')
 	
-	# STAGE_1
+	# apply to all input channels
+	transformed_in_x = rbf_x(all_in[:,0], all_in[:,1], all_in[:,2])
+	transformed_in_y = rbf_y(all_in[:,0], all_in[:,1], all_in[:,2])
+	transformed_in_z = rbf_z(all_in[:,0], all_in[:,1], all_in[:,2])
+	transformed_in = np.vstack((transformed_in_x, transformed_in_y, transformed_in_z)).T
+	
+	# update input's position
+	for i, channel in enumerate(input_order):
+		input_dict[channel]["coord"] = transformed_in[i].tolist()
+	
+	state.update({
+	    "templateByName" : template_dict,
+	    "inputByName" : input_dict,
+	})
+	
+	return state
+
+def fill_channels(state, fill_mode):
+
+	new_idx = state["newOrder"]
+	template_dict = state["templateByName"]
+	input_dict = state["inputByName"]
+	template_order = state["templateByIndex"]
+	input_order = state["inputByIndex"]
+	z_row_idx = state["dataShape"][0]
+	unmatched = [channel for channel in template_dict if template_dict[channel]["matched"]==False]
+	if unmatched == []:
+		return state
+	
+	if fill_mode == 'zero':
+		for channel in unmatched:
+			idx = template_dict[channel]["index"]
+			new_idx[idx] = [z_row_idx]
+	
+	elif fill_mode == 'mean_auto':
+		# use KNN to choose k nearest channels
+		in_coords = [input_dict[channel]["coord"] for channel in input_order]
+		in_coords = np.array([in_coords[i] for i in range(len(in_coords))])
+		
+		k = 4 if len(input_dict)>4 else len(input_dict)
+		knn = NearestNeighbors(n_neighbors=k, metric='euclidean')
+		knn.fit(in_coords)
+		
+		for channel in unmatched:
+			distances, indices = knn.kneighbors(template_dict[channel]["coord"].reshape(1,-1))
+			selected = [input_order[i] for i in indices[0]]
+			print(channel, ':', selected)
+			
+			idx = template_dict[channel]["index"]
+			new_idx[idx] = indices[0].tolist()
+
+	state["newOrder"] = new_idx
+	return state
+
+def mapping_stage1(state, data_file, loc_file, fill_mode):
+	second1 = time.time()
 	
-	# match the template's channel names with the input ones
-	finish_flag = 1
+	template_montage, input_montage, template_dict, input_dict = read_montage_data(loc_file)
+	template_order = template_montage.ch_names
+	new_idx = [[]]*30
+	missing_channels = []
 	alias = {
 		'T3': 'T7',
 		'T4': 'T8',
 		'T5': 'P7',
 		'T6': 'P8',
-		'TP7': 'T5\'',
-		'TP8': 'T6\'',
+		#'TP7': 'T5\'',
+		#'TP8': 'T6\'',
 	}
 	
-	for i in range(30):
-		channel = template_montage.ch_names[i]
-		if channel not in input_dict.keys() | alias.keys():
-			exact_missing_channels.append(i)
-			finish_flag = 0
-			continue
-
-		if channel not in input_dict and channel in alias:
-			if alias[channel] in input_dict:
-				template_montage.rename_channels({channel: alias[channel]})
-				template_dict[alias[channel]] = template_dict.pop(channel)
-				template_dict[alias[channel]].name = alias[channel]
-				channel = alias[channel]
-			else:
-				exact_missing_channels.append(i)
-				finish_flag = 0
-				continue
-
-		new_idx[i] = input_dict[channel].index
-		input_dict[channel].used = True
+	# match the names of input channels -> template channels
+	for i, channel in enumerate(template_order):
+		if channel in alias and alias[channel] in input_dict:
+			template_montage.rename_channels({channel: alias[channel]})
+			template_dict[alias[channel]] = template_dict.pop(channel)
+			template_dict[alias[channel]].name = alias[channel]
+			channel = alias[channel]
+		
+		if channel in input_dict:
+			new_idx[i] = [input_dict[channel].index]
+			
+			template_dict[channel].matched = True
+			input_dict[channel].matched = True
+			input_dict[channel].assigned = True
+		else:
+			missing_channels.append(i)
+	
+	state.update({
+	    "newOrder" : new_idx,
+	    "missingChannelsIndex" : missing_channels,
+	    "templateByName" : {k : v.__dict__ for k,v in template_dict.items()},
+	    "inputByName" : {k : v.__dict__ for k,v in input_dict.items()},
+	    "templateByIndex" : template_montage.ch_names,
+		"inputByIndex" : input_montage.ch_names
+	})
+	
+	# align input, template's coordinates
+	state = align_coords(state, template_montage, input_montage)
+	# fill the unmatched channels
+	state = fill_channels(state, fill_mode)
+	
+	second2 = time.time()
+	print('Mapping (stage1) finished in',second2 - second1,'s.')
+	return state
 
-	if finish_flag == 1:
-		second2 = time.time()
-		print('Finish at stage 1 ! (',second2 - second1,'s)')
-		#print('new idx order:', new_idx)
+def mapping_stage2(state, fill_mode):
+	second1 = time.time()
 
-		channels_obj = pack_data(new_idx, [],
-						template_dict, input_dict,
-						template_montage.ch_names, input_montage.ch_names)
-		channels_obj.update({"CZImputed" : False})
-		return channels_obj
-		
-	elif fill_mode == "mean":
-		channels_obj = pack_data(new_idx, exact_missing_channels,
-						template_dict, input_dict,
-						template_montage.ch_names, input_montage.ch_names)
-		channels_obj.update({"CZImputed" : False})
-		return channels_obj
+	template_dict = state["templateByName"]
+	input_dict = state["inputByName"]
+	template_order = state["templateByIndex"]
+	unassigned = [channel for channel in input_dict if input_dict[channel]["assigned"]==False]
+	if unassigned == []:
+		state["runnigState"] = "finished"
+		return state
 	
+	tpl_coords = np.array([template_dict[channel]["coord"] for channel in template_order])
+	unassigned_coords = np.array([input_dict[channel]["coord"] for channel in unassigned])
 	
+	# set all tpl.matched to False
+	for channel in template_dict:
+		template_dict[channel]["matched"] = False
 	
+	# initialize the cost matrix
+	if len(unassigned) < 30:
+		cost_matrix = np.full((30, 30), 10000) # add dummy channels to ensure num_col > num_row
+	else:
+		cost_matrix = np.zeros((30, len(unassigned)))
+	for i in range(30):
+		for j in range(len(unassigned)):
+			cost_matrix[i][j] = np.linalg.norm(tpl_coords[i] - unassigned_coords[j]) # Euclidean distance
 	
-	# STAGE_2
+	# use Hungarian Algorithm to find the minimum sum of distance of (input's coord to template's coord)...?
+	row_idx, col_idx = linear_sum_assignment(cost_matrix)
 	
-	# store channel positions in a 2-d array
-	template_topo_pos = []  
-	temporal_channels = []
-	temporal_row_prefix = ['FC', 'C', 'CP', 'P']
-
-	cnt = 0
-	for i in range(7):
-		tmp = []
-		for j in range(5):
-			if [i,j] in [[0,0],[0,2],[0,4],[6,0],[6,4]]:
-				tmp.append('')
-			else:
-				channel = template_montage.ch_names[cnt]
-				tmp.append(channel)
-
-				ver = 'front' if i<3 else 'center' if i==3 else 'back'
-				hor = 'left' if j<2 else 'center' if j==2 else 'right'
-				template_dict[channel].topo_index = [i, j]
-				template_dict[channel].topo_position = [ver, hor]
-
-				if i > 1 and j in [0, 4]:
-					temporal_channels.append(channel)
-				cnt += 1
-		template_topo_pos.append(tmp)
-		
-
-	# ensure that CZ is found or imputed by another channel
-	CZ_impute_flag = False
-	if 'CZ' not in input_dict and fill_mode=='adjacent':
-		CZ_impute_flag = True
-		min_dist = 1e5
-		for channel in input_montage.ch_names:
-			curr_x, curr_y, curr_z = input_dict[channel].coord.round(6)
-			if curr_x**2 + curr_y**2 < min_dist:
-				nearest_channel = channel
-				min_dist = curr_x**2 + curr_y**2
-
-		if input_dict[nearest_channel].used == True:
-			missing_channels.append(template_dict['CZ'].index)
-		input_dict[nearest_channel].used = True
-		input_dict['CZ'] = input_dict[nearest_channel]
-		print("CZ's nearest neighbor:", nearest_channel)
-
-
+	matches = []
+	new_idx = [[]]*30
 	for i in range(30):
-		if new_idx[i] != -1:
-			continue
-
-		channel = template_montage.ch_names[i]
-
-		curr_prefix = template_dict[channel].prefix()
-		curr_suffix = template_dict[channel].suffix()
-
-		curr_row = template_dict[channel].topo_index[0]
-		curr_col = template_dict[channel].topo_index[1]
-		curr_ver = template_dict[channel].topo_position[0]
-		curr_hor = template_dict[channel].topo_position[1]
-
-		impute_channel = ''
-
-		# if the current channel is a temporal channel
-		if channel in temporal_channels:
-			curr_prefix = temporal_row_prefix[temporal_channels.index(channel)//2]
-			curr_suffix = 7 if curr_hor=='left' else 8
-
-		if fill_mode == 'zero':
-
-			impute_channel = curr_prefix+str(1) if curr_hor=='center' else curr_prefix+str(curr_suffix-2)
-			if impute_channel not in input_dict or input_dict[impute_channel].used==True:
-				impute_channel = ''
-				new_idx[i] = z_row_idx
-				missing_channels.append(i)
-				continue
-		
-		elif fill_mode == 'adjacent':
+		if col_idx[i] < len(unassigned): # filter out dummy channels
+			matches.append([row_idx[i], col_idx[i]])
 			
-			ver_dir = 1 if curr_ver == 'front' else -1
-            
-			if curr_hor == 'center':  # FZ, FPZ...
-
-				if curr_prefix+str(1) in input_dict:  # ex: FZ<-F1
-					impute_channel = curr_prefix + str(1)
-					
-				elif template_topo_pos[curr_row+ver_dir][curr_col] in input_dict: # ex: front:FZ<-FCZ,
-				    impute_channel = template_topo_pos[curr_row+ver_dir][curr_col]
-					
-				elif curr_prefix+str(3) in input_dict: # ex: FZ<-F3
-					impute_channel = curr_prefix + str(3)
-					
-				else:
-					impute_channel = 'CZ'
-					
-			elif curr_hor == 'left' or curr_hor == 'right':
-
-				ver_ctrl = 1 if curr_ver=='front' else 2 if curr_ver=='back' else 3 # bit0: row+1, bit1: row-1
-
-				# search horizontally
-				cnt = 0
-				tmp_suffix = curr_suffix
-				while tmp_suffix > 0:  # ex: F7<-F5/F3/F1
-					tmp_suffix = curr_suffix - 2*cnt
-					if curr_prefix+str(tmp_suffix) in input_dict:
-						impute_channel = curr_prefix + str(tmp_suffix)
-						break
-
-					if cnt == 2:
-						# check row+1/row-1
-						if ver_ctrl&1 and template_topo_pos[curr_row+1][curr_col] in input_dict:
-							impute_channel = template_topo_pos[curr_row+1][curr_col]
-							break
-						if ver_ctrl&2 and template_topo_pos[curr_row-1][curr_col] in input_dict:
-							impute_channel = template_topo_pos[curr_row-1][curr_col]
-							break
-					cnt += 1
-
-				# search vertically
-				if impute_channel == '':
-					cnt = 0
-					tmp_row = curr_row + ver_dir
-					while tmp_row-ver_dir != 3:  # terminate if the last channel is a middle one
-						if template_topo_pos[tmp_row][curr_col] in input_dict:
-							impute_channel = template_topo_pos[tmp_row][curr_col]
-							break
-						tmp_row += ver_dir
-
-				# if still cannot find available channel...
-				if impute_channel == '':
-					impute_channel = 'CZ'
-
-		new_idx[i] = input_dict[impute_channel].index
-		if input_dict[impute_channel].used == True:  # this channel is shared with others
-			missing_channels.append(i)
-		input_dict[impute_channel].used = True
+			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[i] = [input_dict[in_channel]["index"]]
 	
-	second2 = time.time()
-	print('Finish at stage 2 ! (',second2 - second1,'s)')
-	#print('new_idx:', new_idx)
+	state.update({
+	    "newOrder" : new_idx,
+	    "templateByName" : template_dict,
+	    "inputByName" : input_dict
+	})
 	
-	channels_obj = pack_data(new_idx, missing_channels,
-						template_dict, input_dict,
-						template_montage.ch_names, input_montage.ch_names)
-	channels_obj.update({"CZImputed" : CZ_impute_flag})
-	return channels_obj
+	# fill the unmatched channels
+	state = fill_channels(state, fill_mode)
 	
+	second2 = time.time()
+	print(f'Mapping (stage2-{state["batchCount"]-1}) finished in {second2 - second1}s.')
+	return state
+

utils.py

@@ -42,6 +42,34 @@ def resample(signal, fs):
 
     return signal_new
 
+# original -> 256Hz or 256Hz -> original
+def resample_(signal, current_fs, target_fs):
+    fs = current_fs
+    # downsample the signal to the target sample rate
+    if fs>target_fs:
+        fs_down = target_fs # Desired sample rate
+        q = int(fs / fs_down) # Downsampling factor
+        signal_new = []
+        for ch in signal:
+            x_down = decimate(ch, q)
+            signal_new.append(x_down)
+
+    # upsample the signal to the target sample rate
+    elif fs<target_fs:
+        fs_up = target_fs  # Desired sample rate
+        p = int(fs_up / fs)  # Upsampling factor 
+        signal_new = []
+        for ch in signal:
+            x_up = resample_poly(ch, p, 1)
+            signal_new.append(x_up)
+
+    else:
+        signal_new = signal
+
+    signal_new = np.array(signal_new).astype(np.float64)
+
+    return signal_new
+
 def FIR_filter(signal, lowcut, highcut):
     fs = 256.0
     # Number of FIR filter taps
@@ -96,11 +124,14 @@ def glue_data(file_name, total, output):
             raw_data[:, 1] = smooth
             gluedata = np.append(gluedata, raw_data, axis=1)
     #print(gluedata.shape)
+    '''
     filename2 = output
     with open(filename2, 'w', newline='') as csvfile:
         writer = csv.writer(csvfile)
         writer.writerows(gluedata)
         #print("GLUE DONE!" + filename2)
+    '''
+    return gluedata
 
 
 def save_data(data, filename):
@@ -189,10 +220,11 @@ def preprocessing(filepath, filename, samplerate):
         print(e)
     
     # read data
-    signal = read_train_data(filepath+'/'+filename)
+    signal = read_train_data(filepath+filename)
     #print(signal.shape)
     # resample
     signal = resample(signal, samplerate)
+    #signal = resample_(signal, samplerate, 256)
     #print(signal.shape)
     # FIR_filter
     signal = FIR_filter(signal, 1, 50)
@@ -204,7 +236,7 @@ def preprocessing(filepath, filename, samplerate):
 
 
 # model = tf.keras.models.load_model('./denoise_model/')
-def reconstruct(model_name, total, filepath, outputfile):
+def reconstruct(model_name, total, filepath, outputfile, samplerate):
     # -------------------decode_data---------------------------
     second1 = time.time()
     for i in range(total):
@@ -224,9 +256,14 @@ def reconstruct(model_name, total, filepath, outputfile):
         save_data(d_data, outputname)
 
     # --------------------glue_data----------------------------
-    glue_data(filepath+"/temp2/", total, filepath+'/'+outputfile)
+    signal = glue_data(filepath+"/temp2/", total, filepath+outputfile)
+    #print(signal.shape)
     # -------------------delete_data---------------------------
     dataDelete(filepath+"/temp2/")
+    # --------------------resample-----------------------------
+    signal = resample_(signal, 256, samplerate) # 256Hz -> original sampling rate
+    #print(signal.shape)
+    save_data(signal, filepath+outputfile)
     second2 = time.time()
 
     print("Using", model_name,"model to reconstruct", outputfile, " has been success in", second2 - second1, "sec(s)")