Add UI descriptions and new parameter for sampling frequency

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)
 

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,47 +50,110 @@ 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(
                 label="Input Files (.mat)",
                 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",
                 minimum=1000,
                 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",
     )