{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JHEt8HwIH6_x"
   },
   "source": [
    "# Training and inference with colab\n",
    "[Google colab](https://colab.research.google.com) is a free server for running notebooks with GPU/TPU support - this is a great way to use _DAS_ if you do not have a computer with a GPU.\n",
    "\n",
    "This notebook demonstrates \n",
    "- how to setup _DAS_ \n",
    "- how to load your own datasets\n",
    "- how to train a network and then use that network to label a new audio recording.\n",
    "\n",
    "Open and edit this notebook in colab by clicking this badge:\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/janclemenslab/das/blob/master/docs/tutorials/colab.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M_HarM2deUY2"
   },
   "source": [
    "Install _DAS_:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "xfOlezxOHGHu",
    "outputId": "01161dc3-d318-46e5-df45-6ce1d875c94a"
   },
   "outputs": [],
   "source": [
    "!pip install uv\n",
    "!uv pip install das --system --torch-backend=auto"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GupBBiSXFWjf"
   },
   "source": [
    "Import all the things:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Rm-lDKcqFWjf"
   },
   "outputs": [],
   "source": [
    "import das.train, das.predict, das.utils, das.io.npy_dir\n",
    "import matplotlib.pyplot as plt\n",
    "import flammkuchen\n",
    "import logging\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FevZwGYzeg5k"
   },
   "source": [
    "Mount your google drive so you can access your own datasets - this will ask for authentication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gnBEpxg02dQ4",
    "outputId": "5afc17f2-e111-4f1d-ecdb-4a825f11cb51"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4R29gIh-ex-P"
   },
   "source": [
    "## Train the model\n",
    "Adjust the variable `path_to_data` to point to the dataset on your own google drive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-jMWYYfmFWjf",
    "outputId": "fecc66d2-96f7-4103-f456-0c380128da4f",
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "path_to_data = '/content/drive/MyDrive/Dmoj.wrigleyi.npy'\n",
    "\n",
    "das.train.train(model_name='tcn',\n",
    "                data_dir=path_to_data,\n",
    "                save_dir='res',\n",
    "                nb_hist=1024,\n",
    "                kernel_size=32,\n",
    "                nb_filters=32,\n",
    "                ignore_boundaries=True,\n",
    "                verbose=2,\n",
    "                nb_conv=4,\n",
    "                learning_rate=0.0005,\n",
    "                use_separable=[True, True, False, False],\n",
    "                nb_epoch=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Vnn1RzqfAsR"
   },
   "source": [
    "Adjust the name to point to the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qpKFtj3zFWjf"
   },
   "outputs": [],
   "source": [
    "res_name = '/content/res/20210925_132436'\n",
    "res = flammkuchen.load(f'{res_name}_results.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EjvHmmXyfXLm"
   },
   "source": [
    "Inspect the history of the training and validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "zCH0jn00Bl1n",
    "outputId": "94b864a2-e135-4826-ea88-a5e1cdccdf8a"
   },
   "outputs": [],
   "source": [
    "plt.plot(res['fit_hist']['loss'], label='train')\n",
    "plt.plot(res['fit_hist']['val_loss'], label='val')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uGJ9Xkx_fe-I"
   },
   "source": [
    "Plot the test results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 223
    },
    "id": "rTuFVEKLA_iM",
    "outputId": "05a3be24-e44d-4043-8467-d73a91714d41"
   },
   "outputs": [],
   "source": [
    "# t0, t1 = 1_020_000, 1_040_000  # dmel tutorial dataset\n",
    "t0, t1 = 40_000, 60_000\n",
    "\n",
    "plt.figure(figsize=(40, 8))\n",
    "plt.subplot(311)\n",
    "plt.plot(res['x_test'][t0:t1], 'k')\n",
    "plt.ylabel('Audio')\n",
    "\n",
    "plt.subplot(312)\n",
    "plt.plot(res['y_test'][t0:t1, 1:])\n",
    "plt.ylabel('Prediction targets')\n",
    "\n",
    "plt.subplot(313)\n",
    "plt.plot(res['y_pred'][t0:t1, 1:])\n",
    "plt.ylabel('Predictions')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i3eXlLRsg891"
   },
   "source": [
    "You can download the model results via the file tab on the left, from `/contest/res`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YaBJ-gzZfqKo"
   },
   "source": [
    "## Predict on new data\n",
    "Load a new recording for prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9MouG9Z8B8nl",
    "outputId": "7a48dd7d-1895-4e87-a3cd-38a04b772a92"
   },
   "outputs": [],
   "source": [
    "model, params = das.models.load_model_and_params(res_name)  # load the model and runtime parameters\n",
    "ds = das.io.npy_dir.NpyDir.load('/content/drive/MyDrive/Dmoj.wrigleyi.npy', memmap_dirs=['train','val'])  # load the new data\n",
    "print(ds)\n",
    "x = ds['test']['x']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U0CTfLSQgMxV"
   },
   "source": [
    "Run inference - this will calculate the confidence score and extract segment boundaries and event times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AA1xaQkFDi9W",
    "outputId": "88fc4fe6-67a1-43ba-a388-9fc8196305b1"
   },
   "outputs": [],
   "source": [
    "events, segments, class_probabilities, _ = das.predict.predict(x, model=model, params=params, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 420
    },
    "id": "mx3ejHmMDzPO",
    "outputId": "4af7de93-2356-4aee-9af7-357e1b7778ef"
   },
   "outputs": [],
   "source": [
    "t0, t1 = 42_000, 48_000\n",
    "plt.figure(figsize=(20, 8))\n",
    "plt.plot(x[t0:t1, 0], alpha=0.25, c='k', label='Audio')\n",
    "plt.plot(class_probabilities[t0:t1, 1:], label='Prediction')\n",
    "plt.xlim(0, t1-t0)\n",
    "plt.legend(['Audio', 'Prediction'])\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "colab.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "file_extension": ".py",
  "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.8.11"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}