diff --git a/bcgrun.py b/bcgrun.py index 7b0b766..62ae11f 100644 --- a/bcgrun.py +++ b/bcgrun.py @@ -41,6 +41,9 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("input", nargs="+", type=str, help="Input mat file") parser.add_argument("-o", "--output", default=None, help="Output Path") + parser.add_argument( + "-f", "--frequency", default=5000, type=int, help="Sampling frequency" + ) parser.add_argument( "-i", "--iteration", default=5000, type=int, help="Number of iterations" ) @@ -56,6 +59,21 @@ def main(): action="store_true", help="Disable one cycle scheduler", ) + parser.add_argument( + "--ecg", default="ECG", type=str, help="Variable name for ECG (input)" + ) + parser.add_argument( + "--bce", + default="EEG_before_bcg", + type=str, + help="Variable name for BCG corropted EEG (input)", + ) + parser.add_argument( + "--eeg", + default="EEG_clean", + type=str, + help="Variable name for clean EEG (output)", + ) args = parser.parse_args() print(f"Settings: {args}") @@ -99,8 +117,8 @@ def main(): t = time.time() mat = h5py.File(f, "r") - ECG = np.array(mat["ECG"]).flatten() - EEG = np.array(mat["EEG_before_bcg"]).T + ECG = np.array(mat[args.ecg]).flatten() + EEG = np.array(mat[args.bce]).T # (input_eeg, input_ecg, sfreq=5000, iter_num=5000, winsize_sec=2, lr=1e-3, onecycle=True) EEG_unet = bcgunet.run( @@ -110,9 +128,10 @@ def main(): winsize_sec=args.window_size, lr=args.learning_rate, onecycle=not args.no_one_cycle, + sfreq=args.frequency, ) result = dict() - result["EEG_clean"] = EEG_unet + result[args.eeg] = EEG_unet savemat(ff_output, result, do_compression=True) diff --git a/bcgrun_web.py b/bcgrun_web.py index fec37f4..cde287d 100644 --- a/bcgrun_web.py +++ b/bcgrun_web.py @@ -7,14 +7,26 @@ 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 -) -> list[str]: + files: list[bytes], + sfreq: int, + 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) @@ -28,8 +40,8 @@ def run( output = os.path.join(task, str(i) + "_clean.mat") mat = h5py.File(input, "r") - ECG = np.array(mat["ECG"]).flatten() - EEG = np.array(mat["EEG_before_bcg"]).T + ECG = np.array(mat[ecg]).flatten() + EEG = np.array(mat[bce]).T EEG_unet = bcgunet.run( EEG, @@ -38,20 +50,41 @@ def run( winsize_sec=winsec, lr=lr, onecycle=onecycle, + sfreq=sfreq, ) result = dict() - result["EEG_clean"] = EEG_unet + result[eeg] = EEG_unet savemat(output, result, do_compression=True) outputs.append(output) - return outputs + if i == 0: + double_sfreq = sfreq * 2 + plt.figure(figsize=(12, 6), dpi=300) + plt.plot(EEG[19, :double_sfreq], "b.-", label="Orig EEG") + plt.plot(EEG_unet[19, :double_sfreq], "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 + +matlab_file_spec = """ +ECG Data: The raw ECG data collected inside the MRI scanner. +This array should contain the raw electrocardiogram (ECG) data in shape of (n_points, 1). +EEG Data: The raw EEG data collected inside the MRI scanner. +Each row corresponds to an EEG channel, and each column corresponds to a time point. We expect 31 EEG channels. The shape of the array should be (n_points, 31). + +P.S. n_points = sampling frequency * recording time (in seconds) +""" def main(): app = gr.Interface( title="BCG Unet", - description="BCGunet: Suppressing BCG artifacts on EEG collected inside an MRI scanner", + description="BCGunet: Suppressing BCG artifacts on EEG collected inside an MRI scanner\n\n" + matlab_file_spec, fn=run, inputs=[ gr.File( @@ -59,12 +92,31 @@ def main(): type="binary", file_types=["mat"], file_count=["multiple", "directory"], + info="Upload MATLAB .mat files. You can specify in-file variable names in the options below.", + ), + gr.Slider( + label="Sampling Frequency", + minimum=100, + maximum=10000, + step=100, + value=5000, + info="The sampling frequency of the EEG and ECG data. (Hz)", ), gr.Slider( - label="Learning Rate", minimum=1e-5, maximum=1e-1, step=1e-5, value=1e-3 + label="Learning Rate", + minimum=1e-5, + maximum=1e-1, + step=1e-5, + value=1e-3, + info="Adjusts the step size at each iteration of model training.", ), gr.Slider( - label="Window Size (seconds)", minimum=1, maximum=10, step=1, value=2 + label="Window Size (seconds)", + minimum=1, + maximum=10, + step=1, + value=2, + info="Sets the time window for analyzing the EEG signal.", ), gr.Slider( label="Number of Iterations", @@ -72,13 +124,36 @@ def main(): maximum=10000, step=1000, value=5000, + info="The total number of passes for training.", ), gr.Checkbox( label="One Cycle Scheduler", value=True, + info="Enable the learning rate policy that cycles the learning rate between two boundaries.", + ), + gr.Textbox( + label="Variable name for ECG (input)", + placeholder="Enter the variable name for ECG data", + value="ECG", + info="Specify the variable name in your .mat file that corresponds to the ECG data.", ), + gr.Textbox( + label="Variable name for BCG-corrupted EEG (input)", + placeholder="Enter the variable name for BCG-corrupted EEG data", + value="EEG_before_bcg", + info="Specify the variable name in your .mat file that corresponds to the BCG-corrupted EEG data.", + ), + gr.Textbox( + label="Variable name for clean EEG (output)", + placeholder="Enter the variable name for clean EEG data", + value="EEG_clean", + info="Specify the variable name in your .mat file that corresponds to the clean EEG data.", + ), + ], + outputs=[ + gr.File(label="Output File", file_count="multiple"), + gr.Image(label="Output Image", type="filepath"), ], - outputs=[gr.File(label="Output File", file_count="multiple")], allow_flagging="never", )