Training_and_inference_using_Google_Drive.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Training and inference using Google Drive",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/murthylab/sleap-notebooks/blob/master/Training_and_inference_using_Google_Drive.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K5xp-A8Oc80Q",
        "colab_type": "text"
      },
      "source": [
        "This notebook shows you how to run training and inference on your own SLEAP dataset using the command-line interface.\n",
        "\n",
        "You'll use [Google Drive](https://www.google.com/drive) for moving your data to and from Colab. We'll guide you through the process."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yX9noEb8m8re",
        "colab_type": "text"
      },
      "source": [
        "# Installation\n",
        "\n",
        "Before we install SLEAP we need to set Colab to use TensorFlow 2."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1hJvyI24KXlu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%tensorflow_version 2.x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C5WL10cao-sc",
        "colab_type": "text"
      },
      "source": [
        "Let's confirm that we have a GPU available. This next line should return something like \"/device:GPU:0\". If you instead see an empty string as the result, go to \"Notebook settings\" in the \"Edit\" menu and select \"GPU\" as the hardware accelerator."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9I12Lk6TKZnn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "tf.test.gpu_device_name()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pMjYpcVFqWgV",
        "colab_type": "text"
      },
      "source": [
        "Now let's use `pip` to install SLEAP from PyPI.\n",
        "\n",
        "Note: This installation method should also work on other Linux machines, such as an HPC cluster, or on any system where you aren't planning to use a GPU. To use a GPU on a Windows machine you'll need to install using `conda`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DUfnkxMtLcK3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install sleap==1.0.10"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iq7jrgUksLtR",
        "colab_type": "text"
      },
      "source": [
        "## Getting your training data into Colab\n",
        "\n",
        "You'll need to get your training data into Colab. The easiest way to do this is to export a self-contained **training package** from SLEAP and copy that on to [Google Drive](https://www.google.com/drive).\n",
        "\n",
        "Let's first mount your Google Drive so that it can be accessed from Colab. You'll be prompted to log into your Google account, give Colab access, and the copy an authorization code into a field below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GBWjF4jpMG2N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KQhv_gsdJzaq",
        "colab_type": "text"
      },
      "source": [
        "Now let's create a training package and copy it into your Google Drive.\n",
        "\n",
        "A training package contains both labeled data as well as the labeled images which will be used for training. One advantage to training packages is that it doesn't depend on paths to other files (i.e., videos) to be messed up when you copy your project to another volume.\n",
        "\n",
        "See [this guide](https://sleap.ai/guides/training-package.html) for exporting a training package from SLEAP.\n",
        "\n",
        "\n",
        "\n",
        "The following cell sets your current working directory to the directory with your training package. This will ensure that the output from training (i.e., the models) and from interence (i.e., predictions) will all be saved in this directory.\n",
        "\n",
        "**Important**: For this demo, I'll assume you have a `sleap` directory in the root of your Google drive and you place the exported training package in that directory with the filename `colab.pkg.slp`. If you placed your training pckage somewhere else, you'll need to adjust the path below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9umui-gI2rBz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "os.chdir(\"/content/drive/My Drive/sleap\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xZ-sr67av5uu"
      },
      "source": [
        "## Training\n",
        "\n",
        "Now you're ready to train a model! We'll use the command-line interface for training, and train a model for confidence maps using the default **training profile**. The training profile determines the model architecture, the learning rate for training, and other training hyperparameters.\n",
        "\n",
        "When you start running this cell, you'll see the training parameters listed and then you'll see the training and validation loss for each training epoch.\n",
        "\n",
        "If you're happy with the validation loss you see for an epoch during training, you're welcome to stop training by clicking the stop button next to the notebook cell running training. The version of the model with the lowest validation loss is saved during training, and that's what will be used for inference. 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` parameter in the training profile).\n",
        "\n",
        "**Important**: If your training package isn't named `colab.pkg.slp`, you'll need to adjust the name below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QKf6qzMqNBUi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-train baseline_medium_rf.bottomup.json colab.pkg.slp --run_name \"colab_demo.bottomup\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whOf8PaFxYbt",
        "colab_type": "text"
      },
      "source": [
        "Once training finishes, you'll have a trained model for confidence maps on your Google Drive. There will be a `models/` directory inside your `sleap/` directory (or wherever you had the training package), and inside this there will be a `colab_demo.bottomup` directory (the name of this directory was set by the `--run_name` parameter). \n",
        "\n",
        "This `colab_demo.bottomup` directory contains all the files SLEAP needs to use this model. You can copy it to a local drive if you want to use it for running inference from the SLEAP GUI, copy it to a network drive if you want to run inference from an HPC cluster, or leave it on your Google Drive if you want to run inference on Colab... as we'll do below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q7JUp9P840qB",
        "colab_type": "text"
      },
      "source": [
        "## Training other models\n",
        "\n",
        "The **bottomup** model you trained above can be used for \"bottom up\" inference. You can also train a **centroid** model and a **centered instance** model for \"top down\" inference. See [here](https://sleap.ai/#getting-started-with-sleap) for more information about these two different approaches.\n",
        "\n",
        "Here's how to train centroid and centered instance models using the default training settings:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WZVftKcJ46jU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-train baseline_medium_rf.topdown.json --run_name \"colab_demo.topdown_confmaps\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DsPTDnkW49-5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-train baseline.centroid.json colab.pkg.slp --run_name \"colab_demo.centroid\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nIsKUX661xFK",
        "colab_type": "text"
      },
      "source": [
        "# Inference\n",
        "\n",
        "At this point you should have SLEAP installed, your Google Drive mounted, and trained models saved on your Google Drive. If you've been working through the notebook, you should have a `models` subdirectory inside your current working directory. Let's take a look:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "01HDohE5upwQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ls models"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MNi260mJuyJ7",
        "colab_type": "text"
      },
      "source": [
        "We'll also need a video for which we want predictions. Copy the video onto your Google drive.\n",
        "\n",
        "For this demo we'll just get predictions for the first 200 frames (or you can adjust the `--frames` parameter below or remove it to run on the whole video).\n",
        "\n",
        "**Important**: If your video is not named `colab_demo.mp4`, change this in the following cell to match the name of your video. If you trained top-down models and not a bottom-up model, see the end of the notebook for how to run inference with the pair of top-down models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CLtjtq9E1Znr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-track colab_demo.mp4 \\\n",
        "    --frames 0-200 \\\n",
        "    --tracking.tracker simple \\\n",
        "    -m models/colab_demo.bottomup"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nzObCUToEqwA",
        "colab_type": "text"
      },
      "source": [
        "When inference is finished, it will save the predictions in a file which can be opened in the GUI as a SLEAP project file. The file will be in the same directory as the video and the filename will be `{video filename}.predictions.slp`.\n",
        "\n",
        "Let's inspect the predictions file:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nPfmNMSt-vS7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-inspect colab_demo.mp4.predictions.slp"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LoJ2kNBK-w6k",
        "colab_type": "text"
      },
      "source": [
        "You can copy this file from your Google Drive to a local drive and open it in the SLEAP GUI app (or open it directly if you have your Google Drive mounted on your local machine). If the video is in the same directory as the predictions file, SLEAP will automatically find it; otherwise, you'll be prompted to locate the video (since the path to the video on your local machine will be different than the path to the video on Colab)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mf-MsD7fv5Wy",
        "colab_type": "text"
      },
      "source": [
        "## Inference parameters\n",
        "\n",
        "One important option when running inference is whether (and how) you want to track instance identities. If you omit `--tracking.tracker flow` then the identities will not be tracked. Tracking methods/options are explained [here](https://sleap.ai/guides/proofreading.html#tracking-methods).\n",
        "\n",
        "You can see all of the command-line arguments by calling `sleap-track` with the `--help` argument, like so:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6hh4ogd9wtxP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-track --help"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qW-NoJOFvYHM",
        "colab_type": "text"
      },
      "source": [
        "## Inference with top-down models\n",
        "\n",
        "If you trained the pair of models needed for top-down inference, you can call `sleap-track` with `-m path/to/model` for each model, like so:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QPKnMc1qvim7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-track colab_demo.video.mp4 \\\n",
        "    --frames 0-200 \\\n",
        "    --tracking.tracker simple \\\n",
        "    -m models/colab_demo.topdown_confmaps \\\n",
        "    -m models/colab_demo.centroid"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}