-
Notifications
You must be signed in to change notification settings - Fork 5
/
Copy pathneurofeedback.py
140 lines (118 loc) · 6.26 KB
/
neurofeedback.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import pandas as pd
from matplotlib import style
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import matplotlib
import time
import sys
import brainflow
import numpy as np
from brainflow.board_shim import BoardShim, BrainFlowInputParams, LogLevels, BoardIds
from brainflow.data_filter import DataFilter, FilterTypes, AggOperations
from brainflow.ml_model import MLModel, BrainFlowMetrics, BrainFlowClassifiers, BrainFlowModelParams
def main(i):
BoardShim.enable_dev_board_logger()
BoardShim.disable_board_logger() #optional. take this out for initial setup for your board.
# use synthetic board for demo
params = BrainFlowInputParams()
board_id = BoardIds.SYNTHETIC_BOARD.value
board = BoardShim(board_id, params)
eeg_channels = BoardShim.get_eeg_channels(board_id)
sampling_rate = BoardShim.get_sampling_rate(board_id)
timestamp = BoardShim.get_timestamp_channel(board_id)
board.prepare_session()
board.start_stream()
style.use('fivethirtyeight')
plt.title("Live EEG stream from Brainflow", fontsize=15)
plt.ylabel("Data in millivolts", fontsize=15)
plt.xlabel("\nTime", fontsize=10)
keep_alive = True
eeg1 = [] #lists to store eeg data
eeg2 = []
eeg3 = []
eeg4 = []
timex = [] #list to store timestamp
while keep_alive == True:
while board.get_board_data_count() < 250: #ensures that all data shape is the same
time.sleep(0.005)
data = board.get_current_board_data(250)
# creating a dataframe of the eeg data to extract eeg values later
eegdf = pd.DataFrame(np.transpose(data[eeg_channels]))
eegdf_col_names = ["ch1", "ch2", "ch3", "ch4", "ch5", "ch6", "ch7",
"ch8", "ch9", "ch10", "ch11", "ch12", "ch13", "ch14", "ch15", "ch16"]
eegdf.columns = eegdf_col_names
# to keep it simple, making another dataframe for the timestamps to access later
timedf = pd.DataFrame(np.transpose(data[timestamp]))
print("EEG Dataframe") #easy way to check what data is being streamed and if program is working
print(eegdf) #isn't neccesary.
for count, channel in enumerate(eeg_channels):
# filters work in-place
# Check Brainflow docs for more filters
if count == 0:
DataFilter.perform_bandstop(data[channel], sampling_rate, 60.0, 4.0, 4,
FilterTypes.BUTTERWORTH.value, 0) # bandstop 58 - 62
DataFilter.perform_bandpass(data[channel], sampling_rate, 21.0, 20.0, 4,
FilterTypes.BESSEL.value, 0) # bandpass 11 - 31
if count == 1:
DataFilter.perform_bandstop(data[channel], sampling_rate, 60.0, 4.0, 4,
FilterTypes.BUTTERWORTH.value, 0) # bandstop 58 - 62
DataFilter.perform_bandpass(data[channel], sampling_rate, 21.0, 20.0, 4,
FilterTypes.BESSEL.value, 0) # bandpass 11 - 31
if count == 2:
DataFilter.perform_bandstop(data[channel], sampling_rate, 60.0, 4.0, 4,
FilterTypes.BUTTERWORTH.value, 0) # bandstop 58 - 62
DataFilter.perform_bandpass(data[channel], sampling_rate, 21.0, 20.0, 4,
FilterTypes.BESSEL.value, 0) # bandpass 11 - 31
if count == 3:
DataFilter.perform_bandstop(data[channel], sampling_rate, 60.0, 4.0, 4,
FilterTypes.BUTTERWORTH.value, 0) # bandstop 58 - 62
DataFilter.perform_bandpass(data[channel], sampling_rate, 21.0, 20.0, 4,
FilterTypes.BESSEL.value, 0) # bandpass 11 - 31
# Brainflow ML Model
bands = DataFilter.get_avg_band_powers(
data, eeg_channels, sampling_rate, True)
feature_vector = np.concatenate((bands[0], bands[1]))
# calc concentration
concentration_params = BrainFlowModelParams(
BrainFlowMetrics.CONCENTRATION.value, BrainFlowClassifiers.KNN.value)
concentration = MLModel(concentration_params)
concentration.prepare()
print('Concentration: %f' % concentration.predict(feature_vector))
concentrated_measure = concentration.predict(feature_vector)
concentration.release()
# calc relaxation
relaxation_params = BrainFlowModelParams(
BrainFlowMetrics.RELAXATION.value, BrainFlowClassifiers.KNN.value)
relaxation = MLModel(relaxation_params)
relaxation.prepare()
print('Relaxation: %f' % relaxation.predict(feature_vector))
relaxed_measure = relaxation.predict(feature_vector)
relaxation.release()
#appending eeg data to lists
eeg1.extend(eegdf.iloc[:, 0].values) # I am using OpenBCI Ganglion board, so I only have four channels.
eeg2.extend(eegdf.iloc[:, 1].values) # If you have a different board, you should be able to copy paste
eeg3.extend(eegdf.iloc[:, 2].values) # these commands for more channels.
eeg4.extend(eegdf.iloc[:, 3].values)
timex.extend(timedf.iloc[:, 0].values) # timestamps
plt.cla()
#plotting eeg data
plt.plot(timex, eeg1, label="Channel 1", color="red")
plt.plot(timex, eeg2, label="Channel 2", color="blue")
plt.plot(timex, eeg3, label="Channel 3", color="orange")
plt.plot(timex, eeg4, label="Channel 4", color="purple")
plt.tight_layout()
keep_alive = False #resetting stream so that matplotlib can plot data
if concentrated_measure >= 0.5:
print("GOOD KEEP CONCENTRATING") #a program screaming at you to concentrate should do the trick :)
else:
print("WHERE IS THE CONCENTRATION??")
if relaxed_measure >= 0.5:
print("YES RELAX MORE")
else:
print("NO, START RELAXING")
board.stop_stream()
board.release_session()
ani = FuncAnimation(plt.gcf(), main, interval=1000) #this essentially calls the function several times until keyboard interrupt
plt.tight_layout()
plt.autoscale(enable=True, axis="y", tight=True)
plt.show()