Main_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install wfdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing needed libraries\n",
    "from matplotlib import pyplot as plt\n",
    "from wfdb.io import get_record_list\n",
    "from wfdb import rdsamp, processing\n",
    "import numpy as np\n",
    "import random\n",
    "from scipy.signal import resample_poly\n",
    "from pickle import dump, load\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils import data\n",
    "from sklearn.metrics import recall_score\n",
    "import tqdm\n",
    "\n",
    "import utils\n",
    "from data_generator import dataset_gen\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▍         | 4/84 [01:20<29:41, 22.27s/it]"
     ]
    }
   ],
   "source": [
    "load_flag = True\n",
    "db_name = 'ltafdb'\n",
    "\n",
    "if load_flag:\n",
    "    # import files\n",
    "    records = get_record_list(db_name)\n",
    "    signals, beats, beat_types = utils.data_from_records(records, channel=0, db=db_name)\n",
    "    # save file\n",
    "    #dump(signals, open('./tests db/' + db_name + '_signals.pkl', 'wb'))\n",
    "    #dump(beats, open('./tests db/' + db_name + '_beats.pkl', 'wb'))\n",
    "    #dump(beat_types, open('./tests db/' + db_name + 'beat_types.pkl', 'wb'))\n",
    "else:\n",
    "    signals = load(open('./tests db/' + db_name + '_signals.pkl', 'rb'))\n",
    "    beats = load(open('./tests db/' + db_name + '_beats.pkl', 'rb'))\n",
    "    beat_types = load(open('./tests db/' + db_name + 'beat_types.pkl', 'rb'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Different R-peak types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract beat symbols from all records\n",
    "all_symbols = []\n",
    "for symbols in beat_types:\n",
    "    all_symbols.append(symbols)\n",
    "    \n",
    "all_symbols = [item for sublist in all_symbols for item in sublist]\n",
    "all_symbols = np.asarray(all_symbols)\n",
    "u, c = np.unique(all_symbols, return_counts=True)\n",
    "\n",
    "# Meanings for different heart beat codings\n",
    "label_meanings = {\n",
    "    \"N\": \"Normal beat\",\n",
    "    \"L\": \"Left bundle branch block beat\",\n",
    "    \"R\": \"Right bundle branch block beat\",\n",
    "    \"V\": \"Premature ventricular contraction\",\n",
    "    \"/\": \"Paced beat\",\n",
    "    \"A\": \"Atrial premature beat\",\n",
    "    \"f\": \"Fusion of paced and normal beat\",\n",
    "    \"F\": \"Fusion of ventricular and normal beat\",\n",
    "    \"j\": \"Nodal (junctional) escape beat\",\n",
    "    \"a\": \"Aberrated atrial premature beat\",\n",
    "    \"E\": \"Ventricular escape beat\",\n",
    "    \"J\": \"Nodal (junctional) premature beat\",\n",
    "    \"Q\": \"Unclassifiable beat\",\n",
    "    \"e\": \"Atrial escape beat\",\n",
    "    \"S\": \"Supraventricular premature or ectopic\"\n",
    "}\n",
    "\n",
    "# Print number of instances in each beat type\n",
    "label_counts = [(label, count) for label, count in zip(u.tolist(), c.tolist())]\n",
    "label_counts.sort(key=lambda tup: tup[1], reverse=True)\n",
    "for label in label_counts:\n",
    "    print(label_meanings[label[0]], \"-\"*(40-len(label_meanings[label[0]])), label[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot examples of beat types\n",
    "plt.style.use('ggplot')\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "fig.set_figheight(5), fig.set_figwidth(18)\n",
    "ax1.plot(signals[0][:1000])\n",
    "ax1.set_title('Normal beats')\n",
    "ax2.plot(signals[1][:1000])\n",
    "ax2.set_title('Premature venticular contractions')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting character encodings to numerical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dictionary that encodes characters as numerical\n",
    "label_codings = {}\n",
    "for i in range(0, len(label_counts)):\n",
    "    if label_counts[i][0] == 'N':\n",
    "        label_codings[label_counts[i][0]] = 1\n",
    "    else:\n",
    "        label_codings[label_counts[i][0]] = -1\n",
    "label_codings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save new numerical encodings as list\n",
    "labels = []\n",
    "for beat_types in beat_types:\n",
    "    numerical_symbols = [label_codings[sym] for sym in beat_types]\n",
    "    labels.append(np.asarray(numerical_symbols))\n",
    "    \n",
    "labels[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot signal and beat locations as dashed line\n",
    "plt.figure(figsize=(10,5))\n",
    "peaks = beats[0][0]\n",
    "plt.plot(signals[0][peaks-100:peaks+100])\n",
    "for peak in beats[0][:1]:\n",
    "    plt.axvline(x=100, color='k', linestyle='--', alpha=0.5)\n",
    "plt.title('Normal beats where R-peak location occurs at the bottom of a valley')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = utils.fix_labels(signals, beats, labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data\n",
    "baseline_wander = rdsamp('bw', pn_dir='nstdb')\n",
    "muscle_artifact = rdsamp('ma', pn_dir='nstdb')\n",
    "\n",
    "# Concatenate two channels to make one longer recording\n",
    "ma = np.concatenate((muscle_artifact[0][:,0], muscle_artifact[0][:,1]))\n",
    "bw = np.concatenate((baseline_wander[0][:,0], baseline_wander[0][:,1]))\n",
    "\n",
    "# Resample noise to 250 Hz\n",
    "ma = resample_poly(ma, up=250, down=muscle_artifact[1]['fs'])\n",
    "bw = resample_poly(bw, up=250, down=baseline_wander[1]['fs'])\n",
    "\n",
    "# Plot examples of baseline wandering and muscle artifact noise types\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "fig.set_figheight(7), fig.set_figwidth(18)\n",
    "ax1.plot(bw[:5000])\n",
    "ax1.set_title('Baseline wander')\n",
    "ax2.plot(ma[:5000])\n",
    "ax2.set_title('Muscle artifact')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Training Data\n",
    "n_batch = 84\n",
    "win_size = 1000\n",
    "index_list = list(range(len(signals)))\n",
    "print('Number of examples is %d' % len(index_list))\n",
    "\n",
    "\n",
    "train_dataset = dataset_gen(signals=[signals[i] for i in index_list], \n",
    "                            peaks=[beats[i] for i in index_list], \n",
    "                            labels=[labels[i] for i in index_list], \n",
    "                            ma=None,\n",
    "                            bw=None,\n",
    "                            win_size=win_size,\n",
    "                            add_noise = False)\n",
    "\n",
    "train_loader = data.DataLoader(train_dataset, \n",
    "                               batch_size=n_batch, \n",
    "                               shuffle=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create dataset from loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true = []\n",
    "X = []\n",
    "epochs = 50\n",
    "\n",
    "# turn off gradients for evalute\n",
    "for epoch in tqdm.tqdm(range(epochs)):\n",
    "    with torch.no_grad():\n",
    "        for i,batch in enumerate(train_loader):\n",
    "            # get batch images and labels\n",
    "            X_batch,y_true_batch = batch \n",
    "\n",
    "            # convert to cpu numpy\n",
    "            X_batch = X_batch.cpu().squeeze(2).numpy()\n",
    "            y_true_batch = y_true_batch.cpu().squeeze(2).numpy()\n",
    "\n",
    "            # append true\n",
    "            y_true.append(y_true_batch)\n",
    "            X.append(X_batch)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save file to pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dump(X, open('./tests db/X_' + db_name + '.pkl', 'wb'))\n",
    "dump(y_true, open('./tests db/y_true_' + db_name + '.pkl', 'wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Different Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = ['nsrdb', 'svdb', 'incartdb', 'edb']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_name = \"edb\"\n",
    "X = load(open('./tests db/X_' + db_name + '.pkl', 'rb'))\n",
    "y_true = load(open('./tests db/y_true_' + db_name + '.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load mode;\n",
    "model = torch.load('./transformer_inception_label_correction_model.pt')\n",
    "model.to(device)\n",
    "\n",
    "# empty lists\n",
    "y = []\n",
    "y_pred = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn off gradients for evalute\n",
    "with torch.no_grad():\n",
    "    for X_batch,y_true_batch in tqdm.tqdm(zip(X,y_true)):\n",
    "        # get batch images and labels\n",
    "        X_batch = torch.from_numpy(X_batch).unsqueeze(2).to(device)\n",
    "        y_pred_batch = model(X_batch)\n",
    "        \n",
    "        # save to vectors\n",
    "        y_true_batch = y_true_batch.flatten().tolist()\n",
    "        y_pred_batch = torch.round(y_pred_batch.cpu().flatten()).numpy().tolist()\n",
    "        y.append(y_true_batch)\n",
    "        y_pred.append(y_pred_batch)\n",
    "        \n",
    "# flat arrays\n",
    "y = np.array([item for sublist in y for item in sublist]).astype(int)\n",
    "y_pred = np.array([item for sublist in y_pred for item in sublist]).astype(int)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate recall and specificty\n",
    "recall = recall_score(y, y_pred) * 100\n",
    "specificty = recall_score(y, y_pred, pos_label = 0) * 100\n",
    "print('For gqrs algorithm, Recall is %.3f and Specificity is %.3f' % (recall, specificty))"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.g4dn.xlarge",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.8 Python 3.6 GPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/pytorch-1.8-gpu-py36"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}