Skip to content

Commit

Permalink
fix: fixed import (brainflow updated API)
Browse files Browse the repository at this point in the history
  • Loading branch information
ErikBjare committed Jun 16, 2022
1 parent 5bdfbb1 commit f208e08
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 35 deletions.
57 changes: 29 additions & 28 deletions eegnb/devices/eeg.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,17 @@
import numpy as np
import pandas as pd

from brainflow import BoardShim, BoardIds, BrainFlowInputParams
from brainflow.board_shim import BoardShim, BoardIds, BrainFlowInputParams
from muselsl import stream, list_muses, record, constants as mlsl_cnsts
from pylsl import StreamInfo, StreamOutlet, StreamInlet, resolve_byprop

from eegnb.devices.utils import get_openbci_usb, create_stim_array,SAMPLE_FREQS,EEG_INDICES,EEG_CHANNELS
from eegnb.devices.utils import (
get_openbci_usb,
create_stim_array,
SAMPLE_FREQS,
EEG_INDICES,
EEG_CHANNELS,
)


logger = logging.getLogger(__name__)
Expand All @@ -39,18 +45,19 @@
"notion2",
"freeeeg32",
"crown",
"museS_bfn", # bfn = brainflow with native bluetooth;
"museS_bfb", # bfb = brainflow with BLED dongle bluetooth
"museS_bfn", # bfn = brainflow with native bluetooth;
"museS_bfb", # bfb = brainflow with BLED dongle bluetooth
"muse2_bfn",
"muse2_bfb",
"muse2016_bfn",
"muse2016_bfb"
"muse2016_bfb",
]


class EEG:
device_name: str
stream_started: bool = False

def __init__(
self,
device=None,
Expand Down Expand Up @@ -85,8 +92,8 @@ def initialize_backend(self):
self.timestamp_channel = BoardShim.get_timestamp_channel(self.brainflow_id)
elif self.backend == "muselsl":
self._init_muselsl()
self._muse_get_recent() # run this at initialization to get some
# stream metadata into the eeg class
self._muse_get_recent() # run this at initialization to get some
# stream metadata into the eeg class

def _get_backend(self, device_name):
if device_name in brainflow_devices:
Expand Down Expand Up @@ -126,7 +133,7 @@ def _start_muse(self, duration):
print("will save to file: %s" % self.save_fn)
self.recording = Process(target=record, args=(duration, self.save_fn))
self.recording.start()

time.sleep(5)
self.stream_started = True
self.push_sample([99], timestamp=time.time())
Expand All @@ -137,7 +144,7 @@ def _stop_muse(self):
def _muse_push_sample(self, marker, timestamp):
self.muse_StreamOutlet.push_sample(marker, timestamp)

def _muse_get_recent(self, n_samples: int=256, restart_inlet: bool=False):
def _muse_get_recent(self, n_samples: int = 256, restart_inlet: bool = False):
if self._muse_recent_inlet and not restart_inlet:
inlet = self._muse_recent_inlet
else:
Expand All @@ -157,9 +164,8 @@ def _muse_get_recent(self, n_samples: int=256, restart_inlet: bool=False):
self.info = info
self.n_chans = n_chans

timeout = (n_samples/sfreq)+0.5
samples, timestamps = inlet.pull_chunk(timeout=timeout,
max_samples=n_samples)
timeout = (n_samples / sfreq) + 0.5
samples, timestamps = inlet.pull_chunk(timeout=timeout, max_samples=n_samples)

samples = np.array(samples)
timestamps = np.array(timestamps)
Expand Down Expand Up @@ -260,13 +266,13 @@ def _init_brainflow(self):

elif self.device_name == "muse2_bfn":
self.brainflow_id = BoardIds.MUSE_2_BOARD.value

elif self.device_name == "muse2_bfb":
self.brainflow_id = BoardIds.MUSE_2_BLED_BOARD.value

elif self.device_name == "muse2016_bfn":
self.brainflow_id = BoardIds.MUSE_2016_BOARD.value

elif self.device_name == "muse2016_bfb":
self.brainflow_id = BoardIds.MUSE_2016_BLED_BOARD.value

Expand All @@ -291,11 +297,13 @@ def _start_brainflow(self):
# only start stream if non exists
if not self.stream_started:
self.board.start_stream()

self.stream_started = True

# wait for signal to settle
if (self.device_name.find("cyton") != -1) or (self.device_name.find("ganglion") != -1):
if (self.device_name.find("cyton") != -1) or (
self.device_name.find("ganglion") != -1
):
# wait longer for openbci cyton / ganglion
sleep(10)
else:
Expand All @@ -314,21 +322,19 @@ def _stop_brainflow(self):

# Create a column for the stimuli to append to the EEG data
stim_array = create_stim_array(timestamps, self.markers)
timestamps = timestamps[ ..., None ]
timestamps = timestamps[..., None]

# Add an additional dimension so that shapes match
total_data = np.append(timestamps, eeg_data, 1)

# Append the stim array to data.
total_data = np.append(total_data, stim_array, 1)
total_data = np.append(total_data, stim_array, 1)

# Subtract five seconds of settling time from beginning
total_data = total_data[5 * self.sfreq :]
data_df = pd.DataFrame(total_data, columns=["timestamps"] + ch_names + ["stim"])
data_df.to_csv(self.save_fn, index=False)



def _brainflow_extract(self, data):
"""
Formats the data returned from brainflow to get
Expand Down Expand Up @@ -357,14 +363,12 @@ def _brainflow_extract(self, data):
eeg_data = data[:, BoardShim.get_eeg_channels(self.brainflow_id)]
timestamps = data[:, BoardShim.get_timestamp_channel(self.brainflow_id)]

return ch_names,eeg_data,timestamps

return ch_names, eeg_data, timestamps

def _brainflow_push_sample(self, marker):
last_timestamp = self.board.get_current_board_data(1)[self.timestamp_channel][0]
self.markers.append([marker, last_timestamp])


def _brainflow_get_recent(self, n_samples=256):

# initialize brainflow if not set
Expand Down Expand Up @@ -405,7 +409,6 @@ def start(self, fn, duration=None):
elif self.backend == "muselsl":
self._start_muse(duration)


def push_sample(self, marker, timestamp):
"""
Universal method for pushing a marker and its timestamp to store alongside the EEG data.
Expand Down Expand Up @@ -445,6 +448,4 @@ def get_recent(self, n_samples: int = 256):
sorted_cols = sorted(df.columns)
df = df[sorted_cols]


return df

14 changes: 7 additions & 7 deletions eegnb/devices/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,14 @@
import platform
import serial

from brainflow import BoardShim, BoardIds
from brainflow.board_shim import BoardShim, BoardIds


# Default channel names for the various EEG devices.
EEG_CHANNELS = {
"muse2016": ['TP9', 'AF7', 'AF8', 'TP10', 'Right AUX'],
"muse2": ['TP9', 'AF7', 'AF8', 'TP10', 'Right AUX'],
"museS": ['TP9', 'AF7', 'AF8', 'TP10', 'Right AUX'],
"muse2016": ["TP9", "AF7", "AF8", "TP10", "Right AUX"],
"muse2": ["TP9", "AF7", "AF8", "TP10", "Right AUX"],
"museS": ["TP9", "AF7", "AF8", "TP10", "Right AUX"],
"muse2016_bfn": BoardShim.get_eeg_names(BoardIds.MUSE_2016_BOARD.value),
"muse2016_bfb": BoardShim.get_eeg_names(BoardIds.MUSE_2016_BLED_BOARD.value),
"muse2_bfn": BoardShim.get_eeg_names(BoardIds.MUSE_2_BOARD.value),
Expand All @@ -26,7 +26,7 @@
"notion1": BoardShim.get_eeg_names(BoardIds.NOTION_1_BOARD.value),
"notion2": BoardShim.get_eeg_names(BoardIds.NOTION_2_BOARD.value),
"crown": BoardShim.get_eeg_names(BoardIds.CROWN_BOARD.value),
"freeeeg32": [f'eeg_{i}' for i in range(0,32)],
"freeeeg32": [f"eeg_{i}" for i in range(0, 32)],
}

BRAINFLOW_CHANNELS = {
Expand Down Expand Up @@ -62,9 +62,9 @@
"muse2016": 256,
"muse2": 256,
"museS": 256,
"muse2016_bfn": BoardShim.get_sampling_rate(BoardIds.MUSE_2016_BOARD.value),
"muse2016_bfn": BoardShim.get_sampling_rate(BoardIds.MUSE_2016_BOARD.value),
"muse2016_bfb": BoardShim.get_sampling_rate(BoardIds.MUSE_2016_BLED_BOARD.value),
"muse2_bfn": BoardShim.get_sampling_rate(BoardIds.MUSE_2_BOARD.value),
"muse2_bfn": BoardShim.get_sampling_rate(BoardIds.MUSE_2_BOARD.value),
"muse2_bfb": BoardShim.get_sampling_rate(BoardIds.MUSE_2_BLED_BOARD.value),
"museS_bfn": BoardShim.get_sampling_rate(BoardIds.MUSE_S_BOARD.value),
"museS_bfb": BoardShim.get_sampling_rate(BoardIds.MUSE_S_BLED_BOARD.value),
Expand Down

0 comments on commit f208e08

Please sign in to comment.