Merge pull request #30 from PayalLakra/bio_amptool

Changes done in ffteeg code

Author: lorforlinux

ffteeg.py

@@ -1,13 +1,14 @@
 import numpy as np
-from numpy import hamming
+from collections import deque
 from PyQt5.QtWidgets import QApplication, QVBoxLayout, QHBoxLayout, QMainWindow, QWidget
 from PyQt5.QtCore import Qt
 from pyqtgraph import PlotWidget
 import pyqtgraph as pg
 import pylsl
 import sys
-from scipy.signal import butter, filtfilt, iirnotch
+from scipy.signal import butter, iirnotch, lfilter, lfilter_zi
 from scipy.fft import fft
+import math
 
 class EEGMonitor(QMainWindow):
     def __init__(self): 
@@ -25,8 +26,8 @@ def __init__(self):
         self.eeg_plot_widget.setBackground('w')
         self.eeg_plot_widget.showGrid(x=True, y=True)
         self.eeg_plot_widget.setLabel('bottom', 'EEG Plot')
-        self.eeg_plot_widget.setYRange(0, 15000, padding=0)
-        self.eeg_plot_widget.setXRange(0, 10, padding=0)
+        self.eeg_plot_widget.setYRange(-5000, 5000, padding=0)
+        self.eeg_plot_widget.setXRange(0, 2, padding=0)
         self.eeg_plot_widget.setMouseEnabled(x=False, y=True)  # Disable zoom
         self.main_layout.addWidget(self.eeg_plot_widget)
 
@@ -38,9 +39,10 @@ def __init__(self):
         self.fft_plot.setBackground('w')
         self.fft_plot.showGrid(x=True, y=True)
         self.fft_plot.setLabel('bottom', 'FFT')
-        self.fft_plot.setYRange(0, 25000, padding=0)
+        # self.fft_plot.setYRange(0, 25000, padding=0)
         self.fft_plot.setXRange(0, 50, padding=0)  # Set x-axis to 0 to 50 Hz
         self.fft_plot.setMouseEnabled(x=False, y=False)  # Disable zoom
+        self.fft_plot.setAutoVisible(y=True)  # Allow y-axis to autoscale
         self.bottom_layout.addWidget(self.fft_plot)
 
         # Bar graph for brainwave power bands (right side of the second half)
@@ -71,20 +73,20 @@ def __init__(self):
         self.sampling_rate = int(self.inlet.info().nominal_srate())
         print(f"Sampling rate: {self.sampling_rate} Hz")
 
-        # Data and buffers
-        self.buffer_size = self.sampling_rate * 10  # Fixed-size buffer for 10 seconds
-        self.eeg_data = np.zeros(self.buffer_size)  # Fixed-size array for circular buffer
-        self.time_data = np.linspace(0, 10, self.buffer_size)  # Fixed time array for plotting
-        self.current_index = 0  # Index for overwriting data
+        # Data and Buffers
+        self.one_second_buffer = deque(maxlen=self.sampling_rate)  # 1-second buffer
+        self.buffer_size = self.sampling_rate * 10
+        self.moving_window_size = self.sampling_rate * 2  # 2-second window
 
-        # Moving window for brainwave power
-        self.moving_window_size = self.sampling_rate * 3   # 3-second window
-        self.moving_window_buffer = np.zeros(self.moving_window_size)
+        self.eeg_data = np.zeros(self.buffer_size)
+        self.time_data = np.linspace(0, 10, self.buffer_size)
+        self.current_index = 0
 
-        # Filters
         self.b_notch, self.a_notch = iirnotch(50, 30, self.sampling_rate)
-        self.b_highpass, self.a_highpass = butter(4, 1.5 / (0.5 * self.sampling_rate), btype='high')
-        self.b_lowpass, self.a_lowpass = butter(4, 45 / (0.5 * self.sampling_rate), btype='low')
+        self.b_band, self.a_band = butter(4, [0.5 / (self.sampling_rate / 2), 48.0 / (self.sampling_rate / 2)], btype='band')
+
+        self.zi_notch = lfilter_zi(self.b_notch, self.a_notch) * 0
+        self.zi_band = lfilter_zi(self.b_band, self.a_band) * 0
 
         # Timer for updating the plot
         self.timer = pg.QtCore.QTimer()
@@ -95,68 +97,59 @@ def __init__(self):
         self.fft_curve = self.fft_plot.plot(pen=pg.mkPen('r', width=1))  # FFT Colour is red
 
     def update_plot(self):
-        samples, _ = self.inlet.pull_chunk(timeout=0.0, max_samples=30)
+        samples, _ = self.inlet.pull_chunk(timeout=0.0)
         if samples:
             for sample in samples:
-                # Overwrite the oldest data point in the buffer
-                self.eeg_data[self.current_index] = sample[0]
-                self.current_index = (self.current_index + 1) % self.buffer_size  # Circular increment
-
-            if self.current_index >= self.buffer_size:
-                plot_data = self.eeg_data
-            else:
-                plot_data = np.concatenate((self.eeg_data[self.current_index:], self.eeg_data[:self.current_index]))
-
-            # Apply filters to the full data for EEG plot
-            filtered_eeg = filtfilt(self.b_notch, self.a_notch, plot_data)
-            filtered_eeg = filtfilt(self.b_highpass, self.a_highpass, filtered_eeg)
-            filtered_eeg = filtfilt(self.b_lowpass, self.a_lowpass, filtered_eeg)
-
-            # Update the EEG plot with the filtered data
-            self.eeg_curve.setData(self.time_data, filtered_eeg)
-
-            # Perform FFT on the latest 1-second slice
-            latest_data = filtered_eeg[-self.sampling_rate:]
-            window = hamming(len(latest_data))
-            filtered_eeg_windowed = latest_data * window
-
-            # Apply zero-padding
-            zero_padded_length = 512
-            filtered_eeg_windowed_padded = np.pad(filtered_eeg_windowed, (0, zero_padded_length - len(filtered_eeg_windowed)), 'constant')
-
-            eeg_fft = np.abs(fft(filtered_eeg_windowed_padded))[:len(filtered_eeg_windowed_padded) // 2]
-            freqs = np.fft.fftfreq(len(filtered_eeg_windowed_padded), 1 / self.sampling_rate)[:len(filtered_eeg_windowed_padded) // 2]
-
-            # Update FFT plot
-            self.fft_curve.setData(freqs, eeg_fft)
-
-            # Update the 3-second moving window buffer
-            for sample in latest_data:
-                self.moving_window_buffer = np.roll(self.moving_window_buffer, -1)
-                self.moving_window_buffer[-1] = sample
-
-            # Apply filters to the moving window buffer
-            filtered_window = filtfilt(self.b_notch, self.a_notch, self.moving_window_buffer)
-            filtered_window = filtfilt(self.b_highpass, self.a_highpass, filtered_window)
-            filtered_window = filtfilt(self.b_lowpass, self.a_lowpass, filtered_window)
-
-            # Perform FFT on the moving window buffer
-            windowed_data = filtered_window * hamming(len(filtered_window))
-            fft_data = np.abs(fft(windowed_data))[:len(windowed_data) // 2]
-            window_freqs = np.fft.fftfreq(len(windowed_data), 1 / self.sampling_rate)[:len(windowed_data) // 2]
-
-            brainwave_power = self.calculate_brainwave_power(fft_data, window_freqs)
-            self.brainwave_bars.setOpts(height=brainwave_power)
+                raw_point = sample[0]
+
+                notch_filtered, self.zi_notch = lfilter(self.b_notch, self.a_notch, [raw_point], zi=self.zi_notch)
+                band_filtered, self.zi_band = lfilter(self.b_band, self.a_band, notch_filtered, zi=self.zi_band)
+                band_filtered = band_filtered[-1]  # Get the current filtered point
+
+                # Update the EEG plot
+                self.eeg_data[self.current_index] = band_filtered
+                self.current_index = (self.current_index + 1) % self.buffer_size
+
+                if self.current_index == 0:
+                    plot_data = self.eeg_data
+                else:
+                    plot_data = np.concatenate((self.eeg_data[self.current_index:], self.eeg_data[:self.current_index]))
+
+                recent_data = plot_data[-self.moving_window_size:]
+                recent_time = np.linspace(0, len(recent_data) / self.sampling_rate, len(recent_data))
+                self.eeg_curve.setData(recent_time, recent_data)
+
+                self.one_second_buffer.append(band_filtered)           # Add the filtered point to the 1-second buffer
+                if len(self.one_second_buffer) == self.sampling_rate:  # Process FFT and brainwave power
+                    self.process_fft_and_brainpower()
+                    self.one_second_buffer.clear() 
+                    
+    def process_fft_and_brainpower(self):
+        window = np.hanning(len(self.one_second_buffer))       # Apply Hanning window to the buffer
+        buffer_windowed = np.array(self.one_second_buffer) * window
+
+        # Perform FFT
+        fft_result = np.abs(fft(buffer_windowed))[:len(buffer_windowed) // 2]
+        fft_result /= len(buffer_windowed)
+        freqs = np.fft.fftfreq(len(buffer_windowed), 1 / self.sampling_rate)[:len(buffer_windowed) // 2]
+        self.fft_curve.setData(freqs, fft_result)
+
+        brainwave_power = self.calculate_brainwave_power(fft_result, freqs)
+        self.brainwave_bars.setOpts(height=brainwave_power)
 
     def calculate_brainwave_power(self, fft_data, freqs):
-        delta_power = np.sum(fft_data[(freqs >= 0.5) & (freqs <= 4)])
-        theta_power = np.sum(fft_data[(freqs >= 4) & (freqs <= 8)])
-        alpha_power = np.sum(fft_data[(freqs >= 8) & (freqs <= 13)])
-        beta_power = np.sum(fft_data[(freqs >= 13) & (freqs <= 30)])
-        gamma_power = np.sum(fft_data[(freqs >= 30) & (freqs <= 45)])
-
+        delta_power = math.sqrt(np.sum(((fft_data[(freqs >= 0.5) & (freqs <= 4)])**2)/4))
+        theta_power = math.sqrt(np.sum(((fft_data[(freqs >= 4) & (freqs <= 8)])**2)/5))
+        alpha_power = math.sqrt(np.sum(((fft_data[(freqs >= 8) & (freqs <= 13)])**2)/6))
+        beta_power = math.sqrt(np.sum(((fft_data[(freqs >= 13) & (freqs <=30)])**2)/18))
+        gamma_power = math.sqrt(np.sum(((fft_data[(freqs >= 30) & (freqs <= 45)])**2)/16))
+        print("Delta Power", delta_power)
+        print("Theta Power", theta_power)
+        print("Alpha Power", alpha_power)
+        print("Beta Power", beta_power)
+        print("Gamma Power", gamma_power)
         return [delta_power, theta_power, alpha_power, beta_power, gamma_power]
-
+ 
 if __name__ == "__main__":
     app = QApplication(sys.argv)
     window = EEGMonitor()  

test/ball.py

@@ -0,0 +1,100 @@
+import pygame
+import pylsl
+import numpy as np
+import time
+from pylsl import StreamInlet, resolve_stream
+
+pygame.init()
+
+# Screen setup
+screen = pygame.display.set_mode((800, 600))
+pygame.display.set_caption('Focus Game')
+
+# Ball properties
+ball_x, ball_y = 400, 550
+ball_radius = 20
+focus_speed_upward = 8     # Speed when moving upward
+focus_speed_downward = 4   # Speed when moving downward
+focus_timeout = 2          # Time in seconds to stabilize focus
+focus_threshold = 1000000000  # Threshold for beta power
+
+# LSL stream setup
+streams = resolve_stream('name', 'BioAmpDataStream')
+if not streams:
+    print("No LSL stream found!")
+    pygame.quit()
+    exit()
+
+inlet = StreamInlet(streams[0])
+print(" LSL Stream Started")
+
+# Buffer for EEG data and Focus tracking variables
+buffer = []
+buffer_size = 500
+focus_timer = 0
+last_focus_time = time.time()
+last_time = time.time()
+
+def calculate_focus_level(eeg_data):
+    fft_result = np.fft.fft(eeg_data)
+    freqs = np.fft.fftfreq(len(eeg_data), 1 / 500)
+    positive_freqs = freqs[:len(freqs) // 2]
+    fft_magnitude = np.abs(fft_result[:len(freqs) // 2])
+
+    # Extract beta band power (13-30 Hz)
+    beta_band = np.where((positive_freqs >= 13) & (positive_freqs <= 30))
+    beta_power = np.sum(fft_magnitude[beta_band] ** 2)
+    print(f"Beta Power: {beta_power}")
+
+    return beta_power
+
+def update_ball_position(focus_level, is_focus_stable):
+    global ball_y
+    if is_focus_stable:
+        ball_y = max(0, ball_y - focus_speed_upward)                    # Move upward, bounded by top
+    else:
+        ball_y = min(600 - ball_radius, ball_y + focus_speed_downward)  # Move downward, bounded by bottom
+
+# Game loop
+running = True
+while running:
+    try:
+        for event in pygame.event.get():
+            if event.type == pygame.QUIT:
+                running = False
+
+        # Read EEG data from the LSL stream
+        sample, _ = inlet.pull_sample(timeout=0.1)
+        if sample:
+            buffer.append(sample[:1])               # Append new data to buffer
+
+        current_time = time.time()
+        if current_time - last_time >= 1:
+            last_time = current_time
+            release_count = int(len(buffer) * 0.2)  # Remove oldest 20%
+            buffer = buffer[release_count:]         # Trim buffer
+
+        if len(buffer) >= buffer_size:              # Process EEG data when the buffer is full
+            eeg_data = np.array(buffer)[:, 0]       # Use only the first channel
+            buffer = []                             # Clear the buffer
+
+            focus_level = calculate_focus_level(eeg_data)
+
+            if focus_level > focus_threshold:
+                focus_timer = min(focus_timeout, focus_timer + (current_time - last_focus_time))
+            else:
+                focus_timer = max(0, focus_timer - (current_time - last_focus_time))
+
+            is_focus_stable = focus_timer >= focus_timeout
+            update_ball_position(focus_level, is_focus_stable)
+            last_focus_time = current_time
+
+        screen.fill((0, 0, 0))
+        pygame.draw.circle(screen, (255, 0, 0), (ball_x, ball_y), ball_radius)  # Draw the ball
+        pygame.display.update()
+
+    except KeyboardInterrupt:
+        print("Exiting gracefully...")
+        running = False
+
+pygame.quit()