vlawhern / arl-eegmodels Goto Github PK

path='/home/user'
import logging
import os.path
import time
from collections import OrderedDict
import sys

import numpy as np


from braindecode.datasets.bcic_iv_2a import BCICompetition4Set2A
from braindecode.experiments.monitors import LossMonitor, MisclassMonitor, \
    RuntimeMonitor
from braindecode.experiments.stopcriteria import MaxEpochs, NoDecrease, Or
from braindecode.datautil.iterators import BalancedBatchSizeIterator
from braindecode.datautil.splitters import split_into_two_sets
from braindecode.mne_ext.signalproc import mne_apply,resample_cnt
from braindecode.datautil.signalproc import (bandpass_cnt,
                                             exponential_running_standardize)
from braindecode.datautil.trial_segment import create_signal_target_from_raw_mne
X_train_set=list()
y_train_set=list()
X_test_set=list()
y_test_set=list()

for subject_id in list(range(1,10)):
    data_folder= path
    low_cut_hz=4
    ival = [500, 2500]
    high_cut_hz = 40
    factor_new = 1e-3
    init_block_size = 1000

    train_filename = 'A{:02d}T.gdf'.format(subject_id)
    test_filename = 'A{:02d}E.gdf'.format(subject_id)
    train_filepath = os.path.join(data_folder, train_filename)
    test_filepath = os.path.join(data_folder, test_filename)
    train_label_filepath = train_filepath.replace('.gdf', '.mat')
    test_label_filepath = test_filepath.replace('.gdf', '.mat')

    train_loader = BCICompetition4Set2A(
        train_filepath, labels_filename=train_label_filepath)
    test_loader = BCICompetition4Set2A(
        test_filepath, labels_filename=test_label_filepath)
    train_cnt = train_loader.load()
    test_cnt = test_loader.load()

    # Preprocessing

    train_cnt = train_cnt.drop_channels(['STI 014', 'EOG-left',
                                         'EOG-central', 'EOG-right'])
    assert len(train_cnt.ch_names) == 22
    # lets convert to millvolt for numerical stability of next operations
    train_cnt = resample_cnt(train_cnt, 128)### added by me as suggeste in the paper
    train_cnt = mne_apply(lambda a: a * 1e6, train_cnt)
    train_cnt = mne_apply(
        lambda a: bandpass_cnt(a, low_cut_hz, high_cut_hz, train_cnt.info['sfreq'],
                               filt_order=3,
                               axis=1), train_cnt)
    train_cnt = mne_apply(
        lambda a: exponential_running_standardize(a.T, factor_new=factor_new,
                                                  init_block_size=init_block_size,
                                                  eps=1e-4).T,
        train_cnt)

    test_cnt = test_cnt.drop_channels(['STI 014', 'EOG-left',
                                       'EOG-central', 'EOG-right'])
    assert len(test_cnt.ch_names) == 22
    test_cnt = resample_cnt(test_cnt, 128)### added by me as suggeste in the paper
    test_cnt = mne_apply(lambda a: a * 1e6, test_cnt)
    test_cnt = mne_apply(
        lambda a: bandpass_cnt(a, low_cut_hz, high_cut_hz, test_cnt.info['sfreq'],
                               filt_order=3,
                               axis=1), test_cnt)
    test_cnt = mne_apply(
        lambda a: exponential_running_standardize(a.T, factor_new=factor_new,
                                                  init_block_size=init_block_size,
                                                  eps=1e-4).T,
        test_cnt)

    marker_def = OrderedDict([('Left Hand', [1]), ('Right Hand', [2],),
                              ('Foot', [3]), ('Tongue', [4])])

    train_set = create_signal_target_from_raw_mne(train_cnt, marker_def, ival)
    test_set = create_signal_target_from_raw_mne(test_cnt, marker_def, ival)

    X_train_set.append(train_set.X)
    y_train_set.append(train_set.y)
    X_test_set.append(test_set.X)
    y_test_set.append(test_set.y)



subject='1' # the subject id as string
subject = int(subject)-1 # str to int
# shuffle part 
from random import shuffle
Lt = list(range(0, 9)) 
shuffle(Lt)
Lt.remove(subject)
Lv = Lt[0:3]
[Lt.remove(i) for i in Lv]
X_t = list()
X_v = list()
y_t = list()
y_v = list()
print(len(X_train_set))
for i in Lt:
    X_t.append(X_train_set[i])
    y_t.append(y_train_set[i])


for i in Lv:
    X_v.append(X_train_set[i])
    y_v.append(y_train_set[i])

X_train = np.concatenate(X_t, axis=0).astype('float32')
    
print(Lt)#list training set
print(Lv)#list testing set
print(len(X_t))
print(len(X_v))
print(X_train.shape)
print(y_train.shape)



X_train = np.concatenate(X_t, axis=0).astype('float32')
X_val = np.concatenate(X_v, axis=0).astype('float32')
X_test = X_test_set[subject].astype('float32')

y_train = np.concatenate(y_t, axis=0).astype('int16')
y_val = np.concatenate(y_v, axis=0).astype('int16')
y_test = y_test_set[subject].astype('int16')

y_train = y_train.reshape((288*5))
y_val = y_val.reshape((288*3))
y_test = y_test.reshape((288))
y_train = keras.utils.to_categorical(y_train, 4)
y_val = keras.utils.to_categorical(y_val, 4)
y_test = keras.utils.to_categorical(y_test, 4)

def getdata_A_T_22(filename,ids):
    raw_gdf = mne.io.read_raw_gdf(filename, stim_channel="auto", verbose='ERROR',
                                  exclude=(["EOG-left", "EOG-central", "EOG-right"]))
    #rename
    raw_gdf.rename_channels(
        {'EEG-Fz': 'Fz', 'EEG-0': 'FC3', 'EEG-1': 'FC1', 'EEG-2': 'FCz', 'EEG-3': 'FC2', 'EEG-4': 'FC4',
         'EEG-5': 'C5', 'EEG-C3': 'C3', 'EEG-6': 'C1', 'EEG-Cz': 'Cz', 'EEG-7': 'C2', 'EEG-C4': 'C4', 'EEG-8': 'C6',
         'EEG-9': 'CP3', 'EEG-10': 'CP1', 'EEG-11': 'CPz', 'EEG-12': 'CP2', 'EEG-13': 'CP4',
         'EEG-14': 'P1', 'EEG-15': 'Pz', 'EEG-16': 'P2', 'EEG-Pz': 'POz'})
    # Pre-load the data
    raw_gdf.load_data()
    #Bandpass filter 4-40Hz
    raw_gdf.filter(4.0, 40.0, fir_design="firwin", skip_by_annotation="edge")
    # Select data from 0.5s to 2.5s after the cue.  
    tmin, tmax =.5 , 2.5
    #  events_id
    events, events_id = mne.events_from_annotations(raw_gdf)
    event_id = dict({'769': 7, '770': 8, '771': 9, '772': 10})
    #get epoches
    epochs = mne.Epochs(raw_gdf, events, event_id, tmin, tmax, proj=True, baseline=None, preload=True)
    # print(epochs)
    labels = epochs.events[:, -1]
    data = epochs.get_data()
    # resample to 128Hz
    original_sampling_rate = 250
    target_sampling_rate = 128
    original_length = data.shape[-1]
    target_length = int(original_length * target_sampling_rate / original_sampling_rate)
    # resample    
    resampled_data = np.zeros((data.shape[0], data.shape[1], target_length))
    for i in range(data.shape[0]):
        for j in range(data.shape[1]):
            resampled_data[i, j, :] = resample(data[i, j, :], target_length)
    # print results
    print(f"orginal shape: {data.shape}, new shape: {resampled_data.shape}")
    data = resampled_data
    labels[labels == 7] = 0
    labels[labels == 8] = 1
    labels[labels == 9] = 2
    labels[labels == 10] = 3
    X = torch.from_numpy(data).unsqueeze(1)
    y = torch.from_numpy(labels).long()
    X, y = shuffle(X, y, random_state=42)  
    #Performing four-fold cross-validation
    num_folds = 4
    skf = StratifiedKFold(n_splits=num_folds, shuffle=True, random_state=42)
    folded_data = []
    for train_index, val_index in skf.split(X, y):
        X_train, X_val = X[train_index], X[val_index]
        y_train, y_val = y[train_index], y[val_index]
        #print(y_val)
        folded_data.append((X_train, y_train, X_val, y_val))
    # save
    for i, fold_data in enumerate(folded_data):
        X_train, y_train, X_val, y_val = fold_data
        print('shape',X_train.shape)
        print(f"Fold {i + 1}: Training Set Size = {len(X_train)}, Validation Set Size = {len(X_val)}")
        np.save(f'./2a_T_22s/train_data_a0{ids}_{i+1}.npy', X_train)
        np.save(f'./2a_T_22s/test_data_a0{ids}_{i+1}.npy', X_val)
        np.save(f'./2a_T_22s/train_labels_a0{ids}_{i+1}.npy', y_train)
        np.save(f'./2a_T_22s/test_labels_a0{ids}_{i+1}.npy', y_val)
for i in range(9):
    #The third subject's event differs
    if i==3:
        continue
    print(i+1)
    getdata_A_T_22(f"./BCICIV_2a_gdf/A0{i+1}T.gdf",i+1)

class EEGNet(nn.Module):
    def __init__(self,Chans, T,clas, F1, D, F2 ):
        super(EEGNet, self).__init__()
        self.drop_out = 0.5
        self.block_1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1,  
                out_channels=F1,  
                kernel_size=(1, 32),  
                bias=False,
                padding='same'
            ),  
            nn.BatchNorm2d(F1)  
        )
        self.block_2 = nn.Sequential(
            nn.Conv2d(
                in_channels=F1,  
                out_channels=D*F1,  
                kernel_size=(Chans, 1), 
                groups=F1,
                bias=False, 
                padding='valid'
            ),  
            nn.BatchNorm2d(D*F1),  
            nn.ELU(),
            nn.AvgPool2d((1, 4)),  
            nn.Dropout(self.drop_out)  
        )

        self.block_3 = nn.Sequential(
            nn.Conv2d(
                in_channels=D*F1,  
                out_channels=D*F1,  
                kernel_size=(1, 16),  
                groups=D*F1,
                bias=False,
                padding='same'
            ),  
            nn.Conv2d(
                in_channels=D*F1,  
                out_channels=F2,  
                kernel_size=(1, 1),  
                bias=False,
                padding='same'
            ),  
            nn.BatchNorm2d(F2), 
            nn.ELU(),
            nn.AvgPool2d((1, 8)),  
            nn.Dropout(self.drop_out)
        )

        self.out = nn.Linear((F2 * (T//32)), clas)

    def forward(self, x):
        x = self.block_1(x)
        x = self.block_2(x)
        x = self.block_3(x)

        x = x.view(x.size(0), -1)
        x = self.out(x)
        return x
device = torch.device("cuda")
creation=nn.CrossEntropyLoss()  
creation=creation.to(device)
learning_rate=3e-4
beta1=0.9
beta2=0.999
accs=np.zeros((9,4))
for sub in range(9):
    for fold in range(4):
        print(sub,fold)
        # load data
        loaded_train_data = np.load(f"./2a_T_22s/train_data_a0{sub+1}_{fold+1}.npy")
        #print(f"./2a_T_22s/train_data_a0{sub+1}_{fold+1}.npy")
        loaded_test_data = np.load(f'./2a_T_22s/test_data_a0{sub+1}_{fold+1}.npy')
        loaded_train_labels = np.load(f'./2a_T_22s/train_labels_a0{sub+1}_{fold+1}.npy')
        loaded_test_labels = np.load(f'./2a_T_22s/test_labels_a0{sub+1}_{fold+1}.npy')
        print(loaded_train_data.shape, loaded_train_labels.shape)
        #  PyTorch Tensor
        loaded_train_data_tensor = torch.Tensor(loaded_train_data)
        loaded_test_data_tensor = torch.Tensor(loaded_test_data)
        loaded_train_labels_tensor = torch.Tensor(loaded_train_labels).long()
        loaded_test_labels_tensor = torch.Tensor(loaded_test_labels).long()
        print(loaded_train_data_tensor.shape, loaded_train_labels_tensor.shape)
        # TensorDataset
        loaded_train_dataset = TensorDataset(loaded_train_data_tensor, loaded_train_labels_tensor)
        loaded_test_dataset = TensorDataset(loaded_test_data_tensor, loaded_test_labels_tensor)

        #  DataLoader
        batch_size = 32
        train_loader = DataLoader(loaded_train_dataset, batch_size=batch_size, shuffle=True)
        test_loader = DataLoader(loaded_test_dataset, batch_size=batch_size, shuffle=False)

        model = EEGNet(22,256,4,8, 2, 16)       
        optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)  
        model_save_path = './2a_T_22/models/'
        os.makedirs(model_save_path, exist_ok=True)
        EPOCHS = 300
        train_steps = 0
        model.cuda()
        test_steps = 0
        losses = []
        accuracies = []
        ftrain_label = torch.tensor(loaded_train_labels).to(device)

        for epoch in range(EPOCHS):  
            model.train() 
            epoch_train_loss = 0.0
            for train_batch, (train_image, train_label) in enumerate(train_loader):

                train_image, train_label = train_image.to(device), train_label.to(device)  
                train_predictions = model(train_image)

                batch_train_loss = creation(train_predictions, train_label)
                optimizer.zero_grad()  
                batch_train_loss.backward()
                optimizer.step()
                epoch_train_loss += batch_train_loss.item()
                train_steps += 1
                if train_steps % 500 == 0:
                    print("train{}，train loss{}".format(train_steps, batch_train_loss.item()))
                if epoch%100==0:
                        model.eval()
                        with torch.no_grad():
                            predictions=model(loaded_test_data_tensor.to(device)).cpu()
                            predictions.numpy()
                            test_acc=(predictions.argmax(dim=1)==torch.tensor(loaded_test_labels )).sum()
                            print("test acc",test_acc/len(loaded_test_labels))
            losses.append(epoch_train_loss)
            train_predictions = model(loaded_train_data_tensor.to(device))

            accuracy = (torch.argmax(train_predictions, dim=1) == ftrain_label).float().mean().item()
            accuracies.append(accuracy)


        torch.save(model.state_dict(), model_save_path + f"modela0_{sub+1}{fold+1}.pth".format(epoch + 1))
        print("finish！")
        with torch.no_grad():
            predictions = model(loaded_test_data_tensor.to(device)).cpu()
            print(predictions[10])
            test_loss = creation(predictions, torch.tensor(loaded_test_labels))
            predictions.numpy()

            test_acc = (predictions.argmax(dim=1) == torch.tensor(loaded_test_labels)).sum()
            print("test acc", test_acc / len(loaded_test_labels))
            print("test loss{}".format(test_loss.item()))
            accs[sub][fold]=test_acc/len(loaded_test_labels)
average_per_experiment = np.mean(accs, axis=1)
# box
plt.boxplot(accs.T)

plt.title('Boxplot of Average per Experiment')
plt.xlabel('Experiment')
plt.ylabel('Average Value')
plt.show()

import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.callbacks import ModelCheckpoint

import numpy as np
import pickle
import argparse
import os
import shutil

from arl_eegmodels.EEGModels import EEGNet

from braindecode.datasets.moabb import MOABBDataset
from braindecode.datautil.preprocess import (
    exponential_moving_standardize, preprocess, Preprocessor)
from braindecode.datautil.windowers import create_windows_from_events

# while the default tensorflow ordering is 'channels_last' we set it here
# to be explicit in case if the user has changed the default ordering
K.set_image_data_format('channels_last')


###################
## Configuration ##
###################
# SYSTEM
parser = argparse.ArgumentParser()
parser.add_argument('--name', required=True)
parser.add_argument('--device', default=None)
parser.add_argument('--subject', type=int, required=True)
args = parser.parse_args()

if args.device:
    os.environ["CUDA_VISIBLE_DEVICES"] = args.device # use specific gpu

name = args.name
subject_id = args.subject

def make_results_directory(name, subject_id, base="."):
    results_dir = f"{base}/results/{name}_subject{subject_id}"
    if os.path.exists(results_dir):
        print(f"'{results_dir}' already exists!")
        raise
    os.mkdir(results_dir)
    
    shutil.copy("main.py", results_dir)
    print(f"'{results_dir}' is created!")
    
make_results_directory(name, subject_id, base=".")

###################
#### Load Data ####
###################
dataset = MOABBDataset(dataset_name="BNCI2014001", subject_ids=[subject_id])

low_cut_hz = 4.  # low cut frequency for filtering
high_cut_hz = 40.  # high cut frequency for filtering
# Parameters for exponential moving standardization
factor_new = 1e-3
init_block_size = 1000

preprocessors = [
    Preprocessor('pick_types', eeg=True, meg=False, stim=False),  # Keep EEG sensors
    Preprocessor(lambda x: x * 1e6),  # Convert from V to uV
    Preprocessor('resample', sfreq=128), # Added by me
    Preprocessor('filter', l_freq=low_cut_hz, h_freq=high_cut_hz),  # Bandpass filter
    Preprocessor(exponential_moving_standardize,  # Exponential moving standardization
                 factor_new=factor_new, init_block_size=init_block_size)
]
# Transform the data
preprocess(dataset, preprocessors)

trial_start_offset_seconds = 0.5
# Extract sampling frequency, check that they are same in all datasets
sfreq = dataset.datasets[0].raw.info['sfreq']
assert all([ds.raw.info['sfreq'] == sfreq for ds in dataset.datasets])
# Calculate the trial start offset in samples.
trial_start_offset_samples = int(trial_start_offset_seconds * sfreq)


trial_stop_offset_seconds = -1.5 
trial_stop_offset_samples = int(trial_stop_offset_seconds * sfreq)


# Create windows using braindecode function for this. It needs parameters to define how
# trials should be used.
windows_dataset = create_windows_from_events(
    dataset,
    trial_start_offset_samples=trial_start_offset_samples,
    trial_stop_offset_samples=trial_stop_offset_samples,
    preload=True,
)

splitted = windows_dataset.split('session')
train_set = splitted['session_T']
test_set = splitted['session_E']


X_Train = []
y_Train = []
for run in train_set.datasets:
    for X, y, _ in run:
        X_Train.append(X)
        y_Train.append(y)
X_Train = np.array(X_Train)
y_Train = np.array(y_Train)
        
X_test = []
y_test = []
for run in test_set.datasets:
    for X, y, _ in run:
        X_test.append(X)
        y_test.append(y)
X_test = np.array(X_test)
y_test = np.array(y_test)

print("X_Train shape:", X_Train.shape, "y_Train shape:", y_Train.shape)
print("X_test shape:", X_test.shape, "y_test shape:", y_test.shape)

fold1 = list(range(0,  96))
fold2 = list(range(96, 192))
fold3 = list(range(192,288))
train_val_split = [
    (fold2 + fold3, fold1),    
    (fold3 + fold1, fold2),
    (fold1 + fold2, fold3)
]
    
###########################
#### Cross Validation #####
###########################

fold_hist = []
for fold_step, (train_index, val_index) in enumerate(train_val_split):
    print(f"Start fold {fold_step}")
    print("Train", train_index)
    print("Val", val_index)
    X_train = X_Train[train_index]
    y_train = y_Train[train_index]
    X_val   = X_Train[val_index]
    y_val   = y_Train[val_index]
    
    ###################
    #### Modeling #####
    ###################
    model = EEGNet(nb_classes=4,
                   Chans=22,
                   Samples=256, # 2 seconds at 128 Hz
                   dropoutRate=0.5,
                   kernLength=32, # for SMR data
                   F1=8,
                   D=2,
                   F2=16,
                   norm_rate=0.25, # FC layer
                   dropoutType="Dropout")
    
    ####################
    #    Training      #
    ####################
    # set a valid path for your system to record model checkpoints
    checkpointer = ModelCheckpoint(filepath=f'results/{name}_subject{subject_id}/{name}_subject{subject_id}_fold{fold_step}.h5', verbose=1,
                                   save_best_only=True)

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    train_hist = model.fit(X_train, y_train, 
             batch_size=64, epochs=500,
             verbose=2, validation_data=(X_val, y_val),
             callbacks=[checkpointer])
    
    # load optimal weights
    model.load_weights(f'results/{name}_subject{subject_id}/{name}_subject{subject_id}_fold{fold_step}.h5')
    test_hist = model.evaluate(X_test,  y_test, verbose=2)
    
    fold_hist.append({"train_hist":train_hist.history, "test_hist":test_hist})
    
    fold_best_val_acc = [np.max(_["train_hist"]["val_accuracy"]) for _ in fold_hist]
fold_test_acc = [_["test_hist"][1] for _ in fold_hist]
fold_mean_test_acc = np.mean([_["test_hist"][1] for _ in fold_hist])

print("\n##############################################")
print("subejct", subject_id, "- 3fold best val  accuracy", fold_best_val_acc)
print("subejct", subject_id, "- 3fold      test accuracy", fold_test_acc)
print("subejct", subject_id, "- 3fold mean test accuract:", fold_mean_test_acc)

with open(f"results/{name}_subject{subject_id}/fold_hist.pkl", "wb") as f:
    pickle.dump(fold_hist, f)