FB-EEGNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.regularizers import l1_l2\n",
    "from tensorflow.keras.constraints import max_norm\n",
    "from tensorflow.keras import backend as K\n",
    "import tensorflow.keras.models as models\n",
    "import tensorflow.compat.v1 as tf\n",
    "import numpy as np\n",
    "from tslearn.metrics import soft_dtw\n",
    "import os\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Dense, Activation, Permute, Dropout, Input\n",
    "from tensorflow.keras.layers import Conv2D, MaxPooling2D, AveragePooling2D\n",
    "from tensorflow.keras.layers import SeparableConv2D, DepthwiseConv2D\n",
    "from tensorflow.keras.layers import BatchNormalization\n",
    "from tensorflow.keras.layers import SpatialDropout2D\n",
    "from tensorflow.keras.regularizers import l1_l2\n",
    "from tensorflow.keras.layers import Input, Flatten\n",
    "from tensorflow.keras.constraints import max_norm\n",
    "from tensorflow.keras import backend as K\n",
    "from scipy.signal import butter, filtfilt\n",
    "from sklearn.model_selection import KFold\n",
    "from scipy import io, signal\n",
    "import math\n",
    "import numpy as np\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def buffer(data, duration):\n",
    "    number_segments = int(len(data) / duration)\n",
    "    temp_buf = [data[i:i + duration] for i in range(0, len(data), duration)]\n",
    "    segmented_data = np.vstack(temp_buf[0:number_segments])\n",
    "    return segmented_data\n",
    "\n",
    "\n",
    "def get_segmented_data(data, window_len, sample_rate):\n",
    "    num_classes = data.shape[0]\n",
    "    num_chan = data.shape[1]\n",
    "    num_trials = data.shape[3]\n",
    "\n",
    "    duration = int(window_len * sample_rate)\n",
    "\n",
    "    number_segments = int(data.shape[2] / duration)\n",
    "    \n",
    "    segmented_data = np.zeros((data.shape[0], data.shape[1],\n",
    "                               data.shape[3], number_segments, duration))\n",
    "\n",
    "    for target in range(0, num_classes):\n",
    "        for channel in range(0, num_chan):\n",
    "            for trial in range(0, num_trials):\n",
    "                segmented_data[target, channel, trial, :, :] = buffer(data[target, channel, :, trial],\n",
    "                                                                      duration)\n",
    "\n",
    "    return segmented_data\n",
    "\n",
    "\n",
    "\n",
    "def get_filtered_eeg(eeg, lowcut, highcut, order, sample_rate):\n",
    "    '''\n",
    "    Returns bandpass filtered eeg for all channels and trials.\n",
    "\n",
    "    Args:\n",
    "        eeg (numpy.ndarray): raw eeg data of shape (num_classes, num_channels, num_samples, num_trials).\n",
    "        lowcut (float): lower cutoff frequency (Hz).\n",
    "        highcut (float): lower cutoff frequency (Hz).\n",
    "        order (int): order of the bandpass filter.\n",
    "        sample_rate (float): sampling rate (Hz).\n",
    "\n",
    "    Returns:\n",
    "        (numpy.ndarray): bandpass filtered eeg of shape (num_classes, num_channels, num_samples, num_trials).\n",
    "    '''\n",
    "    \n",
    "    num_classes = eeg.shape[0]\n",
    "    num_chan = eeg.shape[1]\n",
    "    total_trial_len = eeg.shape[2]\n",
    "    num_trials = eeg.shape[3]\n",
    "    \n",
    "    filtered_data = np.zeros((eeg.shape[0], eeg.shape[1], eeg.shape[2], eeg.shape[3]))\n",
    "\n",
    "    for target in range(0, num_classes):\n",
    "        for channel in range(0, num_chan):\n",
    "            for trial in range(0, num_trials):\n",
    "                signal_to_filter = np.squeeze(eeg[target, channel, :, \n",
    "                                               trial])\n",
    "                filtered_data[target, channel, :, trial] = butter_bandpass_filter(signal_to_filter, lowcut, \n",
    "                                                                                  highcut, sample_rate, order)\n",
    "    return filtered_data\n",
    "\n",
    "def butter_bandpass_filter(data, lowcut, highcut, sample_rate, order):\n",
    "    '''\n",
    "    Returns bandpass filtered data between the frequency ranges specified in the input.\n",
    "\n",
    "    Args:\n",
    "        data (numpy.ndarray): array of samples. \n",
    "        lowcut (float): lower cutoff frequency (Hz).\n",
    "        highcut (float): lower cutoff frequency (Hz).\n",
    "        sample_rate (float): sampling rate (Hz).\n",
    "        order (int): order of the bandpass filter.\n",
    "\n",
    "    Returns:\n",
    "        (numpy.ndarray): bandpass filtered data.\n",
    "    '''\n",
    "    \n",
    "    nyq = 0.5 * sample_rate\n",
    "    low = lowcut / nyq\n",
    "    high = highcut / nyq\n",
    "    b, a = butter(order, [low, high], btype='band')\n",
    "    y = filtfilt(b, a, data)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def EEGNet(input_shape):\n",
    "    model = models.Sequential()\n",
    "    model.add(layers.Conv2D(72, kernel_size=(1, 250), \n",
    "                     kernel_regularizer=keras.regularizers.l2(0.0001), \n",
    "                     kernel_initializer=keras.initializers.RandomNormal(mean=0.0, stddev=0.01, seed=None),\n",
    "                     use_bias=False,\n",
    "                     input_shape=input_shape, \n",
    "                     padding=\"same\"))\n",
    "    model.add(layers.BatchNormalization())\n",
    "    \n",
    "    model.add(layers.DepthwiseConv2D((8, 1), \n",
    "                                     kernel_regularizer=keras.regularizers.l2(0.0001), \n",
    "                                     kernel_initializer=keras.initializers.RandomNormal(mean=0.0, stddev=0.01, seed=None),\n",
    "                                     depth_multiplier=1, depthwise_constraint=max_norm(1.),  use_bias=False))\n",
    "    model.add(layers.BatchNormalization())\n",
    "    model.add(layers.Activation('elu'))\n",
    "    model.add(layers.AvgPool2D(pool_size=(1, 4), padding=\"same\"))\n",
    "    model.add(layers.Dropout(0.5))  \n",
    "    \n",
    "    \n",
    "    model.add(layers.SeparableConv2D(filters=72,\n",
    "                                     kernel_regularizer=keras.regularizers.l2(0.0001),\n",
    "                                     kernel_initializer=keras.initializers.RandomNormal(mean=0.0, stddev=0.01, seed=None),\n",
    "                                     kernel_size=(1, 32), padding=\"same\", use_bias=False))\n",
    "    model.add(layers.BatchNormalization())\n",
    "    model.add(layers.Activation(activation=\"elu\"))\n",
    "    model.add(layers.AvgPool2D(pool_size=(1, 8), padding=\"same\"))\n",
    "    model.add(layers.Dropout(0.5))\n",
    "    model.add(layers.Flatten())\n",
    "    model.add(layers.Dense(9,  activation=\"sigmoid\"))\n",
    "    return model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "all_acc_1 = np.zeros((10, 1))\n",
    "all_acc_fusion = np.zeros((10, 1))\n",
    "\n",
    "for n in range(1, 1):\n",
    "    print(f\"subject{n} is training.......\" , flush=True)\n",
    "    data = io.loadmat(f'./S{n}.mat')['eeg']\n",
    "    data = data[:, :, :]\n",
    "    data = get_segmented_data(data, window_length, 250)\n",
    "    print(data.shape)\n",
    "    cv_acc_1 = np.zeros((15, 1))\n",
    "    cv_acc_fusion = np.zeros((20, 1))\n",
    "    for block in range(0, 20):\n",
    "        train_i = []\n",
    "        for k in range(15):\n",
    "            if k != block:\n",
    "                train_i.append(k)\n",
    "        test_i = [block]\n",
    "\n",
    "        train_data = data[:, :, train_i, :, :]\n",
    "        targets = train_data.shape[0]\n",
    "        channels = train_data.shape[1]\n",
    "        trails = train_data.shape[2]\n",
    "        segments = train_data.shape[3]\n",
    "        samples = train_data.shape[4]\n",
    "        train_data = train_data.reshape(targets, channels, trails*segments, samples)\n",
    "        train_data = train_data.transpose(0, 1, 3, 2)\n",
    "\n",
    "        train_data_1 = get_filtered_eeg(train_data, 8, 55, 2, 250)\n",
    "        temp = train_data_1[:, :, :, 0]\n",
    "        for i in range(1, trails*segments):\n",
    "            temp = np.append(temp, train_data_1[:, :, :, i], axis=0)\n",
    "        train_data_1 = temp\n",
    "        train_data_1 = np.expand_dims(train_data_1, 3)\n",
    "\n",
    "        train_data_2 = get_filtered_eeg(train_data, 16, 55, 2, 250)\n",
    "        temp = train_data_2[:, :, :, 0]\n",
    "        for i in range(1, trails*segments):\n",
    "            temp = np.append(temp, train_data_2[:, :, :, i], axis=0)\n",
    "        train_data_2 = temp\n",
    "        train_data_2 = np.expand_dims(train_data_2, 3)\n",
    "\n",
    "\n",
    "        train_data_3 = get_filtered_eeg(train_data, 24, 55, 2, 250)\n",
    "        temp = train_data_3[:, :, :, 0]\n",
    "        for i in range(1, trails*segments):\n",
    "            temp = np.append(temp, train_data_3[:, :, :, i], axis=0)\n",
    "        train_data_3 = temp\n",
    "        train_data_3 = np.expand_dims(train_data_3, 3)\n",
    "\n",
    "        label = np.arange(12)\n",
    "        label = np.tile(label, trails*segments)\n",
    "        label = keras.utils.to_categorical(label)  \n",
    "        for l in label:\n",
    "            for j in range(12):\n",
    "                if l[j] == 0:\n",
    "                    l[j] = 0.25\n",
    "                else:\n",
    "                    l[j] = 1\n",
    "        fusion_label = np.arange(12)\n",
    "        fusion_label = np.tile(fusion_label, trails*segments)\n",
    "        fusion_label = keras.utils.to_categorical(fusion_label)\n",
    "        \n",
    "        \n",
    "        test_data = data[:, :, test_i, :, :]\n",
    "        targets = test_data.shape[0]\n",
    "        channels = test_data.shape[1]\n",
    "        trails = test_data.shape[2]\n",
    "        segments = test_data.shape[3]\n",
    "        samples = test_data.shape[4]\n",
    "\n",
    "        test_data = test_data.reshape(targets, channels, trails*segments, samples)\n",
    "        test_data = test_data.transpose(0, 1, 3, 2)\n",
    "\n",
    "        test_label = np.arange(12)\n",
    "        test_label = np.tile(test_label, trails*segments)\n",
    "        test_label = keras.utils.to_categorical(test_label)  \n",
    "\n",
    "        test_data_1 = get_filtered_eeg(test_data, 8, 55, 2, 250)\n",
    "        temp = test_data_1[:, :, :, 0]\n",
    "        for i in range(1, trails*segments):\n",
    "            temp = np.append(temp, test_data_1[:, :, :, i], axis=0)\n",
    "        test_data_1 = temp\n",
    "        test_data_1 = np.expand_dims(test_data_1, 3)\n",
    "\n",
    "        test_data_2 = get_filtered_eeg(test_data, 16, 55, 2, 250)\n",
    "        temp = test_data_2[:, :, :, 0]\n",
    "        for i in range(1, trails*segments):\n",
    "            temp = np.append(temp, test_data_2[:, :, :, i], axis=0)\n",
    "        test_data_2 = temp\n",
    "        test_data_2 = np.expand_dims(test_data_2, 3)\n",
    "\n",
    "        test_data_3 = get_filtered_eeg(test_data, 24, 55, 2, 250)\n",
    "        temp = test_data_3[:, :, :, 0]\n",
    "        for i in range(1, trails*segments):\n",
    "            temp = np.append(temp, test_data_3[:, :, :, i], axis=0)\n",
    "        test_data_3 = temp\n",
    "        test_data_3 = np.expand_dims(test_data_3, 3)\n",
    "\n",
    "\n",
    "        model_1 = EEGNet(train_data_1.shape[1:])\n",
    "        model_2 = EEGNet(train_data_2.shape[1:])\n",
    "        model_3 = EEGNet(train_data_3.shape[1:])\n",
    "\n",
    "        opt = keras.optimizers.Adam(learning_rate = 0.0005)\n",
    "\n",
    "        model_1.compile(loss='categorical_crossentropy', optimizer = \"Adam\", metrics = [\"accuracy\"])\n",
    "        model_1.fit(train_data_1, label, batch_size=64, epochs=500, shuffle=True, verbose=1, validation_data=(test_data_1, test_label))\n",
    "        score_1 = model_1.evaluate(test_data_1, test_label, verbose=0)\n",
    "        print(f\"EEGNET1 acc: {score_1[1]}\", flush=True)\n",
    "        cv_acc_1[block-1, :] = score_1[1]\n",
    "        model_2.compile(loss='categorical_crossentropy', optimizer = \"Adam\", metrics = [\"accuracy\"])\n",
    "        model_2.fit(train_data_2, label, batch_size=64, epochs=500, shuffle=True, verbose=1, validation_data=(test_data_2, test_label))\n",
    "        score_2 = model_2.evaluate(test_data_2, test_label, verbose=0)\n",
    "        print(f\"EEGNET2 acc: {score_2[1]}\", flush=True)    \n",
    "\n",
    "\n",
    "        model_3.compile(loss='categorical_crossentropy', optimizer = \"Adam\", metrics = [\"accuracy\"])\n",
    "        model_3.fit(train_data_3, label, batch_size=64, epochs=500, shuffle=True, verbose=1, validation_data=(test_data_3, test_label))\n",
    "        score_3 = model_3.evaluate(test_data_3, test_label, verbose=0)\n",
    "        print(f\"EEGNET3 acc: {score_3[1]}\", flush=True)\n",
    "\n",
    "        model_1.trainable = False\n",
    "        model_2.trainable = False\n",
    "        model_3.trainable = False\n",
    "\n",
    "        features = layers.concatenate([model_1.layers[-2].output, model_2.layers[-2].output, model_3.layers[-2].output])\n",
    "        features = layers.Dense(12, activation=\"softmax\")(features)\n",
    "        \n",
    "        fusion_model = models.Model([model_1.input, model_2.input, model_3.input], features)\n",
    "\n",
    "        fusion_model.compile(loss='categorical_crossentropy', optimizer = \"Adam\", metrics = [\"accuracy\"])\n",
    "        fusion_model.fit([train_data_1, train_data_2, train_data_3], fusion_label, batch_size=64, epochs=300, shuffle=True, verbose=1, validation_data=([test_data_1, test_data_2, test_data_3], test_label))\n",
    "        score_fusion = fusion_model.evaluate([test_data_1, test_data_2, test_data_3], test_label, verbose=0)\n",
    "        print(f\"FB-EEGNET acc: {score_fusion[1]}\", flush=True)\n",
    "        cv_acc_fusion[block-1, :] = score_fusion[1]\n",
    "    print(f\"subject{n} EEGNet acc: {np.mean(cv_acc_1)}\", flush=True)\n",
    "    print(f\"subject{n} FB-EEGNet acc: {np.mean(cv_acc_fusion)}\", flush=True)\n",
    "    all_acc_1[n-1, :] = np.mean(cv_acc_1)\n",
    "    all_acc_fusion[n-1, :] = np.mean(cv_acc_fusion)\n",
    "print(f\"EEGNet all-subject mean acc: {np.mean(all_acc_1)}\", flush=True)\n",
    "print(f\"FB-EEGNet all-subject mean acc: {np.mean(all_acc_fusion)}\", flush=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}