notebook/01_DataSet_DataLoader.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Dataset and DataLoader\n",
    "\n",
    "This notebook loads the `CAUEEG` dataset, tests some useful preprocessing, and makes up the PyTorch DataLoader instances for the training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "-----\n",
    "\n",
    "## Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Minjae\\Documents\\GitHub\\caueeg-ceednet\n"
     ]
    }
   ],
   "source": [
    "# for auto-reloading external modules\n",
    "# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Load some packages\n",
    "import os\n",
    "import json\n",
    "import pprint\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "\n",
    "# custom package\n",
    "from datasets.caueeg_dataset import *\n",
    "from datasets.caueeg_script import *\n",
    "from datasets.pipeline import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch version: 1.12.1+cu113\n",
      "cuda is available.\n"
     ]
    }
   ],
   "source": [
    "print('PyTorch version:', torch.__version__)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "if torch.cuda.is_available(): print('cuda is available.')\n",
    "else: print('cuda is unavailable.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Data file path\n",
    "data_path = r'local/dataset/caueeg-dataset/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "\tdataset_name:\n",
      "\t\tCAUEEG dataset\n",
      "\n",
      "\tsignal_header:\n",
      "\t\tFp1-AVG\n",
      "\t\tF3-AVG\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\tPhotic\n",
      "\n",
      "\tdata:\n",
      "\t\t{'serial': '00001', 'age': 78, 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_rf']}\n",
      "\t\t{'serial': '00002', 'age': 56, 'symptom': ['normal', 'smi']}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '01388', 'age': 73, 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_ef']}\n",
      "\n",
      "}\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "{\n",
      "\ttask_name:\n",
      "\t\tCAUEEG-Abnormal benchmark\n",
      "\n",
      "\ttask_description:\n",
      "\t\tClassification of [Normal] and [Abnormal] symptoms\n",
      "\n",
      "\tclass_label_to_name:\n",
      "\t\t['Normal', 'Abnormal']\n",
      "\n",
      "\tclass_name_to_label:\n",
      "\t\t{'Normal': 0, 'Abnormal': 1}\n",
      "\n",
      "\ttrain_split:\n",
      "\t\t{'serial': '01258', 'age': 77, 'symptom': ['dementia', 'vd', 'sivd'], 'class_name': 'Abnormal', 'class_label': 1}\n",
      "\t\t{'serial': '00836', 'age': 80, 'symptom': ['normal', 'smi'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '00105', 'age': 71, 'symptom': ['normal', 'smi'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\n",
      "\tvalidation_split:\n",
      "\t\t{'serial': '00152', 'age': 81, 'symptom': ['dementia', 'ad', 'load'], 'class_name': 'Abnormal', 'class_label': 1}\n",
      "\t\t{'serial': '01034', 'age': 81, 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_ef'], 'class_name': 'Abnormal', 'class_label': 1}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '00684', 'age': 65, 'symptom': ['normal', 'cb_normal'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\n",
      "\ttest_split:\n",
      "\t\t{'serial': '00560', 'age': 57, 'symptom': ['normal', 'cb_normal'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\t\t{'serial': '01156', 'age': 62, 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_rf'], 'class_name': 'Abnormal', 'class_label': 1}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '00888', 'age': 52, 'symptom': ['normal', 'cb_normal'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\n",
      "}\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "{\n",
      "\ttask_name:\n",
      "\t\tCAUEEG-Dementia benchmark\n",
      "\n",
      "\ttask_description:\n",
      "\t\tClassification of [Normal], [MCI], and [Dementia] symptoms.\n",
      "\n",
      "\tclass_label_to_name:\n",
      "\t\t['Normal', 'MCI', 'Dementia']\n",
      "\n",
      "\tclass_name_to_label:\n",
      "\t\t{'Normal': 0, 'MCI': 1, 'Dementia': 2}\n",
      "\n",
      "\ttrain_split:\n",
      "\t\t{'serial': '00587', 'age': 53, 'symptom': ['normal', 'cb_normal'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\t\t{'serial': '01301', 'age': 88, 'symptom': ['dementia', 'ad', 'load'], 'class_name': 'Dementia', 'class_label': 2}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '00153', 'age': 64, 'symptom': ['mci', 'mci_non_amnestic'], 'class_name': 'MCI', 'class_label': 1}\n",
      "\n",
      "\tvalidation_split:\n",
      "\t\t{'serial': '00341', 'age': 80, 'symptom': ['dementia', 'ad', 'load'], 'class_name': 'Dementia', 'class_label': 2}\n",
      "\t\t{'serial': '00510', 'age': 82, 'symptom': ['dementia', 'vd', 'sivd'], 'class_name': 'Dementia', 'class_label': 2}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '01303', 'age': 52, 'symptom': ['normal', 'smi'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\n",
      "\ttest_split:\n",
      "\t\t{'serial': '00789', 'age': 62, 'symptom': ['normal', 'cb_normal'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\t\t{'serial': '00934', 'age': 61, 'symptom': ['normal', 'cb_normal'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t.\n",
      "\t\t{'serial': '00520', 'age': 55, 'symptom': ['normal', 'smi'], 'class_name': 'Normal', 'class_label': 0}\n",
      "\n",
      "}\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for task in ['annotation.json', 'abnormal.json', 'dementia.json']:\n",
    "    task_path = os.path.join(data_path, task)\n",
    "    with open(task_path, 'r') as json_file:\n",
    "        task_dict = json.load(json_file)\n",
    "        \n",
    "    print('{')\n",
    "    for k, v in task_dict.items():\n",
    "        print(f'\\t{k}:')\n",
    "        if isinstance(v, list) and len(v) > 3:\n",
    "            print(f'\\t\\t{v[0]}')\n",
    "            print(f'\\t\\t{v[1]}')\n",
    "            print(f'\\t\\t.')\n",
    "            print(f'\\t\\t.')\n",
    "            print(f'\\t\\t.')\n",
    "            print(f'\\t\\t{v[-1]}')\n",
    "        else:\n",
    "            print(f'\\t\\t{v}')\n",
    "        print()\n",
    "    print('}')\n",
    "    print('\\n' + '-' * 100 + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "-----\n",
    "\n",
    "## Load the CAUEEG dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load the whole CAUEEG data as a PyTorch dataset instance without considering the target task (no train/val/test sets and no class label)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'dataset_name': 'CAUEEG dataset',\n",
      " 'signal_header': ['Fp1-AVG', 'F3-AVG', 'C3-AVG', 'P3-AVG', 'O1-AVG', 'Fp2-AVG', 'F4-AVG', 'C4-AVG', 'P4-AVG', 'O2-AVG', 'F7-AVG', 'T3-AVG', 'T5-AVG', 'F8-AVG', 'T4-AVG', 'T6-AVG', 'FZ-AVG', 'CZ-AVG', 'PZ-AVG', 'EKG', 'Photic']}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 78,\n",
      " 'serial': '00001',\n",
      " 'signal': array([[  0., -11., -13., ...,   0.,   0.,   0.],\n",
      "       [ 29.,  33.,  34., ...,   0.,   0.,   0.],\n",
      "       [ -3.,  -6.,  -3., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [ -4.,  -2.,   1., ...,   0.,   0.,   0.],\n",
      "       [112.,  67.,  76., ...,   0.,   0.,   0.],\n",
      "       [ -1.,  -1.,  -1., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_rf']}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 78,\n",
      " 'serial': '00004',\n",
      " 'signal': array([[ 30.,  34.,  36., ...,   0.,   0.,   0.],\n",
      "       [ 34.,  25.,  24., ...,   0.,   0.,   0.],\n",
      "       [ 11.,  15.,  15., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [-10.,  -6.,  -4., ...,   0.,   0.,   0.],\n",
      "       [ 23.,  32.,  30., ...,   0.,   0.,   0.],\n",
      "       [ -1.,   0.,  -1., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['dementia', 'ad', 'load']}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n"
     ]
    }
   ],
   "source": [
    "config_data, full_eeg_dataset = load_caueeg_full_dataset(dataset_path=data_path, \n",
    "                                                         load_event=False, \n",
    "                                                         file_format='edf',  # can be ommitted\n",
    "                                                         transform=None)\n",
    "\n",
    "pprint.pprint(config_data, width=250)\n",
    "print('\\n', '-' * 100, '\\n')\n",
    "\n",
    "pprint.pprint(full_eeg_dataset[0])\n",
    "print('\\n', '-' * 100, '\\n')\n",
    "\n",
    "pprint.pprint(full_eeg_dataset[3])\n",
    "print('\\n', '-' * 100, '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### When triggered `load_event` option, the dataset also loads event data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'age': 78,\n",
      " 'event': [[0, 'Start Recording'],\n",
      "           [0, 'New Montage - Montage 002'],\n",
      "           [36396, 'Eyes Open'],\n",
      "           [72518, 'Eyes Closed'],\n",
      "           [73862, 'Eyes Open'],\n",
      "           [75248, 'Eyes Closed'],\n",
      "           [76728, 'swallowing'],\n",
      "           [77978, 'Eyes Open'],\n",
      "           [79406, 'Eyes Closed'],\n",
      "           [79996, 'Photic On - 3.0 Hz'],\n",
      "           [80288, 'Eyes Open'],\n",
      "           [81296, 'Eyes Closed'],\n",
      "           [82054, 'Photic Off'],\n",
      "           [84070, 'Photic On - 6.0 Hz'],\n",
      "           [84488, 'Eyes Open'],\n",
      "           [85538, 'Eyes Closed'],\n",
      "           [86086, 'Photic Off'],\n",
      "           [88144, 'Photic On - 9.0 Hz'],\n",
      "           [90160, 'Photic Off'],\n",
      "           [91458, 'Eyes Open'],\n",
      "           [92218, 'Photic On - 12.0 Hz'],\n",
      "           [92762, 'Eyes Closed'],\n",
      "           [94198, 'Photic Off'],\n",
      "           [94742, 'Eyes Open'],\n",
      "           [95708, 'Eyes Closed'],\n",
      "           [96256, 'Photic On - 15.0 Hz'],\n",
      "           [98272, 'Photic Off'],\n",
      "           [100330, 'Photic On - 18.0 Hz'],\n",
      "           [102346, 'Photic Off'],\n",
      "           [102596, 'Eyes Open'],\n",
      "           [103856, 'Eyes Closed'],\n",
      "           [104361, 'Photic On - 21.0 Hz'],\n",
      "           [106420, 'Photic Off'],\n",
      "           [106880, 'Eyes Open'],\n",
      "           [107804, 'Eyes Closed'],\n",
      "           [108435, 'Photic On - 24.0 Hz'],\n",
      "           [110452, 'Photic Off'],\n",
      "           [111080, 'Eyes Open'],\n",
      "           [112004, 'Eyes Closed'],\n",
      "           [112509, 'Photic On - 27.0 Hz'],\n",
      "           [114528, 'Photic Off'],\n",
      "           [114864, 'Eyes Open'],\n",
      "           [116124, 'Eyes Closed'],\n",
      "           [116544, 'Photic On - 30.0 Hz'],\n",
      "           [118602, 'Photic Off'],\n",
      "           [126672, 'artifact'],\n",
      "           [134030, 'Move'],\n",
      "           [135584, 'Eyes Open'],\n",
      "           [136668, 'Eyes Closed'],\n",
      "           [139818, 'Eyes Open'],\n",
      "           [141414, 'Eyes Closed'],\n",
      "           [145000, 'Paused']],\n",
      " 'serial': '00001',\n",
      " 'signal': array([[  0., -11., -13., ...,   0.,   0.,   0.],\n",
      "       [ 29.,  33.,  34., ...,   0.,   0.,   0.],\n",
      "       [ -3.,  -6.,  -3., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [ -4.,  -2.,   1., ...,   0.,   0.,   0.],\n",
      "       [112.,  67.,  76., ...,   0.,   0.,   0.],\n",
      "       [ -1.,  -1.,  -1., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_rf']}\n"
     ]
    }
   ],
   "source": [
    "config_data, full_eeg_dataset = load_caueeg_full_dataset(dataset_path=data_path, \n",
    "                                                         file_format='edf',\n",
    "                                                         transform=None)\n",
    "pprint.pprint(full_eeg_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-----\n",
    "\n",
    "## Load the CAUEEG benchmarks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load the CAUEEG-Abnormal benchmark using the PyTorch dataset instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'class_label_to_name': ['Normal', 'Abnormal'],\n",
      " 'class_name_to_label': {'Abnormal': 1, 'Normal': 0},\n",
      " 'task_description': 'Classification of [Normal] and [Abnormal] symptoms',\n",
      " 'task_name': 'CAUEEG-Abnormal benchmark'}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 57,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'serial': '00560',\n",
      " 'signal': array([[  7.,  41.,  49., ...,   0.,   0.,   0.],\n",
      "       [ 23.,  26.,  27., ...,   0.,   0.,   0.],\n",
      "       [-16., -20., -18., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [ 32.,  31.,  31., ...,   0.,   0.,   0.],\n",
      "       [177., 228., 142., ...,   0.,   0.,   0.],\n",
      "       [  0.,  -1.,   0., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False)\n",
    "pprint.pprint(config_data)\n",
    "print('\\n', '-' * 100, '\\n')\n",
    "pprint.pprint(test_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load the CAUEEG-Dementia benchmark using the PyTorch dataset instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'class_label_to_name': ['Normal', 'MCI', 'Dementia'],\n",
      " 'class_name_to_label': {'Dementia': 2, 'MCI': 1, 'Normal': 0},\n",
      " 'task_description': 'Classification of [Normal], [MCI], and [Dementia] '\n",
      "                     'symptoms.',\n",
      " 'task_name': 'CAUEEG-Dementia benchmark'}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 62,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'serial': '00789',\n",
      " 'signal': array([[-87., -69., -70., ...,   0.,   0.,   0.],\n",
      "       [-25., -18., -19., ...,   0.,   0.,   0.],\n",
      "       [ -6.,   1.,   0., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [  0.,  -5.,  -4., ...,   0.,   0.,   0.],\n",
      "       [-31.,  -8.,  -7., ...,   0.,   0.,   0.],\n",
      "       [  0.,   1.,  -1., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  transform=None)\n",
    "pprint.pprint(config_data)\n",
    "print('\\n', '-' * 100, '\\n')\n",
    "pprint.pprint(test_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load the CAUEEG-Abnormal (no-overlap version) benchmark using the PyTorch dataset instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'class_label_to_name': ['Normal', 'Abnormal'],\n",
      " 'class_name_to_label': {'Abnormal': 1, 'Normal': 0},\n",
      " 'task_description': 'Classification of [Normal] and [Abnormal] symptoms',\n",
      " 'task_name': 'CAUEEG-Abnormal benchmark'}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 57,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'serial': '00560',\n",
      " 'signal': array([[  7.,  41.,  49., ...,   0.,   0.,   0.],\n",
      "       [ 23.,  26.,  27., ...,   0.,   0.,   0.],\n",
      "       [-16., -20., -18., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [ 32.,  31.,  31., ...,   0.,   0.,   0.],\n",
      "       [177., 228., 142., ...,   0.,   0.,   0.],\n",
      "       [  0.,  -1.,   0., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal-no-overlap',\n",
    "                                                                                  load_event=False)\n",
    "pprint.pprint(config_data)\n",
    "print('\\n', '-' * 100, '\\n')\n",
    "pprint.pprint(test_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load the CAUEEG-Dementia (no-overlap version) benchmark using the PyTorch dataset instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'class_label_to_name': ['Normal', 'MCI', 'Dementia'],\n",
      " 'class_name_to_label': {'Dementia': 2, 'MCI': 1, 'Normal': 0},\n",
      " 'task_description': 'Classification of [Normal], [MCI], and [Dementia] '\n",
      "                     'symptoms.',\n",
      " 'task_name': 'CAUEEG-Dementia benchmark'}\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 70,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'serial': '00142',\n",
      " 'signal': array([[-38., -24., -18., ...,   0.,   0.,   0.],\n",
      "       [ -7.,  -4.,  -2., ...,   0.,   0.,   0.],\n",
      "       [ -4.,  -6.,  -7., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [ -4.,  -6.,  -7., ...,   0.,   0.,   0.],\n",
      "       [-33.,  13.,  22., ...,   0.,   0.,   0.],\n",
      "       [ -1.,   0.,   1., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia-no-overlap',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  transform=None)\n",
    "pprint.pprint(config_data)\n",
    "print('\\n', '-' * 100, '\\n')\n",
    "pprint.pprint(test_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### With `load_event` triggered, the benchmark can use the event data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'age': 57,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'event': [[0, 'Start Recording'],\n",
      "           [0, 'New Montage - Montage 002'],\n",
      "           [5408, 'Eyes Open'],\n",
      "           [6332, 'Eyes Closed'],\n",
      "           [16832, 'Eyes Open'],\n",
      "           [17756, 'Eyes Closed'],\n",
      "           [24200, 'Paused'],\n",
      "           [28400, 'Recording Resumed'],\n",
      "           [67442, 'Eyes Open'],\n",
      "           [68366, 'Eyes Closed'],\n",
      "           [82856, 'Eyes Open'],\n",
      "           [83738, 'Eyes Closed'],\n",
      "           [109833, 'Eyes Open'],\n",
      "           [109833, 'Eyes Closed'],\n",
      "           [132600, 'Paused']],\n",
      " 'serial': '00560',\n",
      " 'signal': array([[  7.,  41.,  49., ...,   0.,   0.,   0.],\n",
      "       [ 23.,  26.,  27., ...,   0.,   0.,   0.],\n",
      "       [-16., -20., -18., ...,   0.,   0.,   0.],\n",
      "       ...,\n",
      "       [ 32.,  31.,  31., ...,   0.,   0.,   0.],\n",
      "       [177., 228., 142., ...,   0.,   0.,   0.],\n",
      "       [  0.,  -1.,   0., ...,   0.,   0.,   0.]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=True, \n",
    "                                                                                  file_format='edf', \n",
    "                                                                                  transform=None)\n",
    "pprint.pprint(test_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utilizing `PyArrow.feather` is much faster than directly using `EDF`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: total: 3.84 s\n",
      "Wall time: 4.66 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='edf', \n",
    "                                                                                  transform=None)\n",
    "for i, d in enumerate(test_dataset):\n",
    "    if i > 30:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: total: 0 ns\n",
      "Wall time: 0 ns\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather', \n",
    "                                                                                  transform=None)\n",
    "for i, d in enumerate(test_dataset):\n",
    "    if i > 30:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---\n",
    "\n",
    "## PyTorch Transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Random crop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'age': 62,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'serial': '00789',\n",
      " 'signal': array([[  -26,   -26,   -24, ...,   -25,   -31,   -32],\n",
      "       [  -26,   -30,   -25, ...,   -15,    -9,    -9],\n",
      "       [  -12,    -6,    -1, ...,     5,    10,    11],\n",
      "       ...,\n",
      "       [    7,     8,     5, ...,     7,     0,    -5],\n",
      "       [   61,    52,    43, ...,  -743, -1006, -1129],\n",
      "       [    1,    -1,     0, ...,    -2,    -4,     0]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n",
      "\n",
      ">>> signal shape: (21, 100)\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n",
      "{'age': 62,\n",
      " 'class_label': 0,\n",
      " 'class_name': 'Normal',\n",
      " 'serial': '00789',\n",
      " 'signal': array([[ 11,   2,  -8, ..., -21, -22, -17],\n",
      "       [ 33,  24,  14, ...,   0,   8,  13],\n",
      "       [  7,  -3, -14, ...,   0,   6,   8],\n",
      "       ...,\n",
      "       [ 13,   7,   0, ...,  -2,  -2,  -4],\n",
      "       [ 73,  83,  91, ...,  10,  23,  43],\n",
      "       [ -1,   0,   2, ...,  -2,  -2,   0]]),\n",
      " 'symptom': ['normal', 'cb_normal']}\n",
      "\n",
      ">>> signal shape: (21, 100)\n",
      "\n",
      " ---------------------------------------------------------------------------------------------------- \n",
      "\n"
     ]
    }
   ],
   "source": [
    "transform = EegRandomCrop(crop_length=100)\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path,\n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=transform)\n",
    "for i in range(2):\n",
    "    d = test_dataset[0]\n",
    "    pprint.pprint(d)\n",
    "    print()\n",
    "    print('>>> signal shape:', d['signal'].shape)\n",
    "    print('\\n', '-' * 100, '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "tags": []
   },
   "source": [
    "### Drop channel(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "channel_ekg:  19\n",
      "channel_photic:  20\n"
     ]
    }
   ],
   "source": [
    "anno_path = os.path.join(data_path, 'annotation.json')\n",
    "with open(anno_path, 'r') as json_file:\n",
    "    annotation = json.load(json_file)\n",
    "signal_headers = annotation['signal_header']\n",
    "\n",
    "channel_ekg = signal_headers.index('EKG')\n",
    "channel_photic = signal_headers.index('Photic')\n",
    "\n",
    "print('channel_ekg: ', channel_ekg)\n",
    "print('channel_photic: ', channel_photic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before: (21, 128400)\n",
      "[[  7  41  49 ...  24  27  29]\n",
      " [ 23  26  27 ... -60 -58 -57]\n",
      " [-16 -20 -18 ...  17  18  17]\n",
      " ...\n",
      " [ 32  31  31 ...  18  16  15]\n",
      " [177 228 142 ...   2   1   0]\n",
      " [  0  -1   0 ...  -1   0   1]]\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "after: (20, 128400)\n",
      "[[  7  41  49 ...  24  27  29]\n",
      " [ 23  26  27 ... -60 -58 -57]\n",
      " [-16 -20 -18 ...  17  18  17]\n",
      " ...\n",
      " [  3   7   7 ...  14  14  12]\n",
      " [ 32  31  31 ...  18  16  15]\n",
      " [  0  -1   0 ...  -1   0   1]]\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather', \n",
    "                                                                                  transform=None)\n",
    "print('before:', test_dataset[0]['signal'].shape)\n",
    "print(test_dataset[0]['signal'])\n",
    "\n",
    "print()\n",
    "print('-' * 100)\n",
    "print()\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather', \n",
    "                                                                                  transform=EegDropChannels(channel_ekg))\n",
    "print('after:', test_dataset[0]['signal'].shape)\n",
    "print(test_dataset[0]['signal'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before: (21, 128400)\n",
      "[[  7  41  49 ...  24  27  29]\n",
      " [ 23  26  27 ... -60 -58 -57]\n",
      " [-16 -20 -18 ...  17  18  17]\n",
      " ...\n",
      " [ 32  31  31 ...  18  16  15]\n",
      " [177 228 142 ...   2   1   0]\n",
      " [  0  -1   0 ...  -1   0   1]]\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "after: (20, 128400)\n",
      "[[  7  41  49 ...  24  27  29]\n",
      " [ 23  26  27 ... -60 -58 -57]\n",
      " [-16 -20 -18 ...  17  18  17]\n",
      " ...\n",
      " [  3   7   7 ...  14  14  12]\n",
      " [ 32  31  31 ...  18  16  15]\n",
      " [177 228 142 ...   2   1   0]]\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=None)\n",
    "print('before:', test_dataset[0]['signal'].shape)\n",
    "print(test_dataset[0]['signal'])\n",
    "\n",
    "print()\n",
    "print('-' * 100)\n",
    "print()\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=EegDropChannels(channel_photic))\n",
    "print('after:', test_dataset[0]['signal'].shape)\n",
    "print(test_dataset[0]['signal'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before: (21, 128400)\n",
      "[[  7  41  49 ...  24  27  29]\n",
      " [ 23  26  27 ... -60 -58 -57]\n",
      " [-16 -20 -18 ...  17  18  17]\n",
      " ...\n",
      " [ 32  31  31 ...  18  16  15]\n",
      " [177 228 142 ...   2   1   0]\n",
      " [  0  -1   0 ...  -1   0   1]]\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "after: (19, 128400)\n",
      "[[  7  41  49 ...  24  27  29]\n",
      " [ 23  26  27 ... -60 -58 -57]\n",
      " [-16 -20 -18 ...  17  18  17]\n",
      " ...\n",
      " [-26 -11 -11 ...  18  18  17]\n",
      " [  3   7   7 ...  14  14  12]\n",
      " [ 32  31  31 ...  18  16  15]]\n"
     ]
    }
   ],
   "source": [
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=None)\n",
    "print('before:', test_dataset[0]['signal'].shape)\n",
    "print(test_dataset[0]['signal'])\n",
    "\n",
    "print()\n",
    "print('-' * 100)\n",
    "print()\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=EegDropChannels([channel_ekg, channel_photic]))\n",
    "print('after:', test_dataset[0]['signal'].shape)\n",
    "print(test_dataset[0]['signal'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### To Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before:\n",
      "{'age': 78,\n",
      " 'serial': '00001',\n",
      " 'signal': array([[  0, -11, -13, ...,  18,  21,  22],\n",
      "       [ 29,  33,  34, ...,  -7,  -4,  -4],\n",
      "       [ -3,  -6,  -3, ...,  -1,   1,   2],\n",
      "       ...,\n",
      "       [ -4,  -2,   1, ...,   0,  -1,  -1],\n",
      "       [112,  67,  76, ..., -13, -15, -11],\n",
      "       [ -1,  -1,  -1, ...,  -1,  -1,  -1]]),\n",
      " 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_rf']}\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "After:\n",
      "{'age': tensor(78.),\n",
      " 'serial': '00001',\n",
      " 'signal': tensor([[  0., -11., -13.,  ...,  18.,  21.,  22.],\n",
      "        [ 29.,  33.,  34.,  ...,  -7.,  -4.,  -4.],\n",
      "        [ -3.,  -6.,  -3.,  ...,  -1.,   1.,   2.],\n",
      "        ...,\n",
      "        [ -4.,  -2.,   1.,  ...,   0.,  -1.,  -1.],\n",
      "        [112.,  67.,  76.,  ..., -13., -15., -11.],\n",
      "        [ -1.,  -1.,  -1.,  ...,  -1.,  -1.,  -1.]]),\n",
      " 'symptom': ['mci', 'mci_amnestic', 'mci_amnestic_rf']}\n"
     ]
    }
   ],
   "source": [
    "config_data, full_eeg_dataset = load_caueeg_full_dataset(dataset_path=data_path, \n",
    "                                                         load_event=False, \n",
    "                                                         file_format='feather',\n",
    "                                                         transform=None)\n",
    "print('Before:')\n",
    "pprint.pprint(full_eeg_dataset[0])\n",
    "\n",
    "print()\n",
    "print('-' * 100)\n",
    "print()\n",
    "\n",
    "config_data, full_eeg_dataset = load_caueeg_full_dataset(dataset_path=data_path, \n",
    "                                                         load_event=False, \n",
    "                                                         file_format='feather',\n",
    "                                                         transform=EegToTensor())\n",
    "print('After:')\n",
    "pprint.pprint(full_eeg_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Compose the above all in one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'age': tensor(77.),\n",
      " 'class_label': tensor(1),\n",
      " 'class_name': 'Abnormal',\n",
      " 'serial': '01258',\n",
      " 'signal': tensor([[-36., -39., -39.,  ...,  -3., -10., -10.],\n",
      "        [  3.,   3.,  -4.,  ...,  -6.,   7.,   7.],\n",
      "        [-19., -19., -21.,  ...,  -9.,  -8.,  -5.],\n",
      "        ...,\n",
      "        [ -3.,  -6.,  -9.,  ...,  -9.,  -7.,  -7.],\n",
      "        [  8.,   7.,   7.,  ...,   7.,   5.,   4.],\n",
      "        [-12., -12., -11.,  ...,  19.,   3.,  -4.]]),\n",
      " 'symptom': ['dementia', 'vd', 'sivd']}\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=200*10,       # crop: 10s\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=transform)\n",
    "\n",
    "pprint.pprint(train_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---\n",
    "\n",
    "## PyTorch DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "if device.type == 'cuda':\n",
    "    num_workers = 0  # A number other than 0 causes an error\n",
    "    pin_memory = True\n",
    "else:\n",
    "    num_workers = 0\n",
    "    pin_memory = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'age': tensor([76., 85., 80., 39.]),\n",
      " 'serial': ['00375', '00252', '00651', '01107'],\n",
      " 'signal': tensor([[[  15.,   14.,   19.,  ...,    0.,    3.,    3.],\n",
      "         [   0.,   -1.,    0.,  ...,   19.,   15.,   12.],\n",
      "         [  -3.,   -1.,    1.,  ...,    8.,    7.,    5.],\n",
      "         ...,\n",
      "         [ -11.,  -10.,   -8.,  ...,   -6.,   -5.,   -4.],\n",
      "         [  -5.,   -5.,   -5.,  ...,   -6.,   -6.,   -6.],\n",
      "         [  -7.,  -18.,  -26.,  ...,  -43.,  -39.,  -40.]],\n",
      "\n",
      "        [[-131., -126., -127.,  ...,   58.,   59.,   51.],\n",
      "         [  27.,   24.,   21.,  ...,   20.,   16.,   13.],\n",
      "         [ -45.,  -46.,  -47.,  ...,   13.,    7.,    3.],\n",
      "         ...,\n",
      "         [ -31.,  -31.,  -31.,  ...,  -11.,  -13.,  -15.],\n",
      "         [ -64.,  -58.,  -53.,  ...,   -6.,   -6.,   -4.],\n",
      "         [  -4.,    7.,   16.,  ...,  -17.,  -17.,  -30.]],\n",
      "\n",
      "        [[ -23.,  -20.,  -18.,  ...,    3.,    0.,    1.],\n",
      "         [ -18.,  -15.,  -14.,  ...,    6.,    6.,    9.],\n",
      "         [  14.,   15.,   16.,  ...,  -27.,  -23.,  -18.],\n",
      "         ...,\n",
      "         [   9.,    8.,    2.,  ...,   -2.,   -2.,   -4.],\n",
      "         [  13.,   12.,    8.,  ...,    8.,    7.,    2.],\n",
      "         [  88.,   36.,   35.,  ...,  -27.,  -53.,  -24.]],\n",
      "\n",
      "        [[ -74.,  -78.,  -81.,  ...,   15.,   15.,   15.],\n",
      "         [ -26.,  -28.,  -28.,  ...,   -4.,   -5.,   -8.],\n",
      "         [  -5.,   -5.,   -4.,  ...,   11.,    9.,    6.],\n",
      "         ...,\n",
      "         [ -12.,  -14.,  -16.,  ...,   -1.,   -2.,   -5.],\n",
      "         [   0.,    0.,   -2.,  ...,   -4.,   -6.,   -6.],\n",
      "         [ -55.,  -53.,  -66.,  ...,  -14.,   -5.,  -15.]]]),\n",
      " 'symptom': [['dementia', 'ad', 'load'],\n",
      "             ['dementia', 'vd', 'sivd'],\n",
      "             ['dementia', 'ad', 'load'],\n",
      "             ['tga']]}\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=200*10,       # crop: 10s\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "\n",
    "config_data, full_eeg_dataset = load_caueeg_full_dataset(dataset_path=data_path, \n",
    "                                                         load_event=False, \n",
    "                                                         file_format='feather',\n",
    "                                                         transform=transform)\n",
    "\n",
    "full_loader = DataLoader(full_eeg_dataset,\n",
    "                         batch_size=4,\n",
    "                         shuffle=True,\n",
    "                         drop_last=True,\n",
    "                         num_workers=num_workers,\n",
    "                         pin_memory=pin_memory,\n",
    "                         collate_fn=eeg_collate_fn)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(full_loader):\n",
    "    pprint.pprint(sample_batched)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'age': tensor([78., 66., 75., 84., 80., 61., 82., 54.]),\n",
      " 'class_label': tensor([1, 1, 0, 1, 1, 0, 1, 0]),\n",
      " 'class_name': ['Abnormal', 'Abnormal', 'Normal', 'Abnormal', 'Abnormal', 'Normal', 'Abnormal', 'Normal'],\n",
      " 'serial': ['01134', '01016', '00093', '00039', '00481', '00900', '00229', '00851'],\n",
      " 'signal': tensor([[[ 267.,  264.,  265.,  ...,    6.,    9.,    9.],\n",
      "         [  11.,    6.,    3.,  ...,    3.,   -6.,   -1.],\n",
      "         [   2.,    0.,   -3.,  ...,   13.,   11.,   11.],\n",
      "         ...,\n",
      "         [  -8.,  -10.,  -13.,  ...,  -12.,   -7.,   -6.],\n",
      "         [  10.,   11.,   11.,  ...,   -3.,    1.,    3.],\n",
      "         [  22.,   32.,   19.,  ...,   19.,    4.,    1.]],\n",
      "\n",
      "        [[  13.,   11.,   10.,  ...,   57.,   55.,   53.],\n",
      "         [  12.,   10.,    9.,  ...,   15.,   14.,   14.],\n",
      "         [  -7.,   -5.,   -7.,  ...,   -8.,   -6.,   -7.],\n",
      "         ...,\n",
      "         [  -5.,   -3.,    0.,  ...,   -1.,    1.,    1.],\n",
      "         [  -1.,    1.,    1.,  ...,   -4.,   -2.,   -3.],\n",
      "         [ -65.,  -57.,  -46.,  ...,    1.,   20.,   26.]],\n",
      "\n",
      "        [[  17.,   14.,    6.,  ...,    4.,    7.,   14.],\n",
      "         [  26.,   24.,   17.,  ...,  -13.,  -12.,   -7.],\n",
      "         [ -20.,  -22.,  -22.,  ...,  -11.,   -9.,   -8.],\n",
      "         ...,\n",
      "         [   0.,    6.,    6.,  ...,   -6.,   -9.,  -12.],\n",
      "         [ -14.,  -16.,  -15.,  ...,   -1.,   -4.,   -7.],\n",
      "         [  -1.,   -1.,   -4.,  ...,    2.,    4.,   -4.]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[  -3.,   -4.,   -9.,  ...,   -7.,   -6.,   -7.],\n",
      "         [   6.,    5.,    3.,  ...,   -7.,   -5.,   -3.],\n",
      "         [   2.,    1.,   -1.,  ...,   -4.,   -2.,    0.],\n",
      "         ...,\n",
      "         [   6.,    4.,    0.,  ...,  -13.,  -13.,  -11.],\n",
      "         [  10.,   10.,    9.,  ...,  -17.,  -18.,  -16.],\n",
      "         [ -26.,  -48., -118.,  ...,  -45.,  -49.,  -42.]],\n",
      "\n",
      "        [[  21.,   17.,   14.,  ...,   29.,   27.,   30.],\n",
      "         [  -4.,  -16.,  -12.,  ...,   18.,   26.,   32.],\n",
      "         [  11.,    7.,    8.,  ...,   -3.,    6.,    8.],\n",
      "         ...,\n",
      "         [  -5.,   -2.,    1.,  ...,    2.,    4.,    4.],\n",
      "         [  -1.,    6.,   12.,  ...,    6.,    4.,    2.],\n",
      "         [-159.,   62.,   37.,  ...,  -16.,  -14.,  -16.]],\n",
      "\n",
      "        [[ -86.,  -17.,   18.,  ...,  -24.,  -21.,  -19.],\n",
      "         [   4.,  -11.,  -19.,  ...,    9.,   10.,   10.],\n",
      "         [ -10.,   -9.,   -8.,  ...,   -6.,   -5.,   -3.],\n",
      "         ...,\n",
      "         [  -5.,   -7.,   -8.,  ...,    5.,    8.,    9.],\n",
      "         [  -9.,  -10.,  -10.,  ...,   -8.,   -7.,   -5.],\n",
      "         [   9.,   15.,   -6.,  ...,  -28.,  -14.,    9.]]]),\n",
      " 'symptom': [['mci', 'mci_amnestic'],\n",
      "             ['mci', 'mci_amnestic', 'mci_amnestic_ef'],\n",
      "             ['normal', 'cb_normal'],\n",
      "             ['dementia', 'ad', 'load'],\n",
      "             ['mci', 'mci_amnestic', 'mci_amnestic_ef'],\n",
      "             ['normal', 'cb_normal'],\n",
      "             ['mci', 'mci_amnestic', 'mci_amnestic_rf'],\n",
      "             ['normal', 'cb_normal']]}\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=200*10,       # crop: 10s\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='abnormal',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=transform)\n",
    "\n",
    "test_loader = DataLoader(test_dataset,\n",
    "                          batch_size=8,\n",
    "                          shuffle=True,\n",
    "                          drop_last=True,\n",
    "                          num_workers=num_workers,\n",
    "                          pin_memory=pin_memory,\n",
    "                          collate_fn=eeg_collate_fn)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(test_loader):\n",
    "    pprint.pprint(sample_batched, width=250)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---\n",
    "\n",
    "## Some preprocessing steps can be implemented via the PyTorch Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=200*10,       # crop: 10s\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size=2,\n",
    "                          shuffle=True,\n",
    "                          drop_last=True,\n",
    "                          num_workers=num_workers,\n",
    "                          pin_memory=pin_memory,\n",
    "                          collate_fn=eeg_collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### To GPU device if it is possible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cuda\n",
      "\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      ")\n",
      "- Before -\n",
      "{'age': tensor([76., 64.]),\n",
      " 'class_label': tensor([1, 2]),\n",
      " 'class_name': ['MCI', 'Dementia'],\n",
      " 'serial': ['00378', '00633'],\n",
      " 'signal': tensor([[[-89., -89., -84.,  ..., -15., -15., -15.],\n",
      "         [-79., -73., -70.,  ..., -11., -12., -13.],\n",
      "         [ 29.,  25.,  25.,  ...,   7.,   7.,   6.],\n",
      "         ...,\n",
      "         [ 38.,  38.,  40.,  ...,   7.,   9.,   9.],\n",
      "         [ 32.,  33.,  35.,  ...,   0.,   1.,   0.],\n",
      "         [  0.,  -4., -11.,  ..., -60., -58., -53.]],\n",
      "\n",
      "        [[ -8., -11., -12.,  ...,   5.,   6.,   8.],\n",
      "         [ -5.,  -6.,  -7.,  ...,   4.,   4.,   5.],\n",
      "         [  3.,   3.,   3.,  ...,   0.,   0.,  -1.],\n",
      "         ...,\n",
      "         [ -3.,  -3.,  -4.,  ...,   5.,   5.,   6.],\n",
      "         [ -6.,  -5.,  -5.,  ...,  -5.,  -5.,  -5.],\n",
      "         [-14.,   1.,   8.,  ...,  -4.,  -5.,  -6.]]]),\n",
      " 'symptom': [['mci', 'mci_vascular'], ['dementia', 'ad', 'eoad']]}\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "- After -\n",
      "{'age': tensor([76., 64.], device='cuda:0'),\n",
      " 'class_label': tensor([1, 2], device='cuda:0'),\n",
      " 'class_name': ['MCI', 'Dementia'],\n",
      " 'serial': ['00378', '00633'],\n",
      " 'signal': tensor([[[-89., -89., -84.,  ..., -15., -15., -15.],\n",
      "         [-79., -73., -70.,  ..., -11., -12., -13.],\n",
      "         [ 29.,  25.,  25.,  ...,   7.,   7.,   6.],\n",
      "         ...,\n",
      "         [ 38.,  38.,  40.,  ...,   7.,   9.,   9.],\n",
      "         [ 32.,  33.,  35.,  ...,   0.,   1.,   0.],\n",
      "         [  0.,  -4., -11.,  ..., -60., -58., -53.]],\n",
      "\n",
      "        [[ -8., -11., -12.,  ...,   5.,   6.,   8.],\n",
      "         [ -5.,  -6.,  -7.,  ...,   4.,   4.,   5.],\n",
      "         [  3.,   3.,   3.,  ...,   0.,   0.,  -1.],\n",
      "         ...,\n",
      "         [ -3.,  -3.,  -4.,  ...,   5.,   5.,   6.],\n",
      "         [ -6.,  -5.,  -5.,  ...,  -5.,  -5.,  -5.],\n",
      "         [-14.,   1.,   8.,  ...,  -4.,  -5.,  -6.]]], device='cuda:0'),\n",
      " 'symptom': [['mci', 'mci_vascular'], ['dementia', 'ad', 'eoad']]}\n"
     ]
    }
   ],
   "source": [
    "print('device:', device)\n",
    "print()\n",
    "\n",
    "preprocess_train = transforms.Compose([EegToDevice(device=device)])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    print('- Before -')\n",
    "    pprint.pprint(sample_batched)\n",
    "\n",
    "    print()\n",
    "    print('-' * 100)\n",
    "    print()\n",
    "    \n",
    "    preprocess_train(sample_batched)\n",
    "    \n",
    "    print('- After -')\n",
    "    pprint.pprint(sample_batched)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Normalization per signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizePerSignal(eps=1e-08)\n",
      ")\n",
      "- Before -\n",
      "Mean: tensor([[-6.8880, -4.5790, -0.2805, -1.0875,  1.2830, -5.1060, -0.4220,  1.0950,\n",
      "          2.6655,  1.0315, -3.3815,  0.6835,  0.9020,  4.0190,  0.7710,  1.1295,\n",
      "         -0.5065,  0.9630,  1.2840,  0.5850],\n",
      "        [ 1.3710,  0.9205,  0.1395, -0.2905, -0.7080,  5.2385,  1.1425, -0.0530,\n",
      "         -2.1810, -0.1820,  1.0195, -0.3765, -0.8045,  1.4955, -0.8655, -0.5900,\n",
      "          1.6540,  0.3090, -0.6355,  0.1305]])\n",
      "\n",
      "Std: tensor([[ 68.6684,  27.2979,   6.1905,   9.9731,  14.1283,  59.1413,  15.8267,\n",
      "           7.1583,  10.5741,  14.3079,  25.5349,   8.8883,  14.1888,  13.6482,\n",
      "          13.8810,  15.3526,  12.5770,   6.4207,  10.6819, 121.9661],\n",
      "        [ 31.4508,  15.2986,   8.7377,  10.5960,  23.9189,  31.0532,  12.5155,\n",
      "           8.9206,  10.3289,  27.6918,  11.0976,   7.0854,  13.5130,  10.4345,\n",
      "           8.0944,  16.5765,  14.7400,  11.9985,  10.2284,  32.6769]])\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "- After -\n",
      "Mean: tensor([[ 7.6294e-09,  1.9073e-09, -7.6294e-09, -3.8147e-09,  1.5259e-08,\n",
      "          1.9073e-09,  1.0490e-08, -1.5259e-08, -3.6240e-08,  0.0000e+00,\n",
      "          7.6294e-09, -2.1458e-09, -1.5259e-08, -2.6703e-08,  1.3351e-08,\n",
      "         -7.6294e-09,  1.1444e-08, -5.7220e-09, -9.5367e-09, -1.9073e-09],\n",
      "        [-1.9073e-09,  0.0000e+00,  0.0000e+00, -1.5259e-08,  9.5367e-10,\n",
      "         -4.7684e-09,  7.6294e-09,  9.5367e-09, -1.5259e-08, -1.5736e-08,\n",
      "          7.6294e-09,  3.8147e-09,  5.7220e-09,  1.9073e-09,  1.3351e-08,\n",
      "          1.5259e-08,  7.6294e-09,  0.0000e+00,  1.9073e-09, -6.6757e-09]],\n",
      "       device='cuda:0')\n",
      "\n",
      "Std: tensor([[1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,\n",
      "         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,\n",
      "         1.0000, 1.0000],\n",
      "        [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,\n",
      "         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,\n",
      "         1.0000, 1.0000]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizePerSignal()\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    print('- Before -')\n",
    "    print('Mean:', torch.mean(sample_batched['signal'], axis=-1))\n",
    "    print()\n",
    "    print('Std:', torch.std(sample_batched['signal'], axis=-1))\n",
    "\n",
    "    print()\n",
    "    print('-' * 100)\n",
    "    print()\n",
    "    \n",
    "    preprocess_train(sample_batched)\n",
    "    \n",
    "    print('- After -')\n",
    "    print('Mean:', torch.mean(sample_batched['signal'], axis=-1))\n",
    "    print()\n",
    "    print('Std:', torch.std(sample_batched['signal'], axis=-1))\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Signal normalization using the specified mean and std values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean and standard deviation for signal:\n",
      "tensor([[[ 0.0233],\n",
      "         [-0.0022],\n",
      "         [-0.0132],\n",
      "         [ 0.0054],\n",
      "         [-0.0396],\n",
      "         [ 0.0262],\n",
      "         [-0.0222],\n",
      "         [-0.0088],\n",
      "         [ 0.0500],\n",
      "         [ 0.0818],\n",
      "         [-0.0316],\n",
      "         [ 0.0135],\n",
      "         [-0.0011],\n",
      "         [-0.0490],\n",
      "         [ 0.0508],\n",
      "         [-0.0007],\n",
      "         [ 0.0144],\n",
      "         [ 0.0059],\n",
      "         [-0.0014],\n",
      "         [ 0.0075]]])\n",
      "-\n",
      "tensor([[[45.4992],\n",
      "         [20.4970],\n",
      "         [11.7411],\n",
      "         [11.8701],\n",
      "         [15.6259],\n",
      "         [48.7971],\n",
      "         [19.9753],\n",
      "         [10.6373],\n",
      "         [11.8386],\n",
      "         [15.8078],\n",
      "         [20.8406],\n",
      "         [14.4188],\n",
      "         [13.7448],\n",
      "         [21.7461],\n",
      "         [16.8655],\n",
      "         [14.9501],\n",
      "         [19.5106],\n",
      "         [11.4795],\n",
      "         [11.7924],\n",
      "         [94.0883]]])\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizeMeanStd(mean=tensor([ 0.0233, -0.0022, -0.0132,  0.0054, -0.0396,  0.0262, -0.0222, -0.0088,\n",
      "           0.0500,  0.0818, -0.0316,  0.0135, -0.0011, -0.0490,  0.0508, -0.0007,\n",
      "           0.0144,  0.0059, -0.0014,  0.0075]),std=tensor([45.4992, 20.4970, 11.7411, 11.8701, 15.6259, 48.7971, 19.9753, 10.6373,\n",
      "          11.8386, 15.8078, 20.8406, 14.4188, 13.7448, 21.7461, 16.8655, 14.9501,\n",
      "          19.5106, 11.4795, 11.7924, 94.0883]),eps=1e-08)\n",
      ")\n",
      "- Before -\n",
      "Mean: tensor([[-0.2425,  1.2470, -0.6410, -0.3530, -0.2385, -2.3085, -2.4630,  0.2280,\n",
      "         -1.5865,  0.0805,  0.9285,  6.2535, -0.9320, -1.1390, -0.1945,  0.0460,\n",
      "         -0.2430, -0.6115,  0.3985, -0.2485],\n",
      "        [-0.9215, -0.1580,  2.9025,  0.4720,  1.0350, -4.3200, -1.4030,  1.9075,\n",
      "         -2.7805,  1.1785, -0.9755,  0.9235,  0.1010, -0.4610,  1.7075,  1.5150,\n",
      "          1.2470,  0.7300, -0.1080, -0.0125]])\n",
      "\n",
      "Std: tensor([[ 14.1007,  10.3213,   6.3419,   6.0213,   6.9767,   9.3295,   6.3092,\n",
      "           5.0735,   7.1884,   7.4094,  11.7152,  63.9342,   7.5332,  14.7221,\n",
      "           7.5696,   5.7996,   6.0228,   7.3070,   5.4745,  75.8583],\n",
      "        [ 15.1521,  11.1447,  19.3005,  14.4879,  10.6002,  13.5963,  21.5465,\n",
      "          14.2015,  19.5391,  10.8365,  12.6723,   9.8050,   9.1202,  10.2360,\n",
      "          21.8532,   9.7554,  12.8220,  16.0603,   8.2348, 127.1310]])\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "- After -\n",
      "Mean: tensor([[-5.8424e-03,  6.0943e-02, -5.3472e-02, -3.0191e-02, -1.2731e-02,\n",
      "         -4.7846e-02, -1.2219e-01,  2.2263e-02, -1.3824e-01, -8.0515e-05,\n",
      "          4.6070e-02,  4.3276e-01, -6.7730e-02, -5.0123e-02, -1.4542e-02,\n",
      "          3.1233e-03, -1.3191e-02, -5.3785e-02,  3.3916e-02, -2.7210e-03],\n",
      "        [-2.0766e-02, -7.6034e-03,  2.4833e-01,  3.9311e-02,  6.8768e-02,\n",
      "         -8.9067e-02, -6.9125e-02,  1.8015e-01, -2.3909e-01,  6.9379e-02,\n",
      "         -4.5290e-02,  6.3109e-02,  7.4264e-03, -1.8945e-02,  9.8232e-02,\n",
      "          1.0138e-01,  6.3177e-02,  6.3075e-02, -9.0360e-03, -2.1274e-04]],\n",
      "       device='cuda:0')\n",
      "\n",
      "Std: tensor([[0.3099, 0.5036, 0.5401, 0.5073, 0.4465, 0.1912, 0.3158, 0.4770, 0.6072,\n",
      "         0.4687, 0.5621, 4.4341, 0.5481, 0.6770, 0.4488, 0.3879, 0.3087, 0.6365,\n",
      "         0.4642, 0.8062],\n",
      "        [0.3330, 0.5437, 1.6438, 1.2205, 0.6784, 0.2786, 1.0787, 1.3351, 1.6505,\n",
      "         0.6855, 0.6081, 0.6800, 0.6635, 0.4707, 1.2957, 0.6525, 0.6572, 1.3990,\n",
      "         0.6983, 1.3512]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "signal_mean, signal_std = calculate_signal_statistics(train_loader, repeats=30, verbose=True)\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizeMeanStd(mean=signal_mean, std=signal_std)\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    print('- Before -')\n",
    "    print('Mean:', torch.mean(sample_batched['signal'], axis=-1))\n",
    "    print()\n",
    "    print('Std:', torch.std(sample_batched['signal'], axis=-1))\n",
    "\n",
    "    print()\n",
    "    print('-' * 100)\n",
    "    print()\n",
    "    \n",
    "    preprocess_train(sample_batched)\n",
    "    \n",
    "    print('- After -')\n",
    "    print('Mean:', torch.mean(sample_batched['signal'], axis=-1))\n",
    "    print()\n",
    "    print('Std:', torch.std(sample_batched['signal'], axis=-1))\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Age normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Age mean and standard deviation:\n",
      "tensor([71.2453]) tensor([7.7365])\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizeAge(mean=tensor([71.2453]),std=tensor([7.7365]),eps=1e-08)\n",
      ")\n",
      "- Before -\n",
      "tensor([61., 69.])\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "- After -\n",
      "tensor([-1.3243, -0.2902], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "age_mean, age_std = calculate_age_statistics(train_loader, verbose=True)\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizeAge(mean=age_mean, std=age_std)\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    print('- Before -')\n",
    "    pprint.pprint(sample_batched['age'])\n",
    "\n",
    "    print()\n",
    "    print('-' * 100)\n",
    "    print()\n",
    "    \n",
    "    preprocess_train(sample_batched)\n",
    "    \n",
    "    print('- After -')\n",
    "    pprint.pprint(sample_batched['age'])\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Short time Fourier transform (STFT or spectrogram)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegSpectrogram(n_fft=200, complex_mode=as_real, stft_kwargs={})\n",
      ")\n",
      "- Before -\n",
      "torch.Size([2, 20, 2000])\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "- After -\n",
      "torch.Size([2, 40, 101, 41])\n"
     ]
    }
   ],
   "source": [
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegSpectrogram(n_fft=200, complex_mode='as_real')\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    print('- Before -')\n",
    "    pprint.pprint(sample_batched['signal'].shape)\n",
    "\n",
    "    print()\n",
    "    print('-' * 100)\n",
    "    print()\n",
    "    \n",
    "    preprocess_train(sample_batched)\n",
    "    \n",
    "    print('- After -')\n",
    "    pprint.pprint(sample_batched['signal'].shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "tags": []
   },
   "source": [
    "---\n",
    "\n",
    "## Speed check without STFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "crop_length = 200 * 10\n",
    "batch_size = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### `EDF`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compose(\n",
      "    EegRandomCrop(crop_length=2000, length_limit=10000000, multiple=1, latency=2000, return_timing=False)\n",
      "    EegDropChannels(drop_index=20)\n",
      "    EegToTensor()\n",
      ")\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizeAge(mean=tensor([71.2453]),std=tensor([7.7365]),eps=1e-08)\n",
      "  (2): EegNormalizeMeanStd(mean=tensor([ 0.0233, -0.0022, -0.0132,  0.0054, -0.0396,  0.0262, -0.0222, -0.0088,\n",
      "           0.0500,  0.0818, -0.0316,  0.0135, -0.0011, -0.0490,  0.0508, -0.0007,\n",
      "           0.0144,  0.0059, -0.0014,  0.0075]),std=tensor([45.4992, 20.4970, 11.7411, 11.8701, 15.6259, 48.7971, 19.9753, 10.6373,\n",
      "          11.8386, 15.8078, 20.8406, 14.4188, 13.7448, 21.7461, 16.8655, 14.9501,\n",
      "          19.5106, 11.4795, 11.7924, 94.0883]),eps=1e-08)\n",
      ")\n",
      "CPU times: total: 18min 44s\n",
      "Wall time: 2min 9s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=crop_length,\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "pprint.pprint(transform)\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='edf',\n",
    "                                                                                  transform=transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size=batch_size,\n",
    "                          shuffle=True,\n",
    "                          drop_last=True,\n",
    "                          num_workers=num_workers,\n",
    "                          pin_memory=pin_memory,\n",
    "                          collate_fn=eeg_collate_fn)\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizeAge(mean=age_mean, std=age_std), \n",
    "    EegNormalizeMeanStd(mean=signal_mean, std=signal_std)\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    preprocess_train(sample_batched)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### `PyArrow.feather`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compose(\n",
      "    EegRandomCrop(crop_length=2000, length_limit=10000000, multiple=1, latency=2000, return_timing=False)\n",
      "    EegDropChannels(drop_index=20)\n",
      "    EegToTensor()\n",
      ")\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizeAge(mean=tensor([71.2453]),std=tensor([7.7365]),eps=1e-08)\n",
      "  (2): EegNormalizeMeanStd(mean=tensor([ 0.0233, -0.0022, -0.0132,  0.0054, -0.0396,  0.0262, -0.0222, -0.0088,\n",
      "           0.0500,  0.0818, -0.0316,  0.0135, -0.0011, -0.0490,  0.0508, -0.0007,\n",
      "           0.0144,  0.0059, -0.0014,  0.0075]),std=tensor([45.4992, 20.4970, 11.7411, 11.8701, 15.6259, 48.7971, 19.9753, 10.6373,\n",
      "          11.8386, 15.8078, 20.8406, 14.4188, 13.7448, 21.7461, 16.8655, 14.9501,\n",
      "          19.5106, 11.4795, 11.7924, 94.0883]),eps=1e-08)\n",
      ")\n",
      "CPU times: total: 4.59 s\n",
      "Wall time: 3.66 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=crop_length,\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "pprint.pprint(transform)\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=transform)\n",
    " \n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size=batch_size,\n",
    "                          shuffle=True,\n",
    "                          drop_last=True,\n",
    "                          num_workers=num_workers,\n",
    "                          pin_memory=pin_memory,\n",
    "                          collate_fn=eeg_collate_fn)\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizeAge(mean=age_mean, std=age_std), \n",
    "    EegNormalizeMeanStd(mean=signal_mean, std=signal_std)\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    preprocess_train(sample_batched)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "tags": []
   },
   "source": [
    "---\n",
    "\n",
    "## Speed check with STFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input sequence length: (3000) would become (78, 77) after the STFT with n_fft (155) and hop_length (39).\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "crop_length = 300 * 10\n",
    "n_fft, hop_length, seq_len_2d = calculate_stft_params(seq_length=crop_length, verbose=True)\n",
    "batch_size = 128\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegSpectrogram(n_fft=n_fft, hop_length=hop_length, complex_mode='as_real')\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "signal_2d_mean, signal_2d_std = calculate_signal_statistics(train_loader, preprocess_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### `EDF`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compose(\n",
      "    EegRandomCrop(crop_length=3000, length_limit=10000000, multiple=1, latency=2000, return_timing=False)\n",
      "    EegDropChannels(drop_index=20)\n",
      "    EegToTensor()\n",
      ")\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizeAge(mean=tensor([71.2453]),std=tensor([7.7365]),eps=1e-08)\n",
      "  (2): EegSpectrogram(n_fft=155, complex_mode=as_real, stft_kwargs={'hop_length': 39})\n",
      "  (3): EegNormalizeMeanStd(mean=tensor([[ 2.2792e+01,  3.1216e-01,  1.0842e-01,  ...,  1.3747e-01,\n",
      "            1.2862e-01,  1.2759e-01],\n",
      "          [ 2.7824e+01, -4.2494e-02,  5.1770e-02,  ...,  3.1540e-02,\n",
      "            2.3955e-02,  2.7768e-02],\n",
      "          [ 7.3945e+00,  1.5329e-01, -6.2711e-02,  ..., -2.3968e-02,\n",
      "            7.5325e-03, -9.5869e-04],\n",
      "          ...,\n",
      "          [ 0.0000e+00,  4.6296e-01,  3.1164e-01,  ...,  1.0917e-02,\n",
      "           -1.0444e-02, -3.8881e-04],\n",
      "          [ 0.0000e+00,  5.7710e-01,  3.2173e-01,  ..., -2.2850e-03,\n",
      "            1.9137e-03, -5.3399e-04],\n",
      "          [ 0.0000e+00,  9.0587e-01,  7.9643e-01,  ...,  6.6548e-03,\n",
      "           -8.7216e-04,  3.4029e-03]], device='cuda:0'),std=tensor([[5.7103e+03, 9.8515e+02, 4.9278e+02,  ..., 2.5278e+01, 2.5273e+01,\n",
      "           2.5285e+01],\n",
      "          [2.5968e+03, 3.9856e+02, 2.0727e+02,  ..., 1.2054e+01, 1.2019e+01,\n",
      "           1.2043e+01],\n",
      "          [1.4139e+03, 1.9988e+02, 1.1470e+02,  ..., 7.3799e+00, 7.3298e+00,\n",
      "           7.3657e+00],\n",
      "          ...,\n",
      "          [0.0000e+00, 3.1391e+02, 1.8028e+02,  ..., 2.5639e+00, 2.5912e+00,\n",
      "           2.6011e+00],\n",
      "          [0.0000e+00, 3.3504e+02, 1.8539e+02,  ..., 2.5677e+00, 2.5917e+00,\n",
      "           2.6231e+00],\n",
      "          [0.0000e+00, 1.9271e+03, 2.0630e+03,  ..., 5.7418e+00, 5.4433e+00,\n",
      "           4.6942e+00]], device='cuda:0'),eps=1e-08)\n",
      ")\n",
      "torch.Size([128, 40, 78, 77])\n",
      "CPU times: total: 18min 7s\n",
      "Wall time: 2min 4s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=crop_length,\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "pprint.pprint(transform)\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='edf',\n",
    "                                                                                  transform=transform)\n",
    " \n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size=batch_size,\n",
    "                          shuffle=True,\n",
    "                          drop_last=True,\n",
    "                          num_workers=num_workers,\n",
    "                          pin_memory=pin_memory,\n",
    "                          collate_fn=eeg_collate_fn)\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizeAge(mean=age_mean, std=age_std), \n",
    "    EegSpectrogram(n_fft=n_fft, hop_length=hop_length, complex_mode='as_real'),\n",
    "    EegNormalizeMeanStd(mean=signal_2d_mean, std=signal_2d_std),\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    preprocess_train(sample_batched)\n",
    "    size = sample_batched['signal'].size()\n",
    "    \n",
    "print(size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### `PyArrow.feather`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compose(\n",
      "    EegRandomCrop(crop_length=3000, length_limit=10000000, multiple=1, latency=2000, return_timing=False)\n",
      "    EegDropChannels(drop_index=20)\n",
      "    EegToTensor()\n",
      ")\n",
      "Sequential(\n",
      "  (0): EegToDevice(device=cuda)\n",
      "  (1): EegNormalizeAge(mean=tensor([71.2453]),std=tensor([7.7365]),eps=1e-08)\n",
      "  (2): EegSpectrogram(n_fft=155, complex_mode=as_real, stft_kwargs={'hop_length': 39})\n",
      "  (3): EegNormalizeMeanStd(mean=tensor([[ 2.2792e+01,  3.1216e-01,  1.0842e-01,  ...,  1.3747e-01,\n",
      "            1.2862e-01,  1.2759e-01],\n",
      "          [ 2.7824e+01, -4.2494e-02,  5.1770e-02,  ...,  3.1540e-02,\n",
      "            2.3955e-02,  2.7768e-02],\n",
      "          [ 7.3945e+00,  1.5329e-01, -6.2711e-02,  ..., -2.3968e-02,\n",
      "            7.5325e-03, -9.5869e-04],\n",
      "          ...,\n",
      "          [ 0.0000e+00,  4.6296e-01,  3.1164e-01,  ...,  1.0917e-02,\n",
      "           -1.0444e-02, -3.8881e-04],\n",
      "          [ 0.0000e+00,  5.7710e-01,  3.2173e-01,  ..., -2.2850e-03,\n",
      "            1.9137e-03, -5.3399e-04],\n",
      "          [ 0.0000e+00,  9.0587e-01,  7.9643e-01,  ...,  6.6548e-03,\n",
      "           -8.7216e-04,  3.4029e-03]], device='cuda:0'),std=tensor([[5.7103e+03, 9.8515e+02, 4.9278e+02,  ..., 2.5278e+01, 2.5273e+01,\n",
      "           2.5285e+01],\n",
      "          [2.5968e+03, 3.9856e+02, 2.0727e+02,  ..., 1.2054e+01, 1.2019e+01,\n",
      "           1.2043e+01],\n",
      "          [1.4139e+03, 1.9988e+02, 1.1470e+02,  ..., 7.3799e+00, 7.3298e+00,\n",
      "           7.3657e+00],\n",
      "          ...,\n",
      "          [0.0000e+00, 3.1391e+02, 1.8028e+02,  ..., 2.5639e+00, 2.5912e+00,\n",
      "           2.6011e+00],\n",
      "          [0.0000e+00, 3.3504e+02, 1.8539e+02,  ..., 2.5677e+00, 2.5917e+00,\n",
      "           2.6231e+00],\n",
      "          [0.0000e+00, 1.9271e+03, 2.0630e+03,  ..., 5.7418e+00, 5.4433e+00,\n",
      "           4.6942e+00]], device='cuda:0'),eps=1e-08)\n",
      ")\n",
      "torch.Size([128, 40, 78, 77])\n",
      "CPU times: total: 5.47 s\n",
      "Wall time: 3.65 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    EegRandomCrop(crop_length=crop_length,\n",
    "                  latency=200*10),          # latency: 10s\n",
    "    EegDropChannels(channel_photic), \n",
    "    EegToTensor()\n",
    "])\n",
    "pprint.pprint(transform)\n",
    "\n",
    "config_data, train_dataset, val_dataset, test_dataset = load_caueeg_task_datasets(dataset_path=data_path, \n",
    "                                                                                  task='dementia',\n",
    "                                                                                  load_event=False, \n",
    "                                                                                  file_format='feather',\n",
    "                                                                                  transform=transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size=batch_size,\n",
    "                          shuffle=True,\n",
    "                          drop_last=True,\n",
    "                          num_workers=num_workers,\n",
    "                          pin_memory=pin_memory,\n",
    "                          collate_fn=eeg_collate_fn)\n",
    "\n",
    "preprocess_train = transforms.Compose([\n",
    "    EegToDevice(device=device), \n",
    "    EegNormalizeAge(mean=age_mean, std=age_std), \n",
    "    EegSpectrogram(n_fft=n_fft, hop_length=hop_length, complex_mode='as_real'),\n",
    "    EegNormalizeMeanStd(mean=signal_2d_mean, std=signal_2d_std),\n",
    "])\n",
    "preprocess_train = torch.nn.Sequential(*preprocess_train.transforms).to(device)\n",
    "pprint.pprint(preprocess_train)\n",
    "\n",
    "for i_batch, sample_batched in enumerate(train_loader):\n",
    "    preprocess_train(sample_batched)\n",
    "    size = sample_batched['signal'].size()\n",
    "    \n",
    "print(size)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "37e0b27a9710fa201e26169906fbd4cae3a7652e9e1c64cf2110e48663b0b7d0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}