Training_and_inference_on_an_example_dataset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Training and inference on an example dataset",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "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_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": {
        "id": "LlV70jDuWzea",
        "colab_type": "text"
      },
      "source": [
        "In this notebook we'll show you how to install SLEAP on Colab, download a dataset from the repository of sample datasets, run training and inference on that dataset using the SLEAP command-line interface, and then download the predictions."
      ]
    },
    {
      "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",
        "outputId": "12329081-eaeb-49aa-933e-489d7e6de42a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "import tensorflow as tf\n",
        "tf.test.gpu_device_name()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "''"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pMjYpcVFqWgV",
        "colab_type": "text"
      },
      "source": [
        "Since Colab already has TensorFlow and GPU drivers installed we can install SLEAP using `pip`.\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 onto [Google Drive](https://www.google.com/drive).\n",
        "\n",
        "For this demo we'll instead download a sample dataset from the SLEAP [sample datasets repository](https://github.com/murthylab/sleap-datasets)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fm3cU1Bc0tWc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!apt-get install tree\n",
        "!wget -O dataset.zip https://github.com/murthylab/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"
      ],
      "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 parameters.\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` fields in the training profile)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QKf6qzMqNBUi",
        "colab_type": "code",
        "colab": {}
      },
      "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\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vm3i0ry04IMx",
        "colab_type": "text"
      },
      "source": [
        "Let's also train a centered instance model since we'll need both a centroid model and a top-down confidence maps model to run inference with the top-down pipeline."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ufbULTDw4Hbh",
        "colab_type": "code",
        "colab": {}
      },
      "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\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whOf8PaFxYbt",
        "colab_type": "text"
      },
      "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",
      "metadata": {
        "id": "GBUTQ2Cm44En",
        "colab_type": "code",
        "outputId": "ca298981-af65-43b3-f0f6-573f423acba8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        }
      },
      "source": [
        "!tree models/"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "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"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nIsKUX661xFK",
        "colab_type": "text"
      },
      "source": [
        "# Inference\n",
        "\n",
        "At this point you should have SLEAP installed, the sample dataset downloaded, and trained models for centroids and topdown confidence maps."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CLtjtq9E1Znr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-track \"dataset/drosophila-melanogaster-courtship/20190128_113421.mp4\" --frames 0-100 -m \"models/courtship.centroid\" -m \"models/courtship.topdown_confmaps\""
      ],
      "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. 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",
      "metadata": {
        "id": "n6KVfWDIDEUe",
        "colab_type": "code",
        "outputId": "b0633306-f24f-4e6e-e78f-e968a765a3c6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 136
        }
      },
      "source": [
        "!tree dataset/drosophila-melanogaster-courtship"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "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"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3mf3KZQj_GhH",
        "colab_type": "text"
      },
      "source": [
        "You can inspect your predictions file using `sleap-inspect`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-jbVP_s06hMh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sleap-inspect dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JFfHDVy7_iDz",
        "colab_type": "text"
      },
      "source": [
        "If you're using Chrome you can download your trained models like so:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ej2it8dl_BO_",
        "colab_type": "code",
        "colab": {}
      },
      "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\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iskOQI-r_zNO",
        "colab_type": "text"
      },
      "source": [
        "And you can likewise download your predictions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gdXCYnRV_omC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import files\n",
        "files.download('dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Fy26NVmCWFw",
        "colab_type": "text"
      },
      "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."
      ]
    }
  ]
}