train-ssl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "189a4eda-2c0f-4a99-9a9d-2223d2f08668",
   "metadata": {},
   "source": [
    "### End-to-end SSL stepcounter training with OxWalk\n",
    "\n",
    "This notebook trains the stepcounter that was used in:\n",
    "\n",
    "Small SR, Chan S, Walmsley R, et al. (2023)\n",
    "[Development and Validation of a Machine Learning Wrist-worn Step Detection Algorithm with Deployment in the UK Biobank](https://www.medrxiv.org/content/10.1101/2023.02.20.23285750v1).\n",
    "medRxiv. DOI: 10.1101/2023.02.20.23285750"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "479c96a5-c0e3-4592-ac5e-12a21f0349da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import glob\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pathlib\n",
    "import urllib\n",
    "import shutil\n",
    "import zipfile\n",
    "import torch\n",
    "import joblib\n",
    "from tqdm import tqdm\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "from stepcount.models import StepCounter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6301651-cb06-4255-83cf-08482447ff03",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Download OxWalk dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "15f9f0ef-f66a-4475-a161-1744b1ac8f1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading OxWalk...\n",
      "Unzipping...\n"
     ]
    }
   ],
   "source": [
    "print(f\"Downloading OxWalk...\")\n",
    "url = \"https://ora.ox.ac.uk/objects/uuid:19d3cb34-e2b3-4177-91b6-1bad0e0163e7/files/dcj82k7829\"\n",
    "with urllib.request.urlopen(url) as f_src, open(\"OxWalk_Dec2022.zip\", \"wb\") as f_dst:\n",
    "    shutil.copyfileobj(f_src, f_dst)\n",
    "print(\"Unzipping...\")\n",
    "with zipfile.ZipFile(\"OxWalk_Dec2022.zip\", \"r\") as f:\n",
    "    f.extractall(\".\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd1bbebe-7f2b-427b-a080-3edcc71b481c",
   "metadata": {},
   "source": [
    "Definitions and helper functions to process the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "38eebcba-f2b1-4b2b-a4e8-608547e784a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = './OxWalk_Dec2022/Wrist_100Hz'  # location of OxWalk (the .csv files should be in this folder)\n",
    "DATAFILES = os.path.join(DATA_DIR, 'P*.csv')\n",
    "OUT_DIR = './data/oxwalk_30hz_w10_o0/'  # output location of the processed dataset\n",
    "WINDOW_SEC = 10  # seconds\n",
    "STEP_THRESHOLD = 4  # how many steps per window to consider a step (equal or more) \n",
    "\n",
    "DEVICE_HZ = 100  # original sample rate Hz\n",
    "RESAMPLE_HZ = 30  # Hz\n",
    "WINDOW_OVERLAP_SEC = 0  # seconds\n",
    "WINDOW_LEN = int(DEVICE_HZ * WINDOW_SEC)  # device ticks\n",
    "WINDOW_OVERLAP_LEN = int(DEVICE_HZ * WINDOW_OVERLAP_SEC)  # device ticks\n",
    "WINDOW_STEP_LEN = WINDOW_LEN - WINDOW_OVERLAP_LEN  # device ticks\n",
    "WINDOW_TOL = 0.01  # 1%\n",
    "\n",
    "\n",
    "def resize(x, length, axis=1):\n",
    "    \"\"\"Resize the temporal length using linear interpolation.\n",
    "    X must be of shape (N,M,C) (channels last) or (N,C,M) (channels first),\n",
    "    where N is the batch size, M is the temporal length, and C is the number\n",
    "    of channels.\n",
    "    If X is channels-last, use axis=1 (default).\n",
    "    If X is channels-first, use axis=2.\n",
    "    \"\"\"\n",
    "    from scipy.interpolate import interp1d\n",
    "\n",
    "    length_orig = x.shape[axis]\n",
    "    t_orig = np.linspace(0, 1, length_orig, endpoint=True)\n",
    "    t_new = np.linspace(0, 1, length, endpoint=True)\n",
    "    x = interp1d(t_orig, x, kind=\"linear\", axis=axis, assume_sorted=True)(\n",
    "        t_new\n",
    "    )\n",
    "    return x\n",
    "\n",
    "\n",
    "def is_good_quality(w):\n",
    "    \"\"\" Window quality check \"\"\"\n",
    "\n",
    "    if w.isna().any().any():\n",
    "        return False\n",
    "\n",
    "    if len(w) != WINDOW_LEN:\n",
    "        return False\n",
    "\n",
    "    w_start, w_end = w.index[0], w.index[-1]\n",
    "    w_duration = w_end - w_start\n",
    "    target_duration = pd.Timedelta(WINDOW_SEC, 's')\n",
    "    if np.abs(w_duration - target_duration) > WINDOW_TOL * target_duration:\n",
    "        return False\n",
    "\n",
    "    return True\n",
    "\n",
    "\n",
    "def make(datafile):\n",
    "    X, Y, T, P, = [], [], [], []\n",
    "    y_step = []\n",
    "\n",
    "    data = pd.read_csv(datafile, parse_dates=['timestamp'], index_col='timestamp')\n",
    "\n",
    "    p = re.search(r'(P\\d{2})', datafile, flags=re.IGNORECASE).group()\n",
    "\n",
    "    for i in range(0, len(data), WINDOW_STEP_LEN):\n",
    "        w = data.iloc[i:i + WINDOW_LEN]\n",
    "\n",
    "        if not is_good_quality(w):\n",
    "            continue\n",
    "\n",
    "        t = w.index[0].to_datetime64()\n",
    "        x = w[['x', 'y', 'z']].values\n",
    "        count = int(w['annotation'].sum())\n",
    "\n",
    "        if count >= STEP_THRESHOLD:\n",
    "            y = 'walk'\n",
    "        else:\n",
    "            y = 'notwalk'\n",
    "\n",
    "        X.append(x)\n",
    "        Y.append(y)\n",
    "        T.append(t)\n",
    "        P.append(p)\n",
    "        y_step.append(count)\n",
    "\n",
    "    X = np.asarray(X)\n",
    "    Y = np.asarray(Y)\n",
    "    T = np.asarray(T)\n",
    "    P = np.asarray(P)\n",
    "    y_step = np.asarray(y_step)\n",
    "\n",
    "    if DEVICE_HZ != RESAMPLE_HZ:\n",
    "        X = resize(X, int(RESAMPLE_HZ * WINDOW_SEC))\n",
    "\n",
    "    return X, Y, T, P, y_step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e44b4ac5-f4c9-4ed6-8cb5-c666a7209fbb",
   "metadata": {},
   "source": [
    "### Process dataset\n",
    "\n",
    "Convert the raw data into 10s windowed data `X` at 30Hz, with step count annotations `Y_step` and group labels `groups`. A binary label array `Y` is also generated (for reference, not used in this notebook). Save to disk for later reuse (can be skipped if already done before)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "df614b46-de45-4888-aeef-2fbc4471947e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 39/39 [00:06<00:00,  5.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved in ./data/oxwalk_30hz_w10_o0/\n",
      "X shape: (13613, 300, 3)\n",
      "Y distribution:\n",
      "notwalk    9468\n",
      "walk       4145\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "pathlib.Path(OUT_DIR).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "x, y, t, p, y_step = zip(\n",
    "    *Parallel(n_jobs=4)(\n",
    "        delayed(make)(datafile)\n",
    "        for datafile in tqdm(glob.glob(DATAFILES))\n",
    "    )\n",
    ")\n",
    "\n",
    "X = np.vstack(x)  # data windows\n",
    "Y = np.hstack(y)  # binary labels walk/notwalk (based on Y_step >= STEP_THRESHOLD)\n",
    "Y_step = np.hstack(y_step)  # step counts per window\n",
    "T = np.hstack(t)  # timestamps\n",
    "groups = np.hstack(p)  # group ids\n",
    "\n",
    "np.save(os.path.join(OUT_DIR, 'X'), X)\n",
    "np.save(os.path.join(OUT_DIR, 'Y'), Y)\n",
    "np.save(os.path.join(OUT_DIR, 'time'), T)\n",
    "np.save(os.path.join(OUT_DIR, 'groups'), groups)\n",
    "np.save(os.path.join(OUT_DIR, 'Y_step'), Y_step)\n",
    "\n",
    "print(f\"Saved in {OUT_DIR}\")\n",
    "print(\"X shape:\", X.shape)\n",
    "print(\"Y distribution:\")\n",
    "print(pd.Series(Y).value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d16b8761-daff-43ae-889e-90bec84a0508",
   "metadata": {},
   "source": [
    "### Train stepcounter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "50fe6a93-c120-4623-90db-3e2c0aa0b1d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = './data/oxwalk_30hz_w10_o0/'\n",
    "\n",
    "X = np.load(path + 'X.npy')\n",
    "Y_step = np.load(path + 'Y_step.npy')\n",
    "groups = np.load(path + 'groups.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb7baf72-23fb-4e57-bb01-149284e3263e",
   "metadata": {},
   "outputs": [],
   "source": [
    "wd_params = {\n",
    "    'device': 'cuda' if torch.cuda.is_available() else 'cpu',\n",
    "    'batch_size': 512\n",
    "}\n",
    "\n",
    "step = StepCounter(wd_type='ssl', wd_params=wd_params,\n",
    "                   steptol=STEP_THRESHOLD, verbose=True)\n",
    "\n",
    "step.fit(X, Y_step, groups)  # training\n",
    "\n",
    "step.verbose = False  # don't log spam the user\n",
    "step.wd.verbose = False\n",
    "step.wd.batch_size = 64  # don't blow up user's memory\n",
    "step.wd.device = 'cpu'  # save device-less (cpu) model\n",
    "\n",
    "# save trained model to disk\n",
    "joblib.dump(step, 'ssl.joblib.lzma', compress=('lzma', 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a1f063c",
   "metadata": {},
   "source": [
    "### Usage\n",
    "The trained stepcounter can now be used in a Python script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c09defbe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2000,) [0. 0. 0. ... 0. 0. 0.]\n"
     ]
    }
   ],
   "source": [
    "model = joblib.load('ssl.joblib.lzma')\n",
    "\n",
    "# example data: 2000 windows, 10s at 30Hz, 3-axis\n",
    "X = np.random.rand(2000, 300, 3)  \n",
    "\n",
    "y_pred = model.predict(X)\n",
    "\n",
    "# y_pred is an array of step count values for each of the 2000 windows\n",
    "print(y_pred.shape, y_pred)  \n",
    "\n",
    "# you can optionally run the prediction on a GPU (will run on CPU by default):\n",
    "model.wd.device = 'cuda'\n",
    "y_pred = model.predict(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59039a7e",
   "metadata": {},
   "source": [
    "Or from the command line by passing the path to the joblib file:\n",
    "\n",
    "`stepcount sample.cwa --model-path ssl.joblib.lzma`\n",
    "\n",
    "With GPU:\n",
    "\n",
    "`stepcount sample.cwa --model-path ssl.joblib.lzma --pytorch-device cuda`"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}