update

app.py

@@ -55,7 +55,7 @@ chkbox_js = """
 		position: relative;
 		width: 560px;
 		height: 560px;
-		background: url("file=${app_state.files.raw_montage}");
+		background: url("file=${app_state.filenames.raw_montage}");
 	`;
 	
 	
@@ -191,7 +191,6 @@ with gr.Blocks() as demo:
 										label="Imputation")
 					map_btn = gr.Button("Mapping")
 			
-			#indic
 			chkbox_group = gr.CheckboxGroup(elem_id="chkbox-group", label="", visible=False)
 			next_btn = gr.Button("Next", interactive=False, visible=False)
 			
@@ -228,7 +227,7 @@ with gr.Blocks() as demo:
 									scale=2)
 				run_btn = gr.Button(scale=1, interactive=False)
 			batch_md = gr.Markdown(visible=False)
-			out_denoised_data = gr.File(label="Denoised data")
+			out_denoised_data = gr.File(label="Denoised data", visible=False)
 	
 				
 	with gr.Row():
@@ -245,7 +244,7 @@ with gr.Blocks() as demo:
     
 	#demo.load(js=js)
 	
-	def reset_layout(raw_data, samplerate):
+	def reset1(raw_data, samplerate):
 		# establish temp folder
 		filepath = os.path.dirname(str(raw_data))
 		try:
@@ -259,7 +258,7 @@ with gr.Blocks() as demo:
 		data = utils.read_train_data(raw_data)
 		app_state = {
 			"filepath": filepath+"/temp_data/",
-			"files": {},
+			"filenames": {},
 			"sampleRate": int(samplerate),
 			
 		}
@@ -275,7 +274,8 @@ with gr.Blocks() as demo:
 				tpl_montage : gr.Image(visible=False),
 				map_montage : gr.Image(value=None, visible=False),
 				res_md : gr.Markdown(visible=False),
-				batch_md : gr.Markdown(visible=False)}
+				batch_md : gr.Markdown(visible=False),
+				out_denoised_data : gr.File(visible=False)}
     	
 	def mapping_result(app_state, channel_info, fill_mode):
 	
@@ -283,7 +283,6 @@ with gr.Blocks() as demo:
 		matched_num = 30 - len(channel_info["missingChannelsIndex"])
 		batch_num = math.ceil((in_num-matched_num)/30) + 1
 		app_state.update({
-			"runnigState" : "stage1",
 			"batchCount" : 1,
 			"totalBatchNum" : batch_num
 		})
@@ -295,11 +294,11 @@ with gr.Blocks() as demo:
 			})
 			#print("Missing channels:", channel_info["missingChannelsIndex"])
 			return {app_state_json : app_state,
-					#chkbox_group : gr.CheckboxGroup(visible=True),
 					next_btn : gr.Button(visible=True)}
 		else:
-			app_state["state"] = "finished"
-			
+			app_state.update({
+				"state" : "finished"
+			})
 			return {app_state_json : app_state,
 					res_md : gr.Markdown(visible=True),
 					run_btn : gr.Button(interactive=True)}
@@ -318,7 +317,7 @@ with gr.Blocks() as demo:
 		
 		if app_state["state"] == "initializing":
 			filename = filepath+"raw_montage_"+str(random.randint(1,10000))+".png"
-			app_state["files"]["raw_montage"] = filename
+			app_state["filenames"]["raw_montage"] = filename
 			raw_fig = raw_montage.plot()
 			raw_fig.set_size_inches(5.6, 5.6)
 			raw_fig.savefig(filename, pad_inches=0)
@@ -327,7 +326,7 @@ with gr.Blocks() as demo:
 		
 		elif app_state["state"] == "finished":
 			filename = filepath+"mapped_montage_"+str(random.randint(1,10000))+".png"
-			app_state["files"]["map_montage"] = filename
+			app_state["filenames"]["map_montage"] = filename
 			
 			show_names= []
 			for channel in channel_info["inputByName"]:
@@ -361,9 +360,10 @@ with gr.Blocks() as demo:
 
 	
 	map_btn.click(
-	    fn = reset_layout,
+	    fn = reset1,
 	    inputs = [in_raw_data, in_sample_rate],
-		outputs = [app_state_json, channel_info_json, chkbox_group, next_btn, run_btn, tpl_montage, map_montage, res_md, batch_md]
+		outputs = [app_state_json, channel_info_json, chkbox_group, next_btn, run_btn,
+					tpl_montage, map_montage, res_md, batch_md, out_denoised_data]
 					
 	).success(
 		fn = mapping_stage1, 
@@ -401,7 +401,7 @@ with gr.Blocks() as demo:
 		prev_target_name = channel_info["templateByIndex"][prev_target_idx]
 		
 		selected_idx = [channel_info["inputByName"][channel]["index"] for channel in selected]
-		app_state["newOrder"][prev_target_idx] = selected_idx
+		app_state["stage1NewOrder"][prev_target_idx] = selected_idx
 		
 		#if len(selected)==1 and channel_info["inputByName"][selected[0]]["used"]==False:
 			#channel_info["inputByName"][selected[0]]["used"] = True
@@ -450,22 +450,47 @@ with gr.Blocks() as demo:
 		outputs = []
 	)
 	
-	@run_btn.click(inputs = [app_state_json, channel_info_json, in_raw_data, in_model_name, in_fill_mode],
-				   outputs = [batch_md, out_denoised_data])
-	def run_model(app_state, channel_info, raw_data, model_name, fill_mode):
+	def delete_file(filename):
+		try:
+			os.remove(filename)
+		except OSError as e:
+			print(e)
+	
+	def reset2(app_state, raw_data, model_name):
 		filepath = app_state["filepath"]
-		samplerate = app_state["sampleRate"]
-		
 		input_name = os.path.basename(str(raw_data))
 		output_name = os.path.splitext(input_name)[0]+'_'+model_name+'.csv'
 		
-		while(app_state["runnigState"] != "finished"):
-			if app_state["batchCount"] > app_state["totalBatchNum"]:
-				app_state["runnigState"] = "finished"
-				break
+		app_state["filenames"]["denoised"] = filepath + output_name
+		app_state.update({
+			"runnigState" : "stage1",
+			"batchCount" : 1,
+			"stage2NewOrder" : [[]]*30
+		})
+		
+		delete_file(filepath+'mapped.csv')
+		delete_file(filepath+'denoised.csv')
+		return {app_state_json : app_state,
+				run_btn : gr.Button(interactive=False),
+				batch_md : gr.Markdown(visible=False),
+				out_denoised_data : gr.File(visible=False)}
+	
+	def run_model(app_state, channel_info, raw_data, model_name, fill_mode):
+		filepath = app_state["filepath"]
+		samplerate = app_state["sampleRate"]
+		new_filename = app_state["filenames"]["denoised"]
+		
+		while app_state["runnigState"] != "finished":
+			#if app_state["batchCount"] > app_state["totalBatchNum"]:
+				#app_state["runnigState"] = "finished"
+				#break
+			md = 'Running model('+str(app_state["batchCount"])+'/'+str(app_state["totalBatchNum"])+')...'
+			yield {batch_md : gr.Markdown(md, visible=True)}
+			
 			if app_state["batchCount"] > 1:
-				app_state["runnigState"] = "stage2"
 				app_state, channel_info = mapping_stage2(app_state, channel_info, fill_mode)
+				if app_state["runnigState"] == "finished":
+					break
 			app_state["batchCount"] += 1
 			
 			reorder_to_template(app_state, raw_data)
@@ -473,15 +498,30 @@ with gr.Blocks() as demo:
 			total_file_num = utils.preprocessing(filepath, 'mapped.csv', samplerate)
 			# step2: Signal reconstruction
 			utils.reconstruct(model_name, total_file_num, filepath, 'denoised.csv', samplerate)
-			reorder_to_origin(app_state, channel_info, filepath+'denoised.csv', filepath+output_name)
+			reorder_to_origin(app_state, channel_info, filepath+'denoised.csv', new_filename)
+			
+			#if model_name == "(mapped data)":
+				#return {out_denoised_data : filepath + 'mapped.csv'}
+			#elif model_name == "(denoised data)":
+				#return {out_denoised_data : filepath + 'denoised.csv'}
 			
-			if model_name == "(mapped data)":
-				return {out_denoised_data : filepath + 'mapped.csv'}
-			elif model_name == "(denoised data)":
-				return {out_denoised_data : filepath + 'denoised.csv'}
+			delete_file(filepath+'mapped.csv')
+			delete_file(filepath+'denoised.csv')
 			
-		return {out_denoised_data : filepath + output_name}
+		yield {run_btn : gr.Button(interactive=True),
+				batch_md : gr.Markdown(visible=False),
+				out_denoised_data : gr.File(new_filename, visible=True)}
+	
+	run_btn.click(
+		fn = reset2,
+		inputs = [app_state_json, in_raw_data, in_model_name],
+		outputs = [app_state_json, run_btn, batch_md, out_denoised_data]
 		
+	).success(
+		fn = run_model,
+		inputs = [app_state_json, channel_info_json, in_raw_data, in_model_name, in_fill_mode],
+		outputs = [run_btn, batch_md, out_denoised_data]
+	)
 	
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()

channel_mapping.py

@@ -11,10 +11,12 @@ from scipy.optimize import linear_sum_assignment
 from sklearn.neighbors import NearestNeighbors
 
 def reorder_to_template(app_state, filename):
-	old_idx = app_state["newOrder"]
+	old_idx = app_state["stage1NewOrder"] if app_state["runnigState"]=="stage1" else app_state["stage2NewOrder"]
 	old_data = utils.read_train_data(filename) # original raw data
 	new_data = np.zeros((30, old_data.shape[1])) # reordered raw data
 	new_filename = app_state["filepath"]+'mapped.csv'
+	#print('new order 1:', app_state["stage1NewOrder"])
+	#print('new order 2:', app_state["stage2NewOrder"])
 	
 	zero_arr = np.zeros((1, old_data.shape[1]))
 	old_data = np.concatenate((old_data, zero_arr), axis=0)
@@ -34,7 +36,7 @@ def reorder_to_template(app_state, filename):
 	return
 
 def reorder_to_origin(app_state, channel_info, filename, new_filename):
-	old_idx = app_state["newOrder"]
+	old_idx = app_state["stage1NewOrder"] if app_state["runnigState"]=="stage1" else app_state["stage2NewOrder"]
 	old_data = utils.read_train_data(filename) # denoised data
 	template_order = channel_info["templateByIndex"]
 	
@@ -161,7 +163,7 @@ def align_coords(channel_info, template_montage, input_montage):
 
 def fill_channels(app_state, channel_info, fill_mode):
 
-	new_idx = app_state["newOrder"]
+	new_idx = app_state["stage1NewOrder"] if app_state["runnigState"]=="stage1" else app_state["stage2NewOrder"]
 	template_dict = channel_info["templateByName"]
 	input_dict = channel_info["inputByName"]
 	template_order = channel_info["templateByIndex"]
@@ -186,16 +188,18 @@ def fill_channels(app_state, channel_info, fill_mode):
 		knn.fit(in_coords)
 		
 		for channel in unmatched:
-			distances, indices = knn.kneighbors(template_dict[channel]["coord"].reshape(1,-1))
+			distances, indices = knn.kneighbors(np.array(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()
-
-	app_state.update({
-	    "newOrder" : new_idx
-	})
+	
+	if app_state["runnigState"] == "stage1":
+		app_state["stage1NewOrder"] = new_idx
+	else:
+		app_state["stage2NewOrder"] = new_idx
+    
 	return app_state
 
 def mapping_stage1(app_state, channel_info, data_file, loc_file, fill_mode):
@@ -239,7 +243,8 @@ def mapping_stage1(app_state, channel_info, data_file, loc_file, fill_mode):
 		"inputByIndex" : input_montage.ch_names
 	})
 	app_state.update({
-	    "newOrder" : new_idx
+	    "stage1NewOrder" : new_idx,
+	    "runnigState" : "stage1"
 	})
 	
 	# align input, template's coordinates
@@ -271,12 +276,13 @@ def mapping_stage2(app_state, channel_info, fill_mode):
 	
 	# initialize the cost matrix
 	if len(unassigned) < 30:
-		cost_matrix = np.full((30, 30), 10000) # add dummy channels to ensure num_col > num_row
+		cost_matrix = np.full((30, 30), 1e6) # 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
+			cost_matrix[i][j] = np.linalg.norm((tpl_coords[i]-unassigned_coords[j])*1000) # Euclidean distance
+			#print(cost_matrix[i][j], tpl_coords[i] - unassigned_coords[j])
 	
 	# 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)
@@ -292,13 +298,16 @@ def mapping_stage2(app_state, channel_info, fill_mode):
 			template_dict[tpl_channel]["matched"] = True
 			input_dict[in_channel]["assigned"] = True
 			new_idx[i] = [input_dict[in_channel]["index"]]
+			
+			print(template_order[row_idx[i]], '<-', unassigned[col_idx[i]])
 	
 	channel_info.update({
 	    "templateByName" : template_dict,
 	    "inputByName" : input_dict
 	})
 	app_state.update({
-	    "newOrder" : new_idx
+	    "stage2NewOrder" : new_idx,
+	    "runnigState" : "stage2"
 	})
 	
 	# fill the unmatched channels