csc421_project_dataset_generation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "csc421_project_dataset_generation.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMDiV1ksQR0WZw9nIIy7w4S",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Locrian24/csc421-project-stress-classification/blob/main/csc421_project_dataset_generation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Dataset Generation for Stress Detection with CNN\n"
      ],
      "metadata": {
        "id": "kw3GLtK4NcP8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Retrieving dataset"
      ],
      "metadata": {
        "id": "Mf_fXz5DYwPI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -O wesad.zip https://uni-siegen.sciebo.de/s/HGdUkoNlW1Ub0Gx/download\n",
        "!unzip wesad.zip"
      ],
      "metadata": {
        "id": "46khBagaN19f",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6ee33b28-908e-4d3a-f2f5-489cc3d834f1"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2022-04-09 02:14:10--  https://uni-siegen.sciebo.de/s/HGdUkoNlW1Ub0Gx/download\n",
            "Resolving uni-siegen.sciebo.de (uni-siegen.sciebo.de)... 128.176.1.2\n",
            "Connecting to uni-siegen.sciebo.de (uni-siegen.sciebo.de)|128.176.1.2|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2249444501 (2.1G) [application/zip]\n",
            "Saving to: ‘wesad.zip’\n",
            "\n",
            "wesad.zip           100%[===================>]   2.09G  19.9MB/s    in 1m 47s  \n",
            "\n",
            "2022-04-09 02:15:57 (20.1 MB/s) - ‘wesad.zip’ saved [2249444501/2249444501]\n",
            "\n",
            "Archive:  wesad.zip\n",
            "   creating: WESAD/\n",
            "   creating: WESAD/S10/\n",
            "  inflating: WESAD/S10/S10.pkl       \n",
            "  inflating: WESAD/S10/S10_E4_Data.zip  \n",
            "  inflating: WESAD/S10/S10_quest.csv  \n",
            "  inflating: WESAD/S10/S10_readme.txt  \n",
            "  inflating: WESAD/S10/S10_respiban.txt  \n",
            "   creating: WESAD/S11/\n",
            "  inflating: WESAD/S11/S11.pkl       \n",
            "  inflating: WESAD/S11/S11_E4_Data.zip  \n",
            "  inflating: WESAD/S11/S11_quest.csv  \n",
            "  inflating: WESAD/S11/S11_readme.txt  \n",
            "  inflating: WESAD/S11/S11_respiban.txt  \n",
            "   creating: WESAD/S13/\n",
            "  inflating: WESAD/S13/S13.pkl       \n",
            "  inflating: WESAD/S13/S13_E4_Data.zip  \n",
            "  inflating: WESAD/S13/S13_quest.csv  \n",
            "  inflating: WESAD/S13/S13_readme.txt  \n",
            "  inflating: WESAD/S13/S13_respiban.txt  \n",
            "   creating: WESAD/S14/\n",
            "  inflating: WESAD/S14/S14.pkl       \n",
            "  inflating: WESAD/S14/S14_E4_Data.zip  \n",
            "  inflating: WESAD/S14/S14_quest.csv  \n",
            "  inflating: WESAD/S14/S14_readme.txt  \n",
            "  inflating: WESAD/S14/S14_respiban.txt  \n",
            "   creating: WESAD/S15/\n",
            "  inflating: WESAD/S15/S15.pkl       \n",
            "  inflating: WESAD/S15/S15_E4_Data.zip  \n",
            "  inflating: WESAD/S15/S15_quest.csv  \n",
            "  inflating: WESAD/S15/S15_readme.txt  \n",
            "  inflating: WESAD/S15/S15_respiban.txt  \n",
            "   creating: WESAD/S16/\n",
            "  inflating: WESAD/S16/S16.pkl       \n",
            "  inflating: WESAD/S16/S16_E4_Data.zip  \n",
            "  inflating: WESAD/S16/S16_quest.csv  \n",
            "  inflating: WESAD/S16/S16_readme.txt  \n",
            "  inflating: WESAD/S16/S16_respiban.txt  \n",
            "   creating: WESAD/S17/\n",
            "  inflating: WESAD/S17/S17.pkl       \n",
            "  inflating: WESAD/S17/S17_E4_Data.zip  \n",
            "  inflating: WESAD/S17/S17_quest.csv  \n",
            "  inflating: WESAD/S17/S17_readme.txt  \n",
            "  inflating: WESAD/S17/S17_respiban.txt  \n",
            "   creating: WESAD/S2/\n",
            "  inflating: WESAD/S2/S2.pkl         \n",
            "  inflating: WESAD/S2/S2_E4_Data.zip  \n",
            "  inflating: WESAD/S2/S2_quest.csv   \n",
            "  inflating: WESAD/S2/S2_readme.txt  \n",
            "  inflating: WESAD/S2/S2_respiban.txt  \n",
            "   creating: WESAD/S3/\n",
            "  inflating: WESAD/S3/S3.pkl         \n",
            "  inflating: WESAD/S3/S3_E4_Data.zip  \n",
            "  inflating: WESAD/S3/S3_quest.csv   \n",
            "  inflating: WESAD/S3/S3_readme.txt  \n",
            "  inflating: WESAD/S3/S3_respiban.txt  \n",
            "   creating: WESAD/S4/\n",
            "  inflating: WESAD/S4/S4.pkl         \n",
            "  inflating: WESAD/S4/S4_E4_Data.zip  \n",
            "  inflating: WESAD/S4/S4_quest.csv   \n",
            "  inflating: WESAD/S4/S4_readme.txt  \n",
            "  inflating: WESAD/S4/S4_respiban.txt  \n",
            "   creating: WESAD/S5/\n",
            "  inflating: WESAD/S5/S5.pkl         \n",
            "  inflating: WESAD/S5/S5_E4_Data.zip  \n",
            "  inflating: WESAD/S5/S5_quest.csv   \n",
            "  inflating: WESAD/S5/S5_readme.txt  \n",
            "  inflating: WESAD/S5/S5_respiban.txt  \n",
            "   creating: WESAD/S6/\n",
            "  inflating: WESAD/S6/S6.pkl         \n",
            "  inflating: WESAD/S6/S6_E4_Data.zip  \n",
            "  inflating: WESAD/S6/S6_quest.csv   \n",
            "  inflating: WESAD/S6/S6_readme.txt  \n",
            "  inflating: WESAD/S6/S6_respiban.txt  \n",
            "   creating: WESAD/S7/\n",
            "  inflating: WESAD/S7/S7.pkl         \n",
            "  inflating: WESAD/S7/S7_E4_Data.zip  \n",
            "  inflating: WESAD/S7/S7_quest.csv   \n",
            "  inflating: WESAD/S7/S7_readme.txt  \n",
            "  inflating: WESAD/S7/S7_respiban.txt  \n",
            "   creating: WESAD/S8/\n",
            "  inflating: WESAD/S8/S8.pkl         \n",
            "  inflating: WESAD/S8/S8_E4_Data.zip  \n",
            "  inflating: WESAD/S8/S8_quest.csv   \n",
            "  inflating: WESAD/S8/S8_readme.txt  \n",
            "  inflating: WESAD/S8/S8_respiban.txt  \n",
            "   creating: WESAD/S9/\n",
            "  inflating: WESAD/S9/S9.pkl         \n",
            "  inflating: WESAD/S9/S9_E4_Data.zip  \n",
            "  inflating: WESAD/S9/S9_quest.csv   \n",
            "  inflating: WESAD/S9/S9_readme.txt  \n",
            "  inflating: WESAD/S9/S9_respiban.txt  \n",
            "  inflating: WESAD/wesad_readme.pdf  \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!pip install webdataset"
      ],
      "metadata": {
        "id": "4voPCtOX9eH0"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# RespiBAN used for CNN data, EmpaticaE4 used for MLP-NN"
      ],
      "metadata": {
        "id": "5O3r4O45TEAh"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import webdataset as wds\n",
        "import os\n",
        "from sklearn.model_selection import train_test_split"
      ],
      "metadata": {
        "id": "_ru0UFNNU-h1"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from scipy import stats\n",
        "\n",
        "def generate_sample(data, i, ignore_labels=[]):\n",
        "  keys = ['ECG', 'EDA', 'EMG', 'Temp', 'Resp', 'ACC-x', 'ACC-y', 'ACC-z']\n",
        "\n",
        "  idx = i * 3500\n",
        "  sample = np.ndarray((8, 3500))\n",
        "  for j, key in enumerate(keys):\n",
        "    # Accelerometer data at the end\n",
        "    if key in ['ACC-x', 'ACC-y', 'ACC-z']:\n",
        "      sample[j] = data['signal']['chest']['ACC'][idx:idx+3500, 0]\n",
        "      sample[j+1] = data['signal']['chest']['ACC'][idx:idx+3500, 1]\n",
        "      sample[j+2] = data['signal']['chest']['ACC'][idx:idx+3500, 2]\n",
        "      break\n",
        "\n",
        "    sample[j] = np.squeeze(data['signal']['chest'][key][idx:idx+3500])\n",
        "\n",
        "  filtered = data['label'][idx:idx+3500][np.where(np.isin(data['label'][idx:idx+3500], ignore_labels, invert=True))]\n",
        "  \n",
        "  if len(filtered) == 0:\n",
        "    label = 999\n",
        "  else:\n",
        "    label = stats.mode(filtered)[0][0]\n",
        "\n",
        "  return sample, label"
      ],
      "metadata": {
        "id": "cvHKStspAEXn"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate binary stress classification dataset\n",
        "\n",
        "all_samples = []\n",
        "labels = []\n",
        "ignore_labels = [5,6,7] # Used for binary classification\n",
        "\n",
        "dirs = next(os.walk('./WESAD'))[1]\n",
        "for dir in dirs:\n",
        "  data = pd.read_pickle(f'./WESAD/{dir}/{dir}.pkl')\n",
        "  num_samples = data['label'].shape[0] // 3500\n",
        "\n",
        "  for i in range(num_samples):\n",
        "    sample, label = generate_sample(data, i, ignore_labels)\n",
        "    \n",
        "    if label in ignore_labels or label == 999: # Ignore meaningless samples\n",
        "      continue\n",
        "\n",
        "    labels.append(label)\n",
        "    all_samples.append(sample)\n",
        "\n",
        "train_X, test_X, train_y, test_y = train_test_split(all_samples, labels, test_size=.3, train_size=.7, shuffle=True, stratify=labels)\n",
        "\n",
        "# Write WebDataset tar files\n",
        "sink_train = wds.TarWriter(\"WESAD_RespiBAN_binary_train.tar\")\n",
        "\n",
        "for i in range(len(train_y)):\n",
        "  sink_train.write({\n",
        "      \"__key__\": \"sample%06d\" % i,\n",
        "      \"input.npy\": train_X[i],\n",
        "      \"label.cls\": train_y[i]\n",
        "  })\n",
        "\n",
        "sink_train.close()\n",
        "\n",
        "sink_test = wds.TarWriter(\"WESAD_RespiBAN_binary_test.tar\")\n",
        "for i in range(len(test_y)):\n",
        "  sink_test.write({\n",
        "      \"__key__\": \"sample%06d\" % i,\n",
        "      \"input.npy\": test_X[i],\n",
        "      \"label.cls\": test_y[i]\n",
        "  })\n",
        "\n",
        "sink_test.close()"
      ],
      "metadata": {
        "id": "LyxvNsFiSRhr"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate 3-class emotion classification dataset\n",
        "\n",
        "all_samples = []\n",
        "labels = []\n",
        "ignore_labels = [0,4,5,6,7] # Only 1, 2, 3 useful for emotion classification\n",
        "\n",
        "dirs = next(os.walk('./WESAD'))[1]\n",
        "for dir in dirs:\n",
        "  data = pd.read_pickle(f'./WESAD/{dir}/{dir}.pkl')\n",
        "  num_samples = data['label'].shape[0] // 3500\n",
        "\n",
        "  for i in range(num_samples):\n",
        "    sample, label = generate_sample(data, i, ignore_labels)\n",
        "    \n",
        "    if label in ignore_labels or label == 999: # Ignore meaningless samples\n",
        "      continue\n",
        "\n",
        "    labels.append(label)\n",
        "    all_samples.append(sample)\n",
        "\n",
        "train_X, test_X, train_y, test_y = train_test_split(all_samples, labels, test_size=.3, train_size=.7, shuffle=True, stratify=labels)\n",
        "\n",
        "# Write WebDataset tar files\n",
        "sink_train = wds.TarWriter(\"WESAD_RespiBAN_multiclass_train.tar\")\n",
        "\n",
        "for i in range(len(train_y)):\n",
        "  sink_train.write({\n",
        "      \"__key__\": \"sample%06d\" % i,\n",
        "      \"input.npy\": train_X[i],\n",
        "      \"label.cls\": train_y[i]\n",
        "  })\n",
        "\n",
        "sink_train.close()\n",
        "\n",
        "sink_test = wds.TarWriter(\"WESAD_RespiBAN_multiclass_test.tar\")\n",
        "for i in range(len(test_y)):\n",
        "  sink_test.write({\n",
        "      \"__key__\": \"sample%06d\" % i,\n",
        "      \"input.npy\": test_X[i],\n",
        "      \"label.cls\": test_y[i]\n",
        "  })\n",
        "\n",
        "sink_test.close()"
      ],
      "metadata": {
        "id": "ck8foI97A3-3"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the following to copy the dataset to your Google Drive if needed"
      ],
      "metadata": {
        "id": "wkkRWtvaIqN9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eGJi6fT_FwnV",
        "outputId": "ac2352f4-337a-47d9-c88d-495a20ff67f8"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cp WESAD_RespiBAN_binary_train.tar /content/drive/MyDrive/\n",
        "!cp WESAD_RespiBAN_binary_test.tar /content/drive/MyDrive/\n",
        "!cp WESAD_RespiBAN_multiclass_train.tar /content/drive/MyDrive/\n",
        "!cp WESAD_RespiBAN_multiclass_test.tar /content/drive/MyDrive/"
      ],
      "metadata": {
        "id": "cZ4IWuX5F6gm"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!rm -r WESAD\n",
        "!rm wesad.zip"
      ],
      "metadata": {
        "id": "3SF-OPQvXZS7"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = (\n",
        "      wds.WebDataset('WESAD_RespiBAN_multiclass_train.tar')\n",
        "        .decode()\n",
        "        .to_tuple(\"input.npy\", \"label.cls\")\n",
        ")"
      ],
      "metadata": {
        "id": "y9kkyVFHFCKE"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "i = 0\n",
        "for input, label in dataset:\n",
        "  if label not in [1, 2, 3]:\n",
        "    i += 1\n",
        "    print(label)\n",
        "\n",
        "i"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O99tqe-yFLo3",
        "outputId": "be588d80-b1d4-43bf-9eba-b7d73ecdfa18"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    }
  ]
}