Tensorflow Object Detection/detect_object_in_webcam_video.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Detection Objects in a webcam image stream"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "V8-yl-s-WKMG"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "<table align=\"left\"><td>\n",
    "  <a target=\"_blank\"  href=\"https://colab.research.google.com/github/TannerGilbert/Tutorials/blob/master/Tensorflow%20Object%20Detection/detect_object_in_webcam_video.ipynb\">\n",
    "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab\n",
    "  </a>\n",
    "</td><td>\n",
    "  <a target=\"_blank\"  href=\"https://github.com/TannerGilbert/Tutorials/blob/master/Tensorflow%20Object%20Detection/detect_object_in_webcam_video.ipynb\">\n",
    "    <img width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "</td></table>"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "This notebook is based on the [official Tensorflow Object Detection demo](https://github.com/tensorflow/models/blob/r1.13.0/research/object_detection/object_detection_tutorial.ipynb) and only contains some slight changes. Make sure to follow the [installation instructions](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1.md#installation) before you start."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Imports"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "kFSqkTCdWKMI"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import numpy as np\r\n",
    "import os\r\n",
    "import six.moves.urllib as urllib\r\n",
    "import sys\r\n",
    "import tarfile\r\n",
    "import tensorflow as tf\r\n",
    "from distutils.version import StrictVersion\r\n",
    "\r\n",
    "# This is needed since the notebook is stored in the object_detection folder.\r\n",
    "sys.path.append(\"..\")\r\n",
    "from object_detection.utils import ops as utils_ops\r\n",
    "\r\n",
    "if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):\r\n",
    "  raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "hV4P5gyTWKMI"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Env setup"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "Wy72mWwAWKMK"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "# This is needed to display the images.\r\n",
    "%matplotlib inline"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "v7m_NY_aWKMK"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Object detection imports\n",
    "Here are the imports from the object detection module."
   ],
   "metadata": {
    "colab_type": "text",
    "id": "r5FNuiRPWKMN"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "from utils import label_map_util\r\n",
    "\r\n",
    "from utils import visualization_utils as vis_util"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "C:\\Users\\Gilbert\\Downloads\\models\\research\\object_detection\\utils\\visualization_utils.py:26: UserWarning: matplotlib.pyplot as already been imported, this call will have no effect.\n",
      "  import matplotlib; matplotlib.use('Agg')  # pylint: disable=multiple-statements\n"
     ]
    }
   ],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "bm0_uNRnWKMN"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Model preparation "
   ],
   "metadata": {
    "colab_type": "text",
    "id": "cfn_tRFOWKMO"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Variables\r\n",
    "\r\n",
    "Any model exported using the `export_inference_graph.py` tool can be loaded here simply by changing `PATH_TO_FROZEN_GRAPH` to point to a new .pb file.  \r\n",
    "\r\n",
    "By default we use an \"SSD with Mobilenet\" model here. See the [detection model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md) for a list of other models that can be run out-of-the-box with varying speeds and accuracies."
   ],
   "metadata": {
    "colab_type": "text",
    "id": "X_sEBLpVWKMQ"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "# What model to download.\r\n",
    "MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'\r\n",
    "MODEL_FILE = MODEL_NAME + '.tar.gz'\r\n",
    "DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'\r\n",
    "\r\n",
    "# Path to frozen detection graph. This is the actual model that is used for the object detection.\r\n",
    "PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'\r\n",
    "\r\n",
    "# List of the strings that is used to add correct label for each box.\r\n",
    "PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "VyPz_t8WWKMQ"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Download Model"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "7ai8pLZZWKMS"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "opener = urllib.request.URLopener()\r\n",
    "opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)\r\n",
    "tar_file = tarfile.open(MODEL_FILE)\r\n",
    "for file in tar_file.getmembers():\r\n",
    "  file_name = os.path.basename(file.name)\r\n",
    "  if 'frozen_inference_graph.pb' in file_name:\r\n",
    "    tar_file.extract(file, os.getcwd())"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "KILYnwR5WKMS"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load a (frozen) Tensorflow model into memory."
   ],
   "metadata": {
    "colab_type": "text",
    "id": "YBcB9QHLWKMU"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "detection_graph = tf.Graph()\r\n",
    "with detection_graph.as_default():\r\n",
    "  od_graph_def = tf.GraphDef()\r\n",
    "  with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:\r\n",
    "    serialized_graph = fid.read()\r\n",
    "    od_graph_def.ParseFromString(serialized_graph)\r\n",
    "    tf.import_graph_def(od_graph_def, name='')"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "KezjCRVvWKMV"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Loading label map\n",
    "Label maps map indices to category names, so that when our convolution network predicts `5`, we know that this corresponds to `airplane`.  Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Detection"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "def run_inference_for_single_image(image, graph):\r\n",
    "    if 'detection_masks' in tensor_dict:\r\n",
    "        # The following processing is only for single image\r\n",
    "        detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])\r\n",
    "        detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])\r\n",
    "        # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.\r\n",
    "        real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)\r\n",
    "        detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])\r\n",
    "        detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])\r\n",
    "        detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(\r\n",
    "            detection_masks, detection_boxes, image.shape[0], image.shape[1])\r\n",
    "        detection_masks_reframed = tf.cast(\r\n",
    "            tf.greater(detection_masks_reframed, 0.5), tf.uint8)\r\n",
    "        # Follow the convention by adding back the batch dimension\r\n",
    "        tensor_dict['detection_masks'] = tf.expand_dims(\r\n",
    "            detection_masks_reframed, 0)\r\n",
    "    image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')\r\n",
    "\r\n",
    "    # Run inference\r\n",
    "    output_dict = sess.run(tensor_dict,\r\n",
    "                            feed_dict={image_tensor: np.expand_dims(image, 0)})\r\n",
    "\r\n",
    "    # all outputs are float32 numpy arrays, so convert types as appropriate\r\n",
    "    output_dict['num_detections'] = int(output_dict['num_detections'][0])\r\n",
    "    output_dict['detection_classes'] = output_dict[\r\n",
    "        'detection_classes'][0].astype(np.uint8)\r\n",
    "    output_dict['detection_boxes'] = output_dict['detection_boxes'][0]\r\n",
    "    output_dict['detection_scores'] = output_dict['detection_scores'][0]\r\n",
    "    if 'detection_masks' in output_dict:\r\n",
    "        output_dict['detection_masks'] = output_dict['detection_masks'][0]\r\n",
    "    return output_dict"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "92BHxzcNWKMf"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "import cv2\r\n",
    "cap = cv2.VideoCapture(0)\r\n",
    "try:\r\n",
    "    with detection_graph.as_default():\r\n",
    "        with tf.Session() as sess:\r\n",
    "                # Get handles to input and output tensors\r\n",
    "                ops = tf.get_default_graph().get_operations()\r\n",
    "                all_tensor_names = {output.name for op in ops for output in op.outputs}\r\n",
    "                tensor_dict = {}\r\n",
    "                for key in [\r\n",
    "                  'num_detections', 'detection_boxes', 'detection_scores',\r\n",
    "                  'detection_classes', 'detection_masks'\r\n",
    "                ]:\r\n",
    "                    tensor_name = key + ':0'\r\n",
    "                    if tensor_name in all_tensor_names:\r\n",
    "                        tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(\r\n",
    "                      tensor_name)\r\n",
    "\r\n",
    "                while True:\r\n",
    "                    ret, image_np = cap.read()\r\n",
    "                    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]\r\n",
    "                    image_np_expanded = np.expand_dims(image_np, axis=0)\r\n",
    "                    # Actual detection.\r\n",
    "                    output_dict = run_inference_for_single_image(image_np, detection_graph)\r\n",
    "                    # Visualization of the results of a detection.\r\n",
    "                    vis_util.visualize_boxes_and_labels_on_image_array(\r\n",
    "                        image_np,\r\n",
    "                        output_dict['detection_boxes'],\r\n",
    "                        output_dict['detection_classes'],\r\n",
    "                        output_dict['detection_scores'],\r\n",
    "                        category_index,\r\n",
    "                        instance_masks=output_dict.get('detection_masks'),\r\n",
    "                        use_normalized_coordinates=True,\r\n",
    "                        line_thickness=8)\r\n",
    "                    cv2.imshow('object_detection', cv2.resize(image_np, (800, 600)))\r\n",
    "                    if cv2.waitKey(25) & 0xFF == ord('q'):\r\n",
    "                        cap.release()\r\n",
    "                        cv2.destroyAllWindows()\r\n",
    "                        break\r\n",
    "except Exception as e:\r\n",
    "    print(e)\r\n",
    "    cap.release()"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "3a5wMHN8WKMh"
   }
  }
 ],
 "metadata": {
  "colab": {
   "default_view": {},
   "name": "object_detection_tutorial.ipynb?workspaceId=ronnyvotel:python_inference::citc",
   "provenance": [],
   "version": "0.3.2",
   "views": {}
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}