SSD_Pytorch_Demo_using_YouTube_Video.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SSD_Pytorch_Demo using YouTube Video.ipynb",
      "provenance": [],
      "private_outputs": true,
      "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/tushar-31093/aiotize_research/blob/master/SSD_Pytorch_Demo_using_YouTube_Video.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v933049X43oZ",
        "colab_type": "text"
      },
      "source": [
        "# Object Detection with SSD on Video\n",
        "\n",
        "This notebook uses a [PyTorch port](https://github.com/amdegroot/ssd.pytorch) of [SSD: Single Shot MultiBox Detector](https://arxiv.org/abs/1512.02325) to detect objects on a given youtube video.\n",
        "\n",
        "For other deep-learning Colab notebooks, visit [tugstugi/dl-colab-notebooks](https://github.com/tugstugi/dl-colab-notebooks).\n",
        "\n",
        "\n",
        "## Install amdegroot/ssd.pytorch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-AxByFTdSbcC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "from os.path import exists, join, basename, splitext\n",
        "\n",
        "git_repo_url = 'https://github.com/amdegroot/ssd.pytorch.git'\n",
        "project_name = splitext(basename(git_repo_url))[0]\n",
        "if not exists(project_name):\n",
        "  # clone and install dependencies\n",
        "  !git clone -q $git_repo_url\n",
        "  #!cd $project_name && pip install -q -r requirement.txt\n",
        "  !pip install -q youtube-dl\n",
        "  \n",
        "import sys\n",
        "sys.path.append(project_name)\n",
        "import time\n",
        "import matplotlib\n",
        "import matplotlib.pylab as plt\n",
        "plt.rcParams[\"axes.grid\"] = False\n",
        "\n",
        "from IPython.display import YouTubeVideo"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uXCFWDLm6Fdn",
        "colab_type": "text"
      },
      "source": [
        "## Download the SSD300 pretrained model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZYm1oIWUSnTg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "PRETRAINED_MODEL_URL = 'https://s3.amazonaws.com/amdegroot-models/ssd300_mAP_77.43_v2.pth'\n",
        "pretrained_model = basename(PRETRAINED_MODEL_URL)\n",
        "if not exists(pretrained_model):\n",
        "  !wget -q $PRETRAINED_MODEL_URL"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wTccKFeT6Lbr",
        "colab_type": "text"
      },
      "source": [
        "## Prepare model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0OQCE6fUS9td",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# fix some erros\n",
        "!sed -i 's/HOME = os.path.expanduser(\"~\")/HOME = \"$project_name\"/' $project_name/data/config.py\n",
        "!sed -i \"s/COCO_ROOT = osp.join(HOME, 'data\\/coco\\/')/COCO_ROOT = '$project_name\\/data'/\" $project_name/data/coco.py\n",
        "!sed -i 's/if scores.dim() == 0/if scores.size(0) == 0/' $project_name/layers/functions/detection.py\n",
        "\n",
        "import random\n",
        "import cv2\n",
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "from ssd import build_ssd\n",
        "from data import VOC_CLASSES as labels\n",
        "colors = [[random.randint(0, 255) for _ in range(3)] for _ in labels]\n",
        "\n",
        "# build SS model\n",
        "net = build_ssd('test', 300, 21)    # initialize SSD\n",
        "net.load_weights(pretrained_model)\n",
        "net.eval()\n",
        "torch.set_grad_enabled(False)\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "  net = net.cuda()\n",
        "# why?\n",
        "if torch.cuda.is_available():\n",
        "    torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "else:\n",
        "    torch.set_default_tensor_type('torch.FloatTensor')\n",
        "\n",
        "# test on a given image, see: https://github.com/amdegroot/ssd.pytorch/blob/master/demo/demo.ipynb\n",
        "def test(rgb_image, threshold=0.6):\n",
        "  # resize to 300x300\n",
        "  x = cv2.resize(rgb_image, (300, 300)).astype(np.float32)\n",
        "  # image net mean substract\n",
        "  x -= (123.0, 117.0, 104.0)\n",
        "  # image tensor\n",
        "  x = torch.from_numpy(x.astype(np.float32)).permute(2, 0, 1)\n",
        "  x = x.unsqueeze(0)\n",
        "  if torch.cuda.is_available():\n",
        "    x = x.cuda()\n",
        "  \n",
        "  # forward\n",
        "  t = time.time()\n",
        "  detections = net(x)\n",
        "  #print(\"executed in %.3fs\" % (time.time()-t))\n",
        "  \n",
        "  # draw boxes and labels\n",
        "  result_img = np.array(rgb_image)\n",
        "  scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)\n",
        "  for i in range(detections.size(1)):\n",
        "    j = 0\n",
        "    while detections[0,i,j,0] >= threshold:\n",
        "      score = detections[0,i,j,0]\n",
        "      label = labels[i-1]\n",
        "      display_txt = '%s: %.2f'%(label, score)\n",
        "      pt = (detections[0,i,j,1:]*scale).cpu().numpy()\n",
        "      color = colors[i-1]\n",
        "\n",
        "      # draw box\n",
        "      tl = round(0.002 * max(result_img.shape[0:2])) + 1  # line thickness\n",
        "      c1, c2 = (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3]))\n",
        "      cv2.rectangle(result_img, c1, c2, color, thickness=tl)\n",
        "\n",
        "      # draw text\n",
        "      tf = max(tl - 1, 1)  # font thickness\n",
        "      t_size = cv2.getTextSize(display_txt, 0, fontScale=tl / 3, thickness=tf)[0]\n",
        "      c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3\n",
        "      cv2.rectangle(result_img, c1, c2, color, -1)  # filled\n",
        "      cv2.putText(result_img, display_txt, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)\n",
        "\n",
        "      j+=1\n",
        "      \n",
        "  return result_img"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uBAuSOHQ6WUt",
        "colab_type": "text"
      },
      "source": [
        "## Detect objects on a test video\n",
        "\n",
        "We are going to detect objects on the following youtube video:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8IJtxhhWYUTo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Change this ID to anything which you want to test.\n",
        "# YOUTUBE_ID = 'bQ53oTL2AzU'\n",
        "YOUTUBE_ID = 'zc8hbSM1zVo'\n",
        "\n",
        "\n",
        "YouTubeVideo(YOUTUBE_ID)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EZuvnJcJ6ftp",
        "colab_type": "text"
      },
      "source": [
        "First, download the above youtube video:\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FnkZOPMPaHBb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -df youtube.mp4\n",
        "# download the youtube with the given ID\n",
        "!youtube-dl -f 'bestvideo[ext=mp4]' --output \"youtube.%(ext)s\" https://www.youtube.com/watch?v=$YOUTUBE_ID\n",
        "# do object detection only on the first 20 seconds...\n",
        "!ffmpeg -y -loglevel info -i youtube.mp4 -t 20 video.mp4"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h-_8nk1iB8Wg",
        "colab_type": "text"
      },
      "source": [
        "Now detect objects on the downloaded video:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HRvI5nOU97CH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "video_capture = cv2.VideoCapture()\n",
        "if video_capture.open('video.mp4'):\n",
        "  width, height = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "  fps = video_capture.get(cv2.CAP_PROP_FPS)\n",
        "  !rm -f output.mp4 output.avi\n",
        "  # can't write out mp4, so try to write into an AVI file\n",
        "  video_writer = cv2.VideoWriter(\"output.avi\", cv2.VideoWriter_fourcc(*'MJPG'), fps, (width, height))\n",
        "  while video_capture.isOpened():\n",
        "    ret, frame = video_capture.read()\n",
        "    if not ret:\n",
        "      break\n",
        "      \n",
        "    start = time.time()\n",
        "    \n",
        "    rgb_frame = test(frame[:,:,::-1])\n",
        "    frame = rgb_frame[:,:,::-1]\n",
        "\n",
        "    end = time.time()\n",
        "    print(\"time: {}s, fps: {}\".format(end-start, 1/(end-start)))\n",
        "            \n",
        "    video_writer.write(frame)\n",
        "  video_capture.release()\n",
        "  video_writer.release()\n",
        "  \n",
        "  # convert AVI to MP4\n",
        "  !ffmpeg -y -loglevel info -i output.avi output.mp4\n",
        "else:\n",
        "  print(\"can't open the given input video file!\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wd3Zn3G-CTyP",
        "colab_type": "text"
      },
      "source": [
        "Show the result video:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "auffCg9iHIeJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def show_local_mp4_video(file_name, width=640, height=480):\n",
        "  import io\n",
        "  import base64\n",
        "  from IPython.display import HTML\n",
        "  video_encoded = base64.b64encode(io.open(file_name, 'rb').read())\n",
        "  return HTML(data='''<video width=\"{0}\" height=\"{1}\" alt=\"test\" controls>\n",
        "                        <source src=\"data:video/mp4;base64,{2}\" type=\"video/mp4\" />\n",
        "                      </video>'''.format(width, height, video_encoded.decode('ascii')))\n",
        "\n",
        "show_local_mp4_video('output.mp4', width=960, height=720)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}