Tensorflow Object Detection/object_detection_with_own_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Using your own object detector for detection images"
   ],
   "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/object_detection_with_own_model.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/object_detection_with_own_model.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 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": [],
   "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\n",
    "\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.  "
   ],
   "metadata": {
    "colab_type": "text",
    "id": "X_sEBLpVWKMQ"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "MODEL_NAME = 'inference_graph'\r\n",
    "PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'\r\n",
    "PATH_TO_LABELS = 'training/labelmap.pbtxt'"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load a (frozen) Tensorflow model into memory."
   ],
   "metadata": {
    "colab_type": "text",
    "id": "YBcB9QHLWKMU"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": 6,
   "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": 8,
   "source": [
    "def run_inference_for_single_image(image, graph):\n",
    "    if 'detection_masks' in tensor_dict:\n",
    "        # The following processing is only for single image\n",
    "        detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])\n",
    "        detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])\n",
    "        # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.\n",
    "        real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)\n",
    "        detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])\n",
    "        detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])\n",
    "        detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(\n",
    "            detection_masks, detection_boxes, image.shape[0], image.shape[1])\n",
    "        detection_masks_reframed = tf.cast(\n",
    "            tf.greater(detection_masks_reframed, 0.5), tf.uint8)\n",
    "        # Follow the convention by adding back the batch dimension\n",
    "        tensor_dict['detection_masks'] = tf.expand_dims(\n",
    "            detection_masks_reframed, 0)\n",
    "    image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')\n",
    "\n",
    "    # Run inference\n",
    "    output_dict = sess.run(tensor_dict,\n",
    "                            feed_dict={image_tensor: np.expand_dims(image, 0)})\n",
    "\n",
    "    # all outputs are float32 numpy arrays, so convert types as appropriate\n",
    "    output_dict['num_detections'] = int(output_dict['num_detections'][0])\n",
    "    output_dict['detection_classes'] = output_dict[\n",
    "        'detection_classes'][0].astype(np.uint8)\n",
    "    output_dict['detection_boxes'] = output_dict['detection_boxes'][0]\n",
    "    output_dict['detection_scores'] = output_dict['detection_scores'][0]\n",
    "    if 'detection_masks' in output_dict:\n",
    "        output_dict['detection_masks'] = output_dict['detection_masks'][0]\n",
    "    return output_dict"
   ],
   "outputs": [],
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "92BHxzcNWKMf"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import cv2\n",
    "cap = cv2.VideoCapture(0)\n",
    "try:\n",
    "    with detection_graph.as_default():\n",
    "        with tf.Session() as sess:\n",
    "                # Get handles to input and output tensors\n",
    "                ops = tf.get_default_graph().get_operations()\n",
    "                all_tensor_names = {output.name for op in ops for output in op.outputs}\n",
    "                tensor_dict = {}\n",
    "                for key in [\n",
    "                  'num_detections', 'detection_boxes', 'detection_scores',\n",
    "                  'detection_classes', 'detection_masks'\n",
    "                ]:\n",
    "                    tensor_name = key + ':0'\n",
    "                    if tensor_name in all_tensor_names:\n",
    "                        tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(\n",
    "                      tensor_name)\n",
    "\n",
    "                while True:\n",
    "                    ret, image_np = cap.read()\n",
    "                    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]\n",
    "                    image_np_expanded = np.expand_dims(image_np, axis=0)\n",
    "                    # Actual detection.\n",
    "                    output_dict = run_inference_for_single_image(image_np, detection_graph)\n",
    "                    # Visualization of the results of a detection.\n",
    "                    vis_util.visualize_boxes_and_labels_on_image_array(\n",
    "                        image_np,\n",
    "                        output_dict['detection_boxes'],\n",
    "                        output_dict['detection_classes'],\n",
    "                        output_dict['detection_scores'],\n",
    "                        category_index,\n",
    "                        instance_masks=output_dict.get('detection_masks'),\n",
    "                        use_normalized_coordinates=True,\n",
    "                        line_thickness=8)\n",
    "                    cv2.imshow('object_detection', cv2.resize(image_np, (800, 600)))\n",
    "                    if cv2.waitKey(25) & 0xFF == ord('q'):\n",
    "                        cap.release()\n",
    "                        cv2.destroyAllWindows()\n",
    "                        break\n",
    "except Exception as e:\n",
    "    print(e)\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
}