Update app.py

app.py

@@ -4,7 +4,6 @@ import numpy as np
 import pandas as pd
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from mne.time_frequency import psd_welch
 import torch
 
 # Load LLM
@@ -18,6 +17,8 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 def compute_band_power(psd, freqs, fmin, fmax):
+    # psd shape: (n_channels, n_freqs)
+    # freqs shape: (n_freqs,)
     freq_mask = (freqs >= fmin) & (freqs <= fmax)
     band_psd = psd[:, freq_mask].mean()
     return float(band_psd)
@@ -35,19 +36,16 @@ def inspect_file(file):
     file_ext = file_ext.lower()
 
     if file_ext == ".fif":
-        # FIF files: MNE compatible, no columns needed
         return (
             "FIF file detected. No need for time column selection. The file's sampling frequency will be used.",
             [],
             "FIF file doesn't require column inspection."
         )
     elif file_ext == ".csv":
-        # Read a small portion of the CSV to determine columns
         try:
             df = pd.read_csv(file_path, nrows=5)
         except Exception as e:
             return f"Error reading CSV: {e}", [], "Could not read CSV preview."
-
         cols = list(df.columns)
         preview = df.head().to_markdown()
         return (
@@ -62,7 +60,7 @@ def load_eeg_data(file_path, default_sfreq=256.0, time_col='time'):
     """
     Load EEG data with flexibility.
     If FIF: Use MNE's read_raw_fif directly.
-    If CSV: 
+    If CSV:
       - If time_col is given and present in the file, use it.
       - Otherwise, assume default_sfreq.
     """
@@ -79,19 +77,14 @@ def load_eeg_data(file_path, default_sfreq=256.0, time_col='time'):
             # Use the selected time column to compute sfreq
             time = df[time_col].values
             data_df = df.drop(columns=[time_col])
-
-            # Drop non-numeric columns
             for col in data_df.columns:
                 if not pd.api.types.is_numeric_dtype(data_df[col]):
                     data_df = data_df.drop(columns=[col])
 
             if len(time) < 2:
-                # Not enough time points to compute sfreq, fallback
                 sfreq = default_sfreq
             else:
-                # Compute sfreq from time
                 dt = np.mean(np.diff(time))
-                # Ensure dt is positive
                 if dt <= 0:
                     sfreq = default_sfreq
                 else:
@@ -104,13 +97,11 @@ def load_eeg_data(file_path, default_sfreq=256.0, time_col='time'):
             data_df = df
             sfreq = default_sfreq
 
-        # Ensure sfreq is positive
         if sfreq <= 0:
-            sfreq = 256.0  # fallback if something odd happens
+            sfreq = 256.0
 
         ch_names = list(data_df.columns)
-        data = data_df.values.T  # shape: (n_channels, n_samples)
-
+        data = data_df.values.T  # (n_channels, n_samples)
         ch_types = ['eeg'] * len(ch_names)
         info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
         raw = mne.io.RawArray(data, info)
@@ -130,8 +121,9 @@ def analyze_eeg(file, default_sfreq, time_col):
 
     raw = load_eeg_data(file.name, default_sfreq=fs, time_col=time_col)
 
-    # Use the directly imported psd_welch function
-    psd, freqs = psd_welch(raw, fmin=1, fmax=40)
+    # Use raw.compute_psd instead of psd_welch
+    psd_obj = raw.compute_psd(fmin=1, fmax=40, method='welch')
+    psd, freqs = psd_obj.get_data(return_freqs=True)
     alpha_power = compute_band_power(psd, freqs, 8, 12)
     beta_power = compute_band_power(psd, freqs, 13, 30)
 
@@ -154,11 +146,8 @@ Provide a concise, user-friendly interpretation of these findings in simple term
 def preview_file(file):
     msg, cols, preview = inspect_file(file)
     # Always include (No time column) as the first choice
-    # If no columns were found, we still have (No time column) as an option
     cols = ["(No time column)"] + cols
     default_value = "(No time column)"
-
-    # Return an update dict for the dropdown
     return msg, gr.update(choices=cols, value=default_value), preview
 
 with gr.Blocks() as demo:
@@ -173,11 +162,14 @@ with gr.Blocks() as demo:
     file_input = gr.File(label="Upload your EEG data (FIF or CSV)")
     preview_button = gr.Button("Inspect File")
     msg_output = gr.Markdown()
-    # Allow custom values in case something goes off
     cols_dropdown = gr.Dropdown(label="Select Time Column (optional)", allow_custom_value=True, interactive=True)
     preview_output = gr.Markdown()
 
-    preview_button.click(preview_file, inputs=[file_input], outputs=[msg_output, cols_dropdown, preview_output])
+    preview_button.click(
+        preview_file, 
+        inputs=[file_input], 
+        outputs=[msg_output, cols_dropdown, preview_output]
+    )
 
     default_sfreq_input = gr.Textbox(label="Default Sampling Frequency (Hz) if no time column", value="100")
     analyze_button = gr.Button("Run Analysis")