Video_Object_Detection_Project.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Video Object Detection Project",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AlzIlBsScJJ_"
      },
      "source": [
        "# Video Object Detection Project\n",
        "by Amanda Ma"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nTirVS4FWaPx"
      },
      "source": [
        "This project uses [OpenCV](https://opencv.org/) for video and image \n",
        "processing, as well as [SSD MobileNet V1 Coco](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md), a trained tensorflow model from [TensorflowZoo](https://github.com/tensorflow/models), to identify objects from videos. Pictures and videos are obtained from [pixabay](https://pixabay.com/).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "En2ZrxY7qcBz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import cv2 as cv"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "48e3yFraQWfm",
        "colab_type": "text"
      },
      "source": [
        "#### Load the Labels\n",
        "The object labels are stored in a text file with the following format: <br>\n",
        "0: unlabeled, 1: person, 2: bicycle ...... 75: remote, 76: keyboard ...... 180: window-blind, 181: window-other, 182: wood <br> <br>\n",
        "Objects and their associated label numbers are stored in the dictionary _labels_ for future lookup. (key = label number, value = label)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JmgGAtEToKEt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "labels = {}\n",
        "with open('labels.txt', 'r') as labelsFile:\n",
        "  for line in labelsFile.readlines():\n",
        "    split = line.strip().split(': ')\n",
        "    labels[int(split[0])] = split[1]\n",
        "  print(labels)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AioBgivbQnCP",
        "colab_type": "text"
      },
      "source": [
        "#### Obtain the Model File"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aZ-5x4NrQriv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import urllib.request\n",
        "import os"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cwjrAyAZcWH5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "base_url = 'http://download.tensorflow.org/models/object_detection/'\n",
        "file_name = 'ssd_mobilenet_v1_coco_2018_01_28.tar.gz'\n",
        "url = base_url + file_name\n",
        "urllib.request.urlretrieve(url, file_name)\n",
        "os.listdir()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CPPjRljFQoYl",
        "colab_type": "text"
      },
      "source": [
        "#### Extract the Model Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LIgJzPhSQwmW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tarfile\n",
        "import shutil"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B60xkKQ7O5D2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "dir_name = file_name[0:-len('.tar.gz')]\n",
        "if os.path.exists(dir_name):\n",
        "  shutil.rmtree(dir_name) \n",
        "tarfile.open(file_name, 'r:gz').extractall('./')\n",
        "os.listdir(dir_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xvEGo39oRDlY",
        "colab_type": "text"
      },
      "source": [
        "#### Load the Model\n",
        "The model file is read in binary format and parsed into a GraphDef object"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cbLRP6_skOom",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "frozen_graph = os.path.join(dir_name, 'frozen_inference_graph.pb')\n",
        "with tf.gfile.FastGFile(frozen_graph,'rb') as f:\n",
        "  graph_def = tf.GraphDef()\n",
        "  graph_def.ParseFromString(f.read());"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6sdJkU-bRK4Y",
        "colab_type": "text"
      },
      "source": [
        "#### Video Processing and Frame Labeling"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fOnfju1ER1ij",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from PIL import Image"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rZKm8-7gj_rD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def process_video(input_filename, output_filename, frames): \n",
        "  \"\"\" Takes an input video and creates an output video with labels\n",
        "      on various objects detected in the video. \n",
        "  \"\"\"\n",
        "  assert type(input_filename) is str \n",
        "  assert type(output_filename) is str\n",
        "  assert type(frames) is int \n",
        "  \n",
        "  input_video = cv.VideoCapture(input_filename)\n",
        "  height = int(vid.get(cv.CAP_PROP_FRAME_HEIGHT))\n",
        "  width = int(vid.get(cv.CAP_PROP_FRAME_WIDTH))\n",
        "  fps = int(vid.get(cv.CAP_PROP_FPS))\n",
        "  total_frames = int(vid.get(cv.CAP_PROP_FRAME_COUNT))\n",
        "  \n",
        "  fourcc = cv.VideoWriter_fourcc(*'MP4V')\n",
        "  output = cv.VideoWriter(output_filename, fourcc, fps, (width, height))\n",
        "  \n",
        "  with tf.Session() as sess: \n",
        "    sess.graph.as_default()\n",
        "    tf.import_graph_def(graph_def, name = '')\n",
        "    \n",
        "    for i in range(0, total_frames, frames): \n",
        "      input_video.set(cv.CAP_PROP_POS_FRAMES, i)\n",
        "      ret, frame = input_video.read()\n",
        "      if not ret: \n",
        "        raise Exception(\"Problem reading frame\", i, \" from video\")\n",
        "      else: \n",
        "        labeled_frame = labelFrame(frame, sess)\n",
        "        output.write(labeled_frame)\n",
        "        \n",
        "  input_video.release()\n",
        "  output.release()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w4R2ycYIjNub",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def labelFrame(frame, sess): \n",
        "  \"\"\" Takes a single frame and draws bounding boxes with labels around \n",
        "      detected objects (in the context of the tensorflow session).\n",
        "  \"\"\"\n",
        "  frame_scaled = cv.resize(frame, (300, 300))\n",
        "  \n",
        "  detections = sess.run(\n",
        "                [\n",
        "                  sess.graph.get_tensor_by_name('num_detections:0'), \n",
        "                  sess.graph.get_tensor_by_name('detection_scores:0'), \n",
        "                  sess.graph.get_tensor_by_name('detection_boxes:0'), \n",
        "                  sess.graph.get_tensor_by_name('detection_classes:0'),\n",
        "                ],\n",
        "                  feed_dict = {\n",
        "                    'image_tensor:0': frame_scaled.reshape(1, \n",
        "                      frame_scaled.shape[0], frame_scaled.shape[1], 3)\n",
        "                  }\n",
        "               )\n",
        "  \n",
        "  num_detections = int(detections[0][0])\n",
        "  \n",
        "  height, width = frame.shape[0], frame.shape[1]\n",
        "  r_frame, g_frame, b_frame = 0, 0, 255\n",
        "  r_text, g_text, b_text = 255, 255, 255\n",
        "  scale = 0.8\n",
        "  \n",
        "  threshold = 0.4 \n",
        "  \n",
        "  for i in range(num_detections): \n",
        "    confidence = detections[1][0][i]\n",
        "    if confidence > threshold: \n",
        "      box = detections[2][0][i]\n",
        "      left_h, left_w = int(box[0]*height), int(box[1]*width)\n",
        "      right_h, right_w = int(box[2]*height), int(box[3]*width)\n",
        "      \n",
        "      cv.rectangle(frame, (left_w, left_h), (right_w, right_h), (r_frame, g_frame, b_frame), 2)\n",
        "      cv.rectangle(frame, (left_w, left_h), (right_w, left_h - 35), (r_frame, g_frame, b_frame), cv.FILLED)\n",
        "      \n",
        "      label_id = int(detections[3][0][i])\n",
        "      label = labels[label_id]\n",
        "      \n",
        "      cv.putText(frame, label, (left_w, left_h - 3), cv.FONT_HERSHEY_SIMPLEX, scale, [r_text, g_text, b_text], 2\n",
        "                )\n",
        "    else: \n",
        "      break \n",
        "      \n",
        "  return frame "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aCuZXx8CRW9H",
        "colab_type": "text"
      },
      "source": [
        "#### Video Samples\n",
        "Here are a few video samples uploaded on Vimeo. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "afVMCYihOiEm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from IPython.display import IFrame"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cJHIXr0Kcenu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "process_video('cars.mp4', 'cars_labeled.mp4', 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q0hHXHsAD9t0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "IFrame('https://player.vimeo.com/video/357028716', width = '660', height = '515')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xif-czZ_BRjN",
        "colab_type": "text"
      },
      "source": [
        "Here, the labels are harder to see. The dancers are classified as 'person's."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P9Im_pphBALS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "process_video('dance.mp4', 'dance_labeled.mp4', 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yzBztnpmBBol",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "IFrame('https://player.vimeo.com/video/357039770', width = '660', height = '515')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wxcxeq_uSmKm",
        "colab_type": "text"
      },
      "source": [
        "Not all labels are perfect! In the labeled video below, some animals were mistakenly labeled as dogs or sports balls."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2kQifV4vOtVn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "process_video('rooster.mp4', 'roster_labeled.mp4', 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "If4gWZsKQZMJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "IFrame('https://player.vimeo.com/video/357029971', width = '660', height = '515')"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}