| import numpy as np | |
| from scipy.io import savemat | |
| import h5py | |
| from bcgunet import bcgunet | |
| import platform | |
| import os | |
| import time | |
| import os | |
| import gradio as gr | |
| import matplotlib | |
| import matplotlib.pyplot as plt | |
| matplotlib.use("agg") | |
| dir = os.path.dirname(os.path.realpath(__file__)) + "/tmp" | |
| os.makedirs(dir, exist_ok=True) | |
| def run( | |
| files: list[bytes], | |
| lr: float, | |
| winsec: int, | |
| iters: int, | |
| onecycle: bool, | |
| ecg: str, | |
| bce: str, | |
| eeg: str, | |
| ) -> tuple[list[str], str]: | |
| task = os.path.join(dir, str(int(time.time()))) | |
| os.makedirs(task) | |
| outputs = [] | |
| for i, file in enumerate(files): | |
| input = os.path.join(task, str(i) + ".mat") | |
| with open(input, "wb") as o: | |
| o.write(file) | |
| output = os.path.join(task, str(i) + "_clean.mat") | |
| mat = h5py.File(input, "r") | |
| ECG = np.array(mat[ecg]).flatten() | |
| EEG = np.array(mat[bce]).T | |
| EEG_unet = bcgunet.run( | |
| EEG, | |
| ECG, | |
| iter_num=iters, | |
| winsize_sec=winsec, | |
| lr=lr, | |
| onecycle=onecycle, | |
| ) | |
| result = dict() | |
| result[eeg] = EEG_unet | |
| savemat(output, result, do_compression=True) | |
| outputs.append(output) | |
| if i == 0: | |
| plt.figure(figsize=(12, 6), dpi=300) | |
| plt.plot(EEG[19, :10000], "b.-", label="Orig EEG") | |
| plt.plot(EEG_unet[19, :10000], "g.-", label="U-Net") | |
| plt.legend() | |
| plt.title("BCG Unet") | |
| plt.xlabel("Time (samples)") | |
| plot = os.path.join(task, str(i) + ".png") | |
| plt.savefig(plot) | |
| return outputs, plot | |
| def main(): | |
| app = gr.Interface( | |
| title="BCG Unet", | |
| description="BCGunet: Suppressing BCG artifacts on EEG collected inside an MRI scanner", | |
| fn=run, | |
| inputs=[ | |
| gr.File( | |
| label="Input Files (.mat)", | |
| type="binary", | |
| file_types=["mat"], | |
| file_count=["multiple", "directory"], | |
| ), | |
| gr.Slider( | |
| label="Learning Rate", minimum=1e-5, maximum=1e-1, step=1e-5, value=1e-3 | |
| ), | |
| gr.Slider( | |
| label="Window Size (seconds)", minimum=1, maximum=10, step=1, value=2 | |
| ), | |
| gr.Slider( | |
| label="Number of Iterations", | |
| minimum=1000, | |
| maximum=10000, | |
| step=1000, | |
| value=5000, | |
| ), | |
| gr.Checkbox( | |
| label="One Cycle Scheduler", | |
| value=True, | |
| ), | |
| gr.Textbox( | |
| label="Variable name for ECG (input)", | |
| value="ECG", | |
| ), | |
| gr.Textbox( | |
| label="Variable name for BCG corropted EEG (input)", | |
| value="EEG_before_bcg", | |
| ), | |
| gr.Textbox( | |
| label="Variable name for clean EEG (output)", | |
| value="EEG_clean", | |
| ), | |
| ], | |
| outputs=[ | |
| gr.File(label="Output File", file_count="multiple"), | |
| gr.Image(label="Output Image", type="filepath"), | |
| ], | |
| allow_flagging="never", | |
| ) | |
| app.launch() | |
| if __name__ == "__main__": | |
| main() | |
| if platform.system() == "Windows": | |
| os.system("pause") | |