3 - Hardware Control.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import scipy.signal\n",
    "import numpy as np\n",
    "import threading\n",
    "import time\n",
    "import usb\n",
    "import sys\n",
    "\n",
    "from sklearn.cluster import KMeans\n",
    "\n",
    "import Jetson.GPIO as GPIO\n",
    "\n",
    "from torch.nn import functional as F\n",
    "from torch import nn\n",
    "import torch\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# our Autoencoder!\n",
    "class AE(nn.Module):\n",
    "    def __init__(self, ae_width):\n",
    "        super(AE, self).__init__()\n",
    "\n",
    "        self.fc1 = nn.Linear(1024, 256)\n",
    "        self.fc2 = nn.Linear(256, 64)\n",
    "        self.fc3 = nn.Linear(64, ae_width)\n",
    "        self.fc4 = nn.Linear(ae_width, 64)\n",
    "        self.fc5 = nn.Linear(64, 256)\n",
    "        self.fc6 = nn.Linear(256, 1024)\n",
    "\n",
    "    def encode(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        return self.fc3(x)\n",
    "\n",
    "    def decode(self, x):\n",
    "        x = F.relu(self.fc4(x))\n",
    "        x = F.relu(self.fc5(x))\n",
    "        return self.fc6(x)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.encode(x.view(-1, 1024))\n",
    "        return self.decode(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a trained model\n",
    "model = AE(20)\n",
    "model.load_state_dict(torch.load('Autoencoder_EEG.pt'))\n",
    "model.eval()\n",
    "\n",
    "# Push it onto the GPU\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up our NIA device\n",
    "VENDOR_ID     = 0x1234   #: Vendor Id\n",
    "PRODUCT_ID    = 0x0000   #: Product Id for the bridged usb cable\n",
    "INTERFACE_ID  = 0x81     #: The interface we use to talk to the device\n",
    "PACKET_LENGTH = 0x40     #: 64 bytes\n",
    "\n",
    "class Nia():\n",
    "    \"\"\" Attaches the NIA device, and provides low level data collection\n",
    "    \"\"\" \n",
    "\n",
    "    def __init__(self, seconds=1) :\n",
    "        self.read_length = 1024\n",
    "        self.incoming_data = np.zeros(self.read_length*2, dtype=np.uint32)\n",
    "        self.current_data = np.zeros(self.read_length, dtype=np.uint32)\n",
    "        self.device = usb.core.find(idVendor=VENDOR_ID, idProduct=PRODUCT_ID)\n",
    "        self.device.reset()\n",
    "    \n",
    "    def _bulk_read(self):\n",
    "        \"\"\" Read data off the NIA from its internal buffer, of up to 16 samples\"\"\"\n",
    "        read_bytes = self.device.read(INTERFACE_ID, PACKET_LENGTH, timeout=30)\n",
    "        return read_bytes\n",
    "        \n",
    "    def _get_data(self):\n",
    "        count = 0\n",
    "        while True:\n",
    "            bytes_data = self._bulk_read()\n",
    "            point_count = int(bytes_data[54])\n",
    "\n",
    "            for i in range(point_count):\n",
    "                self.incoming_data[count + i] = int.from_bytes(bytes_data[i*3:i*3+3], byteorder='little')\n",
    "               \n",
    "            count = count + point_count\n",
    "\n",
    "            if count >= self.read_length:\n",
    "                break\n",
    "\n",
    "        self.current_data = self.incoming_data[:self.read_length]\n",
    "        \n",
    "    def update(self):\n",
    "        data_thread = threading.Thread(target=self._get_data)\n",
    "        data_thread.start()\n",
    "        data_thread.join()\n",
    "        \n",
    "def butter_bandpass(lowcut, highcut, fs, order=5):\n",
    "        nyq = 0.5 * fs\n",
    "        low = lowcut / nyq\n",
    "        high = highcut / nyq\n",
    "        sos = scipy.signal.butter(order, [low, high], analog=False, btype='band', output='sos')\n",
    "        return sos\n",
    "\n",
    "def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):\n",
    "        sos = butter_bandpass(lowcut, highcut, fs, order=order)\n",
    "        y = scipy.signal.sosfilt(sos, data)\n",
    "        return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capture data, clean it up and feed it through the pre-trained Autoencoder\n",
    "def process_eeg_data():\n",
    "    # Capture 250 ms of EEG Data, and run it through the Aurtoencoder\n",
    "\n",
    "    nia.update()\n",
    "    dataset_unfiltered = nia.current_data\n",
    "    mean, std = np.mean(dataset_unfiltered), np.std(dataset_unfiltered)\n",
    "    dataset_unfiltered = np.where(dataset_unfiltered > mean + std*10, dataset_unfiltered - 2**16, dataset_unfiltered)\n",
    "    dataset_unfiltered = np.where(dataset_unfiltered < mean - std*10, dataset_unfiltered + 2**16, dataset_unfiltered)\n",
    "    dataset_unfiltered = dataset_unfiltered - mean\n",
    "    dataset_unfiltered = dataset_unfiltered / std\n",
    "\n",
    "    # make sure we are the right data type\n",
    "    dataset_unfiltered = dataset_unfiltered.astype(np.float32)\n",
    "    dataset_filtered = butter_bandpass_filter(dataset_unfiltered, lowcut=5, highcut=35, fs=4096, order=6)\n",
    "    dataset_filtered = dataset_filtered.astype(np.float32)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        autoencoder = model.encode(torch.from_numpy(dataset_filtered).cuda())\n",
    "    \n",
    "    return autoencoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate our NIA class\n",
    "if 'nia' in locals():\n",
    "    nia.device.reset()\n",
    "nia = Nia()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Record 25 seconds of data in State 1\n",
    "state_1 = np.zeros((100, 20))\n",
    "for i in range(100):\n",
    "    data = process_eeg_data()\n",
    "    state_1[i,:] = data.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Record 25 seconds of data in State 2\n",
    "state_2 = np.zeros((100, 20))\n",
    "for i in range(100):\n",
    "    data = process_eeg_data()\n",
    "    state_2[i,:] = data.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine the EEG data, and use it to train the Classifier\n",
    "combined = np.vstack((state_1, state_2))\n",
    "kmeans = KMeans(n_clusters=2, random_state=0).fit(combined)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the GPIO pins on the Jetson nano\n",
    "GPIO.setmode(GPIO.BOARD)\n",
    "GPIO.setup([11, 12], GPIO.OUT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Finally, run the Classifier, and control those LEDs!!!\n",
    "while True:\n",
    "    autoencoder_output = process_eeg_data().cpu().numpy()\n",
    "    state = kmeans.predict(autoencoder_output.reshape(1, -1))\n",
    "\n",
    "    if state == 0:\n",
    "        GPIO.output(11, GPIO.HIGH)\n",
    "        GPIO.output(12, GPIO.LOW)\n",
    "    else:\n",
    "        GPIO.output(12, GPIO.HIGH)\n",
    "        GPIO.output(11, GPIO.LOW)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}