tracking.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AR4oetmCFHUa"
      },
      "source": [
        "# VGG Follow Things Around\n",
        "\n",
        "Follow Things Around is a program to detect and track multiple\n",
        "objects in videos.  We currently support the tracking of chimpanzee\n",
        "faces, birds, cats, dogs, horses, sheep, cows, elephants, bears,\n",
        "zebras, and giraffes and are expanding the list of things that it\n",
        "can track.  Please contact us at if you are interested in tracking\n",
        "other things.\n",
        "\n",
        "Follow Things Around works by performing a two-step approach named\n",
        "\"tracking by detection\".  First, it detect the \"things\" of interest\n",
        "in all frames of a video and then, it tracks those detections\n",
        "merging them into individual tracks and filling in gaps from the\n",
        "detector."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O3vImnLxFHUf"
      },
      "source": [
        "## 1 - Read Me First\n",
        "\n",
        "This project is a [Jupyter](https://jupyter.org/) notebook to\n",
        "perform tracking by detection of multiple objects in videos and was\n",
        "designed to run in [Google\n",
        "Colab](https://colab.research.google.com/).  If you are not reading\n",
        "this notebook in Google Colab, click\n",
        "[here](https://colab.research.google.com/github/ox-vgg/follow-things-around/blob/main/tracking.ipynb).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-tfDPTizFHUi"
      },
      "source": [
        "### 1.1 - What is, and how to use, a Jupyter notebook\n",
        "\n",
        "A Jupyter notebook is a series of \"cells\".  Each cell contains either\n",
        "text (like this one) or code (like others below).  A cell that\n",
        "contains code will have a \"Run cell\" button on the left side like this\n",
        "\"<img height=\"18rem\" alt=\"The 'Run cell' button in Colab\"\n",
        "src=\"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADAAAAAwCAQAAAD9CzEMAAABTklEQVRYw+2XMU7DMBRAX6ss3VA7VV25AFNWzsDQXoAzVDlBKw6QDJwhTO3OCVjaka0VXVKJDUVC4jOgiMHYcRx9S0j9f7XfS5x8+xsu8R9iQEpGyY4TgnBiR0lGyqA/fMaaI2LJI2tm4fAxObUV3mRNzjgEP+fcCm/yzLwbPKHwhjdZkPjiR2w64wVhw8jv6bdBeEHY+rxFEYz/WaiWWPTCC8LChZ9Q9RZUTOyCvDdeEHJ71drL6o43b0Ftq+6VYxJc8ciXp2L1F37IwSkAuOXVS3BgaApS55TfInzg00ORmoLMSwBww0urIDMFpbcAEpZ8OMeXpmDfQQBwzbNj/N6cUHUUANzzbi03I+oAAUx5stRCfIH6Eql/ZPXfVL3Q1LcK9c1OfbuOcOCoH5kRDn31tiVC4xWhdVRvfiO07xEuIFGuUBEugVGusZfQj28NImRviDLNnQAAAABJRU5ErkJggg==\">\".\n",
        "When you click the \"Run cell\" button, the code in that cell will run\n",
        "and when it finishes, a green check mark appears next to the \"Run\n",
        "cell\" button\".  You need to wait for the code in that cell to finish\n",
        "before \"running\" the next cell.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4kkvU97kDjYh"
      },
      "source": [
        "### 1.2 - Particulars of this notebook\n",
        "\n",
        "This notebook was designed to run in Google Colab and to analyse\n",
        "videos in Google Drive.  It will also save back the analysis results\n",
        "in Google Drive.  As such, it requires a Google account.\n",
        "\n",
        "You must run the cells on this notebook one after the other since each\n",
        "cell is dependent on the results of the previous cell.\n",
        "\n",
        "This notebook also some interactive cells, namely in the options\n",
        "sections.  After setting their values, these cells must be run, just\n",
        "like other code cells.  Setting their values only has effect after you\n",
        "\"run\" their cell."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUaiJq1V-tI"
      },
      "source": [
        "### 1.3 - Testing this notebook\n",
        "\n",
        "We recommend you first test this notebook with a short video, less\n",
        "than 20 seconds long.  First, try our own sample video to track\n",
        "chimpanzee faces (download it\n",
        "[here](https://thor.robots.ox.ac.uk/software/chimpanzee-tracking/test-data/19-mini.mp4)).\n",
        "Then try a video fragment of your own video.  Finally, run it on\n",
        "your full video."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iXk_VjsyDmtS"
      },
      "source": [
        "### 1.4 - Results files\n",
        "\n",
        "This notebook will save all results for a single video in one\n",
        "directory.  It will generate the following files:\n",
        "\n",
        "- `frames` - a directory with the individual video frames.  You may\n",
        "  want to delete them after validating the results.  They take up a\n",
        "  lot of space and can be regenerated later.\n",
        "\n",
        "- `detections.pkl` - the initial detections in [Python's pickle\n",
        "  format](https://docs.python.org/3/library/pickle.html).\n",
        "\n",
        "- `detections-via.json` - the initial detections as a\n",
        "  [VIA 2](https://www.robots.ox.ac.uk/~vgg/software/via/) project.\n",
        "  This requires the images in the `frames` directory.\n",
        "\n",
        "- `tracks.pkl` - the detected tracks in [Python's pickle\n",
        "  format](https://docs.python.org/3/library/pickle.html).\n",
        "\n",
        "- `results-via-project.json` - the final detections as a [VIA\n",
        "  2](https://www.robots.ox.ac.uk/~vgg/software/via/) project.  This\n",
        "  requires the images in the `frames` directory.\n",
        "\n",
        "- `results.csv` - the final detections in CSV format.\n",
        "\n",
        "- `tracks.mp4` - video with tracks (see Section 6).\n",
        "\n",
        "Note that none of those files includes the input video filename.\n",
        "Our recommendation is to create separate results directory for each\n",
        "video.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vuE9bu6GDpOv"
      },
      "source": [
        "### 1.5 - GPU access\n",
        "\n",
        "A GPU is required to run this pipeline in a sensible manner.  For\n",
        "example, without a GPU, a two minutes video will take close to two\n",
        "hours to process.\n",
        "\n",
        "By default, this notebook will run with a GPU.  However, it is\n",
        "possible that you were not allocated one, typically because you've\n",
        "used up all your GPU resources.  You can confirm this, and possibly\n",
        "change it, manually.  To do that, navigate to \"Edit\" -> \"Notebook\n",
        "Settings\" and select \"GPU\" from the \"Hardware Accelerator\" menu.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AU4JuuRqDrxQ"
      },
      "source": [
        "### 1.6 - Moving forward\n",
        "\n",
        "You can run this notebook on Google Colab but if you have a large\n",
        "collection of videos or if your videos are particularly long or have\n",
        "a high number of objects to track, you may end up using up all of\n",
        "your free Colab resources.\n",
        "\n",
        "Once you've tested Follow Things Around with a few of your videos,\n",
        "it may be worth running this on your own computers.  Contact us if\n",
        "you need help to do that.  Alternatively, you can purchase\n",
        "additional \"compute units\" with [Google Colab\n",
        "Plans](https://colab.research.google.com/signup).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wWhspjzoFHUl"
      },
      "source": [
        "## 2 - Setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z9eagfzTBgMh"
      },
      "source": [
        "### 2.1 - Check for GPU access\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "k72IXhXhFHUn"
      },
      "outputs": [],
      "source": [
        "#@markdown By default, this notebook will run with a GPU.  However, it\n",
        "#@markdown is possible that you were not allocated one.  If you get a\n",
        "#@markdown message saying that you do not have access to a GPU,\n",
        "#@markdown navigate to \"Edit\" -> \"Notebook Settings\" and select \"GPU\"\n",
        "#@markdown from the \"Hardware Accelerator\" menu.  Once you change it,\n",
        "#@markdown you need to run this cell again.\n",
        "\n",
        "import torch.cuda\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "    USE_GPU = True\n",
        "    print(\"You are using this GPU:\")\n",
        "    print(\n",
        "        \"GPU %d: %s (%d GB)\"\n",
        "        % (\n",
        "            torch.cuda.current_device(),\n",
        "            torch.cuda.get_device_name(),\n",
        "            torch.cuda.get_device_properties(\n",
        "                torch.cuda.current_device()\n",
        "            ).total_memory\n",
        "            * 1e-9,\n",
        "        )\n",
        "    )\n",
        "else:\n",
        "    USE_GPU = False\n",
        "    print(\"You are NOT connected to a GPU.  This will run very slow.\")\n",
        "    print(\"Consider reconnecting to a runtime with GPU access.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0jKiMdsdBpQO"
      },
      "source": [
        "### 2.2 - Install and load dependencies\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y3Nw3Km_FHUp"
      },
      "outputs": [],
      "source": [
        "#@markdown This step can take a few of minutes to finish.\n",
        "\n",
        "# We only need install follow-things-around but that depends on SVT\n",
        "# and Detectron2 which are not available on PyPI, so pip won't be able\n",
        "# to find them, and we have to install them ourselves from their git\n",
        "# repos.\n",
        "\n",
        "# Using `pip install --quiet` is not enough, it still prints out a\n",
        "# mysterious \"Preparing metadata (setup.py)\" message which is why we\n",
        "# redirect stdout to `/dev/null`.  Important messages should go to\n",
        "# stderr anyway.\n",
        "\n",
        "# The SVT package\n",
        "# https://www.robots.ox.ac.uk/~vgg/projects/seebibyte/software/svt/\n",
        "print(\"Installing SVT\")\n",
        "!pip install --quiet git+https://gitlab.com/vgg/svt/ > /dev/null\n",
        "\n",
        "# The Detectron2 package\n",
        "# https://github.com/facebookresearch/detectron2/\n",
        "print(\"Installing detectron2 (this will take a few minutes to finish)\")\n",
        "!pip install --quiet git+https://github.com/facebookresearch/detectron2.git > /dev/null\n",
        "\n",
        "# Finally, Follow-Things-Around\n",
        "print(\"Installing Follow-Things-Around\")\n",
        "!pip install --quiet 'follow-things-around[colab] @ git+https://github.com/ox-vgg/follow-things-around.git'\n",
        "\n",
        "import glob\n",
        "import json\n",
        "import logging\n",
        "import os\n",
        "import os.path\n",
        "import pickle\n",
        "\n",
        "import google.colab.drive\n",
        "import google.colab.output\n",
        "import requests\n",
        "import torch\n",
        "\n",
        "import follow_things_around\n",
        "from follow_things_around import (\n",
        "    detect,\n",
        "    ffmpeg_video_to_frames,\n",
        "    filter_detections,\n",
        "    make_video_with_tracks,\n",
        "    track,\n",
        "    FramesDirDataset,\n",
        ")\n",
        "from follow_things_around.via import detections_to_via2\n",
        "from follow_things_around.widgets import display_detections\n",
        "\n",
        "logging.basicConfig()\n",
        "_logger = logging.getLogger()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "chRTgjgMCJdd"
      },
      "source": [
        "### 2.3 - Mount Google Drive\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iOD4aWrjFHUr"
      },
      "outputs": [],
      "source": [
        "#@markdown When you run this cell, a dialog will appear about a\n",
        "#@markdown request for access to your Google Drive Files.  This is\n",
        "#@markdown required to access the videos for analysis and to then save\n",
        "#@markdown the results.  Once you click on \"Connect to Google Drive\",\n",
        "#@markdown a pop-up window will appear to choose a Google Account and\n",
        "#@markdown then to allow access to \"Google Drive for desktop\".\n",
        "\n",
        "google.colab.drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Elhia5LHRFZf"
      },
      "source": [
        "### 2.4 - Select \"thing\" to track\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WFYckNS4Ztud"
      },
      "outputs": [],
      "source": [
        "#@markdown A model trained on the \"thing\" of interest is required.  We\n",
        "#@markdown currently have models for these things but it is possible\n",
        "#@markdown to train new ones.  Please contact us if you require a\n",
        "#@markdown detector for something else or if you require to fine-tune\n",
        "#@markdown them.\n",
        "#@markdown\n",
        "#@markdown Do not forget to \"run\" this cell after making your\n",
        "#@markdown selection.\n",
        "\n",
        "THING_TO_TRACK = \"Chimpanzee faces\"  #@param [\"Chimpanzee faces\", \"Orangutan faces\", \"Chimpanzee bodies\", \"Birds\", \"Cats\", \"Dogs\", \"Horses\", \"Sheep\", \"Cows\", \"Elephants\", \"Bears\", \"Zebras\", \"Giraffes\"]\n",
        "\n",
        "THING_TO_MODEL_CONFIG = {\n",
        "    \"Chimpanzee faces\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_LSJ-CFbootstrap.py\",\n",
        "        \"class-idx\": 0,\n",
        "    },\n",
        "    \"Orangutan faces\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_LSJ-orangutan-faces.py\",\n",
        "        \"class-idx\": 0,\n",
        "    },\n",
        "    \"Chimpanzee bodies\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_LSJ-imdb_5k_sup.py\",\n",
        "        \"class-idx\": 0,\n",
        "    },\n",
        "    \"Birds\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 14,\n",
        "    },\n",
        "    \"Cats\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 15,\n",
        "    },\n",
        "    \"Dogs\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 16,\n",
        "    },\n",
        "    \"Horses\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 17,\n",
        "    },\n",
        "    \"Sheep\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 18,\n",
        "    },\n",
        "    \"Cows\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 19,\n",
        "    },\n",
        "    \"Elephants\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 20,\n",
        "    },\n",
        "    \"Bears\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 21,\n",
        "    },\n",
        "    \"Zebras\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 22,\n",
        "    },\n",
        "    \"Giraffes\": {\n",
        "        \"config-url\": \"https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/mask_rcnn_R_50_FPN_400ep_LSJ-coco2017.py\",\n",
        "        \"class-idx\": 23,\n",
        "    },\n",
        "}\n",
        "\n",
        "if THING_TO_TRACK not in THING_TO_MODEL_CONFIG:\n",
        "    raise Exception(\n",
        "        'THING_TO_TRACK \\'%s\\' has no pre-configuration \\'%s\\''\n",
        "        % (THING_TO_TRACK, list(THING_TO_MODEL_CONFIG.keys()))\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lb_oICb3CPpC"
      },
      "source": [
        "### 2.5 - Video file and results folder\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Gk90XC-8FHUs"
      },
      "outputs": [],
      "source": [
        "#@markdown To find the correct path, open the \"Files\" menu in the left\n",
        "#@markdown sidebar.  The `drive` directory contains your Google Drive\n",
        "#@markdown files.  Navigate the files, right click on the wanted file\n",
        "#@markdown or directory, and select \"Copy path\".  Then paste the path\n",
        "#@markdown in this form.  Do not forget to then \"run\" this cell.\n",
        "\n",
        "VIDEO_FILE = ''  #@param {type:\"string\"}\n",
        "RESULTS_DIRECTORY = ''  #@param {type:\"string\"}\n",
        "\n",
        "if not VIDEO_FILE:\n",
        "    raise Exception('VIDEO_FILE is empty, you must set it.')\n",
        "if not RESULTS_DIRECTORY:\n",
        "    raise Exception('RESULTS_DIRECTORY is empty, you must set it.')\n",
        "if not os.path.isfile(VIDEO_FILE):\n",
        "    raise Exception('The VIDEO_FILE \\'%s\\' does not exist' % VIDEO_FILE)\n",
        "if not os.path.isdir(RESULTS_DIRECTORY):\n",
        "    raise Exception(\n",
        "        'The RESULTS_DIRECTORY \\'%s\\' does not exist' % RESULTS_DIRECTORY\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7c1gbeVrFHUu"
      },
      "source": [
        "### 2.6 - Advanced options\n",
        "\n",
        "The cells hidden in this section expose the advanced options for this\n",
        "pipeline and perform the final setup.  In most cases you do not need\n",
        "to change their values.  You can click the \"Run cell\" button to run\n",
        "all the hidden cells inside.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SH-IeyipFHUv"
      },
      "outputs": [],
      "source": [
        "#@markdown #### 2.6.1 - Detection step\n",
        "\n",
        "#@markdown The detection step is the first step.  It detects the\n",
        "#@markdown location of the \"things\" of interest in all frames of the\n",
        "#@markdown video.\n",
        "\n",
        "#@markdown A detection model is required.  You can either train\n",
        "#@markdown your own model, or you can use one of our pre-trained\n",
        "#@markdown models.  This variable is the URL, or file path, for a\n",
        "#@markdown detectron2 model configuration (the URI and file path for\n",
        "#@markdown the model weights are specified in the config file).\n",
        "#@markdown\n",
        "\n",
        "#@markdown Leave this empty to use the pre-configured model to detect\n",
        "#@markdown the `THING_TO_TRACK` choosen in Section 2.4.\n",
        "\n",
        "DETECTION_MODEL_CONFIG_URL = ''  #@param {type: \"string\"}\n",
        "DETECTION_CLASS_IDX = 0\n",
        "if not DETECTION_MODEL_CONFIG_URL:\n",
        "    DETECTION_MODEL_CONFIG_URL = THING_TO_MODEL_CONFIG[THING_TO_TRACK][\"config-url\"]\n",
        "    DETECTION_CLASS_IDX = THING_TO_MODEL_CONFIG[THING_TO_TRACK][\"class-idx\"]\n",
        "\n",
        "#@markdown When the model detects something, that detection is\n",
        "#@markdown made with a confidence score.  Detections with a confidence\n",
        "#@markdown score lower than the threshold will be ignored during\n",
        "#@markdown tracking.  If you set the threshold too high, you may miss\n",
        "#@markdown detections in some frames which need to be filled by the\n",
        "#@markdown tracker..  If you set it too low, false detections will\n",
        "#@markdown lead to false tracks that need to be manually removed\n",
        "#@markdown later.\n",
        "\n",
        "DETECTION_THRESHOLD = 0.9  #@param {type: \"slider\", min: 0.0, max: 1.0, step: 0.01}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nmuWC94sFHUw"
      },
      "outputs": [],
      "source": [
        "#@markdown #### 2.6.2 - Tracking step\n",
        "\n",
        "#@markdown The final step is to merge the detections into individual\n",
        "#@markdown tracks.\n",
        "\n",
        "#@markdown You will need to provide a model for the tracker.  We\n",
        "#@markdown recommend you use [this\n",
        "#@markdown one](https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/tracking-model-20181031_e45.pth).\n",
        "#@markdown Alternatively, you can also specify a path in your Google\n",
        "#@markdown Drive.\n",
        "TRACKING_MODEL_URL = 'https://thor.robots.ox.ac.uk/models/staging/chimp-tracking/tracking-model-20181031_e45.pth'  #@param {type: \"string\"}\n",
        "\n",
        "MATCH_OVERLAP_THRESHOLD = 0.2  #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "\n",
        "NONMATCH_TRACKING_THRESHOLD = 0.9  #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "\n",
        "UNKNOWN_TRACK_ID_MARKER = -1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YUIcftRMFHUy"
      },
      "outputs": [],
      "source": [
        "#@markdown #### 2.6.3 - Verbosity\n",
        "\n",
        "#@markdown How chatty do you want the notebook to be?  INFO is a good\n",
        "#@markdown choice if you want to have a feeling for progress.\n",
        "LOG_LEVEL = 'INFO'  #@param [\"CRITICAL\", \"ERROR\", \"WARNING\", \"INFO\", \"DEBUG\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OWKXpxAuFHUy"
      },
      "outputs": [],
      "source": [
        "#@markdown #### 2.6.4 - The final setup step\n",
        "\n",
        "#@markdown Run this cell to perform the final pipeline setup based on\n",
        "#@markdown the given options.\n",
        "\n",
        "logging.getLogger().setLevel(LOG_LEVEL)\n",
        "\n",
        "FRAMES_DIR = os.path.join(RESULTS_DIRECTORY, 'frames')\n",
        "DETECTIONS_PKL_FPATH = os.path.join(RESULTS_DIRECTORY, 'detections.pkl')\n",
        "DETECTIONS_VIA_FPATH = os.path.join(RESULTS_DIRECTORY, 'detections-via.json')\n",
        "TRACKS_PKL_FPATH = os.path.join(RESULTS_DIRECTORY, 'tracks.pkl')\n",
        "RESULTS_VIA_FPATH = os.path.join(RESULTS_DIRECTORY, 'results-via-project.json')\n",
        "RESULTS_CSV_FPATH = os.path.join(RESULTS_DIRECTORY, 'results.csv')\n",
        "TRACKS_VIDEO_FPATH = os.path.join(RESULTS_DIRECTORY, 'tracks.mp4')\n",
        "\n",
        "\n",
        "# These should never be true because USE_GPU was set automatically\n",
        "# based on whether a GPU is available.\n",
        "if USE_GPU and not torch.cuda.is_available():\n",
        "    raise Exception('Your runtime does not have a GPU.')\n",
        "elif torch.cuda.is_available() and not USE_GPU:\n",
        "    _logger.warn('You have a GPU but chose to not use it.  Are you sure?')\n",
        "\n",
        "if USE_GPU:\n",
        "    DEFAULT_DEVICE = 'cuda'\n",
        "    torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "else:\n",
        "    DEFAULT_DEVICE = 'cpu'\n",
        "    torch.set_default_tensor_type('torch.FloatTensor')\n",
        "\n",
        "_logger.info('Will use %s device.', DEFAULT_DEVICE.upper())\n",
        "\n",
        "\n",
        "# Required to display the tracking results with plotly or matplotlib.\n",
        "google.colab.output.enable_custom_widget_manager()\n",
        "\n",
        "\n",
        "follow_things_around.DEFAULT_DEVICE = DEFAULT_DEVICE\n",
        "follow_things_around.FRAMES_DIR = FRAMES_DIR\n",
        "follow_things_around.MATCH_OVERLAP_THRESHOLD = MATCH_OVERLAP_THRESHOLD\n",
        "follow_things_around.NONMATCH_TRACKING_THRESHOLD = NONMATCH_TRACKING_THRESHOLD\n",
        "follow_things_around.UNKNOWN_TRACK_ID_MARKER = UNKNOWN_TRACK_ID_MARKER\n",
        "follow_things_around.USE_GPU = USE_GPU\n",
        "follow_things_around.VIDEO_FILE = VIDEO_FILE\n",
        "\n",
        "logging2ffmpeg_loglevel = {\n",
        "    'CRITICAL': 'fatal',\n",
        "    'ERROR': 'error',\n",
        "    'WARNING': 'warning',\n",
        "    'INFO': 'info',\n",
        "    'DEBUG': 'debug',\n",
        "}\n",
        "follow_things_around.FFMPEG_LOG_LEVEL = logging2ffmpeg_loglevel[LOG_LEVEL]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XIgX5CS8FHU0"
      },
      "source": [
        "## 3 - Convert video to frames\n",
        "\n",
        "The pipeline needs the video frames as individual image files.  This\n",
        "cell will create a `frames` directory and save the individual images\n",
        "there.  You may skip running this cell if you already have a `frames`\n",
        "directory with images.  This cell will error if the `frames` directory\n",
        "already exists to prevent overwriting any existing data.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E4xj4YRWFHU0"
      },
      "outputs": [],
      "source": [
        "#@markdown Skip this cell if you already have the frames.  If you run\n",
        "#@markdown this cell and the `frames` directory already exists, it\n",
        "#@markdown errors to avoid overwriting any previous images.\n",
        "\n",
        "os.makedirs(FRAMES_DIR, exist_ok=False)\n",
        "\n",
        "ffmpeg_video_to_frames(VIDEO_FILE, FRAMES_DIR)\n",
        "\n",
        "_logger.info('Finished extracting individual frames to \\'%s\\'', FRAMES_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AAhN5CkMFHU0"
      },
      "source": [
        "## 4 - Detection step\n",
        "\n",
        "The detection of \"things\" is the first step in the pipeline.  If you\n",
        "have previously run the detection step then you will have a\n",
        "`detections.pkl` file in the results directory.  If so, skip the\n",
        "\"Run detection\" cell (section 4.1) and run the \"Load previous\n",
        "detections results\" cell instead (section 4.2).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jk_beBJTFHU2"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 4.1 - Run detection (option 1)\n",
        "\n",
        "dataset = FramesDirDataset(FRAMES_DIR)\n",
        "if len(dataset) == 0:\n",
        "    raise Exception(\n",
        "        \"No files in '%s'.  Did you run the previous section which converts\"\n",
        "        \" the video to frames?\" % FRAMES_DIR\n",
        "    )\n",
        "\n",
        "detections = detect(dataset, DETECTION_MODEL_CONFIG_URL, DETECTION_CLASS_IDX)\n",
        "\n",
        "with open(DETECTIONS_PKL_FPATH, 'wb') as fh:\n",
        "    pickle.dump({'detections': detections}, fh)\n",
        "_logger.info(\"Detection results saved to '%s'\", DETECTIONS_PKL_FPATH)\n",
        "\n",
        "with open(DETECTIONS_VIA_FPATH, 'w') as fh:\n",
        "    json.dump(detections_to_via2(dataset, detections), fh)\n",
        "_logger.info(\"Detection VIA project saved to '%s'\", DETECTIONS_VIA_FPATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xHVc5I4EFHU2"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 4.2 - Load previous detection results (option 2)\n",
        "\n",
        "dataset = FramesDirDataset(FRAMES_DIR)\n",
        "if len(dataset) == 0:\n",
        "    raise Exception(\n",
        "        \"No files in '%s'.  Did you run the previous section which converts\"\n",
        "        \" the video to frames?\" % FRAMES_DIR\n",
        "    )\n",
        "\n",
        "with open(DETECTIONS_PKL_FPATH, 'rb') as fh:\n",
        "    loaded_detections = pickle.load(fh)\n",
        "detections = loaded_detections['detections']\n",
        "\n",
        "_logger.info('Detection results loaded from \\'%s\\'', DETECTIONS_PKL_FPATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FN5FGZxJFHU3"
      },
      "source": [
        "## 5 - Tracking step\n",
        "\n",
        "The final step in the pipeline is to track the things detected in\n",
        "the detection step.  At the end of this step, the tracking results\n",
        "will be saved in a CSV file and as a\n",
        "[VIA](https://www.robots.ox.ac.uk/~vgg/software/via/) project.\n",
        "\n",
        "If you have previously run the tracking step then you will have a\n",
        "`tracks.pkl` file in the results directory.  If so, skip the \"Run\n",
        "tracking\" cell (section 5.1) and run the \"Load previous tracking\n",
        "results\" cell instead (section 5.2)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "GodRuQQ4FHU3"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 5.1 - Run tracking (option 1)\n",
        "\n",
        "tracks = track(\n",
        "    dataset,\n",
        "    filter_detections(detections, DETECTION_THRESHOLD),\n",
        "    TRACKING_MODEL_URL,\n",
        ")\n",
        "\n",
        "tracks.export_via_project(\n",
        "    RESULTS_VIA_FPATH,\n",
        "    config={\n",
        "        'frame_img_dir': dataset.frames_dir,\n",
        "        'via_project_name': '',\n",
        "    },\n",
        ")\n",
        "tracks.export_plain_csv(RESULTS_CSV_FPATH, {})\n",
        "\n",
        "with open(TRACKS_PKL_FPATH, 'wb') as fh:\n",
        "    pickle.dump({'tracks': tracks}, fh)\n",
        "_logger.info('Tracking results saved to \\'%s\\'', TRACKS_PKL_FPATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_TQC37GrNTFw"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 5.2 - Load previous tracking results (option 2)\n",
        "\n",
        "with open(TRACKS_PKL_FPATH, 'rb') as fh:\n",
        "    loaded_tracks = pickle.load(fh)\n",
        "tracks = loaded_tracks['tracks']\n",
        "\n",
        "_logger.info('Tracking results loaded from \\'%s\\'', TRACKS_PKL_FPATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1XoTxVlsFHU3"
      },
      "source": [
        "## 6 - Visualise results\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LL5er1FZCiBo"
      },
      "source": [
        "### 6.1 - Visualise in Google Colab (option 1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4sbn8DsHIi1C"
      },
      "outputs": [],
      "source": [
        "#@markdown You can see the tracking results right here, inside this\n",
        "#@markdown Google Colab notebook, but the interface is a bit slow.\n",
        "#@markdown This is fine if you want to have a quick look at some of\n",
        "#@markdown of frames only.\n",
        "\n",
        "#@markdown Run this cell and then click on the arrow buttons to\n",
        "#@markdown display the next or previous frame, and you can move the\n",
        "#@markdown slider to move to a specific frame.  When you dragging the\n",
        "#@markdown slider, the displayed frame is only updated once the slider\n",
        "#@markdown is released.  Expect a couple of seconds wait for the frame\n",
        "#@markdown to be updated.\n",
        "\n",
        "display_detections(dataset, tracks.detection_data['0'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i5FjIafRC8te"
      },
      "source": [
        "### 6.2 - Visualise locally with VIA (option 2)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QBdj3CToLPVG"
      },
      "source": [
        "[VIA](https://www.robots.ox.ac.uk/~vgg/software/via/) is a web\n",
        "application to view and perform annotations of image, audio, and\n",
        "video.  It is free software and runs locally on the web browser.  You\n",
        "can view the tracking results on the individual frames with VIA.\n",
        "\n",
        "This is much more responsive than viewing inside the notebook but\n",
        "requires download the frames locally (either manually or with [Google\n",
        "Drive for\n",
        "Desktop](https://support.google.com/a/users/answer/13022292)).\n",
        "\n",
        "1. Download [VIA\n",
        "   2](https://www.robots.ox.ac.uk/~vgg/software/via/downloads/via-2.0.12.zip).\n",
        "   This is a zip file.  Open it.  Inside there is a `via.html` file.\n",
        "   Open it in your web browser to start VIA.\n",
        "\n",
        "2. Download the `results-via-project.json` from your results diretcory\n",
        "   and the whole frames directory.  If you are using Google Drive for\n",
        "   Desktop sync it now.  The frames directory is pretty large and this\n",
        "   step may take a long time.\n",
        "\n",
        "3. Navigate to \"Project\" -> \"Load\" and select the\n",
        "   `results-via-project.json` file.  A \"File Not Found\" error message\n",
        "   will appear.  This means that VIA does not know where the images\n",
        "   are.\n",
        "\n",
        "4. Navigate to \"Project\" -> \"Settings\".  Set the \"Default Path\" to the\n",
        "   `frames` directory in your computer.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5kfdOPAfCpyx"
      },
      "source": [
        "### 6.3 - Create video file with tracks (option 3)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "l44oU0LWA7nw"
      },
      "outputs": [],
      "source": [
        "#@markdown You may also generate a video file with the detections\n",
        "#@markdown superimposed.  The video file will be named `tracks.mp4`\n",
        "#@markdown and saved in the `RESULTS_DIRECTORY` in your Google Drive.\n",
        "\n",
        "make_video_with_tracks(\n",
        "    VIDEO_FILE, TRACKS_VIDEO_FPATH, FRAMES_DIR, RESULTS_CSV_FPATH\n",
        ")\n",
        "_logger.info('Video file with tracks created \\'%s\\'', TRACKS_VIDEO_FPATH)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "7c1gbeVrFHUu"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 2
}