{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/talmolab/sleap/blob/main/docs/notebooks/Training_and_inference_on_an_example_dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training and inference on an example dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LlV70jDuWzea"
   },
   "source": [
    "In this notebook we'll install SLEAP, download a sample dataset, run training and inference on that dataset using the SLEAP command-line interface, and then download the predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yX9noEb8m8re"
   },
   "source": [
    "## Install SLEAP\n",
    "Note: Before installing SLEAP check [SLEAP releases](https://github.com/talmolab/sleap/releases) page for the latest version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "DUfnkxMtLcK3",
    "outputId": "a6340ef1-eaac-42ef-f8d4-bcc499feb57b"
   },
   "outputs": [],
   "source": [
    "!pip uninstall -y opencv-python opencv-contrib-python\n",
    "!pip install sleap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iq7jrgUksLtR"
   },
   "source": [
    "## Download sample training data into Colab\n",
    "Let's download a sample dataset from the SLEAP [sample datasets repository](https://github.com/talmolab/sleap-datasets) into Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fm3cU1Bc0tWc",
    "outputId": "c0ac5677-e3c5-477c-a2f7-44d619208b22"
   },
   "outputs": [],
   "source": [
    "!apt-get install tree\n",
    "!wget -O dataset.zip https://github.com/talmolab/sleap-datasets/releases/download/dm-courtship-v1/drosophila-melanogaster-courtship.zip\n",
    "!mkdir dataset\n",
    "!unzip dataset.zip -d dataset\n",
    "!rm dataset.zip\n",
    "!tree dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xZ-sr67av5uu"
   },
   "source": [
    "## Train models\n",
    "For the top-down pipeline, we'll need train two models: a centroid model and a centered-instance model.\n",
    "\n",
    "Using the command-line interface, we'll first train a model for centroids using the default **training profile**. The training profile determines the model architecture, the learning rate, and other parameters.\n",
    "\n",
    "When you start training, you'll first see the training parameters and then the training and validation loss for each training epoch. \n",
    "\n",
    "As soon as you're satisfied with the validation loss you see for an epoch during training, you're welcome to stop training by clicking the stop button. The version of the model with the lowest validation loss is saved during training, and that's what will be used for inference.\n",
    "\n",
    "If you don't stop training, it will run for 200 epochs or until validation loss fails to improve for some number of epochs (controlled by the `early_stopping` fields in the training profile)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QKf6qzMqNBUi"
   },
   "outputs": [],
   "source": [
    "!sleap-train baseline.centroid.json \"dataset/drosophila-melanogaster-courtship/courtship_labels.slp\" --run_name \"courtship.centroid\" --video-paths \"dataset/drosophila-melanogaster-courtship/20190128_113421.mp4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vm3i0ry04IMx"
   },
   "source": [
    "Let's now train a centered-instance model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ufbULTDw4Hbh"
   },
   "outputs": [],
   "source": [
    "!sleap-train baseline_medium_rf.topdown.json \"dataset/drosophila-melanogaster-courtship/courtship_labels.slp\" --run_name \"courtship.topdown_confmaps\" --video-paths \"dataset/drosophila-melanogaster-courtship/20190128_113421.mp4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "whOf8PaFxYbt"
   },
   "source": [
    "The models (along with the profiles and ground truth data used to train and validate the model) are saved in the `models/` directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 306
    },
    "id": "GBUTQ2Cm44En",
    "outputId": "ca298981-af65-43b3-f0f6-573f423acba8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "models/\n",
      "├── courtship.centroid\n",
      "│   ├── best_model.h5\n",
      "│   ├── initial_config.json\n",
      "│   ├── labels_gt.train.slp\n",
      "│   ├── labels_gt.val.slp\n",
      "│   ├── training_config.json\n",
      "│   └── training_log.csv\n",
      "└── courtship.topdown_confmaps\n",
      "    ├── best_model.h5\n",
      "    ├── initial_config.json\n",
      "    ├── labels_gt.train.slp\n",
      "    ├── labels_gt.val.slp\n",
      "    ├── training_config.json\n",
      "    └── training_log.csv\n",
      "\n",
      "2 directories, 12 files\n"
     ]
    }
   ],
   "source": [
    "!tree models/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nIsKUX661xFK"
   },
   "source": [
    "## Inference\n",
    "Let's run inference with our trained models for centroids and centered instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CLtjtq9E1Znr"
   },
   "outputs": [],
   "source": [
    "!sleap-track \"dataset/drosophila-melanogaster-courtship/20190128_113421.mp4\" --frames 0-100 -m \"models/courtship.centroid\" -m \"models/courtship.topdown_confmaps\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nzObCUToEqwA"
   },
   "source": [
    "When inference is finished, predictions are saved in a file. Since we didn't specify a path, it will be saved as `<video filename>.predictions.slp` in the same directory as the video:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 136
    },
    "id": "n6KVfWDIDEUe",
    "outputId": "b0633306-f24f-4e6e-e78f-e968a765a3c6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset/drosophila-melanogaster-courtship\n",
      "├── 20190128_113421.mp4\n",
      "├── 20190128_113421.mp4.predictions.slp\n",
      "├── courtship_labels.slp\n",
      "└── example.jpg\n",
      "\n",
      "0 directories, 4 files\n"
     ]
    }
   ],
   "source": [
    "!tree dataset/drosophila-melanogaster-courtship"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3mf3KZQj_GhH"
   },
   "source": [
    "You can inspect your predictions file using `sleap-inspect`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-jbVP_s06hMh"
   },
   "outputs": [],
   "source": [
    "!sleap-inspect dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JFfHDVy7_iDz"
   },
   "source": [
    "If you're using Chrome you can download your trained models like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ej2it8dl_BO_"
   },
   "outputs": [],
   "source": [
    "# Zip up the models directory\n",
    "!zip -r trained_models.zip models/\n",
    "\n",
    "# Download.\n",
    "from google.colab import files\n",
    "files.download(\"/content/trained_models.zip\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iskOQI-r_zNO"
   },
   "source": [
    "And you can likewise download your predictions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gdXCYnRV_omC"
   },
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "files.download('dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7Fy26NVmCWFw"
   },
   "source": [
    "In some other browsers (Safari) you might get an error and you can instead download using the \"Files\" tab in the side panel (it has a folder icon). Select \"Show table of contents\" in the \"View\" menu if you don't see the side panel."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Training and inference on an example dataset",
   "provenance": [],
   "toc_visible": true
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}