test_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Object Detection From TF2 Saved Model"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "This demo will take you through the steps of running an \"out-of-the-box\" TensorFlow 2 compatible detection model on a collection of images.\r\n",
    "It is based on tutorial at https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/auto_examples/plot_object_detection_saved_model.html#sphx-glr-auto-examples-plot-object-detection-saved-model-py\r\n",
    "\r\n",
    "It is simpified to use a local model previously trained and exported with the `train_model.ipynb`-notebook."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Options"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Test Images\r\n",
    "Glob all files from a given directory."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "IMAGE_DIRECTORY = \"test_images\"\r\n",
    "IMAGE_PATHS = [os.path.join(IMAGE_DIRECTORY, x) for x in os.listdir(IMAGE_DIRECTORY)]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "OR use some images from the training set..."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "source": [
    "IMAGE_PATHS = [\"DoorDetect-Dataset/images/0bf01ee25888cb45.jpg\", \"DoorDetect-Dataset/images/11739cf765a3f3d5.jpg\"]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Model and Labels"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "source": [
    "PATH_TO_MODEL_DIR = \"exported_models/ssd_mobilenet_v2_fpnlite_320x320\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "PATH_TO_LABELS = \"labelmap.pbtxt\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load the model\r\n",
    "Next we load the exported model."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import time\r\n",
    "from object_detection.utils import label_map_util\r\n",
    "from object_detection.utils import visualization_utils as viz_utils\r\n",
    "\r\n",
    "PATH_TO_SAVED_MODEL = PATH_TO_MODEL_DIR + \"/saved_model\"\r\n",
    "\r\n",
    "print('Loading model...', end='')\r\n",
    "start_time = time.time()\r\n",
    "\r\n",
    "# Load saved model and build the detection function\r\n",
    "detect_fn = tf.saved_model.load(PATH_TO_SAVED_MODEL)\r\n",
    "\r\n",
    "end_time = time.time()\r\n",
    "elapsed_time = end_time - start_time\r\n",
    "print('Done! Took {} seconds'.format(elapsed_time))"
   ],
   "outputs": [],
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load label map data (for plotting)\r\n",
    "Label maps correspond index numbers to category names, so that when our convolution network\r\n",
    "predicts `1`, we know that this corresponds to `door`.  Here we use internal utility\r\n",
    "functions, but anything that returns a dictionary mapping integers to appropriate string labels\r\n",
    "would be fine."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "source": [
    "category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS,\r\n",
    "                                                                    use_display_name=True)"
   ],
   "outputs": [],
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Putting everything together\n",
    "The code shown below loads an image, runs it through the detection model and visualizes the\n",
    "detection results, including the keypoints.\n",
    "\n",
    "Note that this will take a long time (several minutes) the first time you run this code due to\n",
    "tf.function's trace-compilation --- on subsequent runs (e.g. on new images), things will be\n",
    "faster.\n",
    "\n",
    "Here are some simple things to try out if you are curious:\n",
    "\n",
    "* Modify some of the input images and see if detection still works. Some simple things to try out here (just uncomment the relevant portions of code) include flipping the image horizontally, or converting to grayscale (note that we still expect the input image to have 3 channels).\n",
    "* Print out `detections['detection_boxes']` and try to match the box locations to the boxes in the image.  Notice that coordinates are given in normalized form (i.e., in the interval [0, 1]).\n",
    "* Set ``min_score_thresh`` to other values (between 0 and 1) to allow more detections in or to filter out more detections.\n",
    "\n"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%matplotlib inline\r\n",
    "import numpy as np\r\n",
    "from PIL import Image\r\n",
    "import matplotlib.pyplot as plt\r\n",
    "import warnings\r\n",
    "warnings.filterwarnings('ignore')   # Suppress Matplotlib warnings\r\n",
    "\r\n",
    "def load_image_into_numpy_array(path):\r\n",
    "    \"\"\"Load an image from file into a numpy array.\r\n",
    "\r\n",
    "    Puts image into numpy array to feed into tensorflow graph.\r\n",
    "    Note that by convention we put it into a numpy array with shape\r\n",
    "    (height, width, channels), where channels=3 for RGB.\r\n",
    "\r\n",
    "    Args:\r\n",
    "      path: the file path to the image\r\n",
    "\r\n",
    "    Returns:\r\n",
    "      uint8 numpy array with shape (img_height, img_width, 3)\r\n",
    "    \"\"\"\r\n",
    "    return np.array(Image.open(path))\r\n",
    "\r\n",
    "\r\n",
    "for image_path in IMAGE_PATHS:\r\n",
    "\r\n",
    "    print('Running inference for {}... '.format(image_path), end='')\r\n",
    "\r\n",
    "    image_np = load_image_into_numpy_array(image_path)\r\n",
    "\r\n",
    "    # Things to try:\r\n",
    "    # Flip horizontally\r\n",
    "    # image_np = np.fliplr(image_np).copy()\r\n",
    "\r\n",
    "    # Convert image to grayscale\r\n",
    "    # image_np = np.tile(\r\n",
    "    #     np.mean(image_np, 2, keepdims=True), (1, 1, 3)).astype(np.uint8)\r\n",
    "\r\n",
    "    # The input needs to be a tensor, convert it using `tf.convert_to_tensor`.\r\n",
    "    input_tensor = tf.convert_to_tensor(image_np)\r\n",
    "    # The model expects a batch of images, so add an axis with `tf.newaxis`.\r\n",
    "    input_tensor = input_tensor[tf.newaxis, ...]\r\n",
    "\r\n",
    "    # input_tensor = np.expand_dims(image_np, 0)\r\n",
    "    detections = detect_fn(input_tensor)\r\n",
    "\r\n",
    "    # All outputs are batches tensors.\r\n",
    "    # Convert to numpy arrays, and take index [0] to remove the batch dimension.\r\n",
    "    # We're only interested in the first num_detections.\r\n",
    "    num_detections = int(detections.pop('num_detections'))\r\n",
    "    detections = {key: value[0, :num_detections].numpy()\r\n",
    "                   for key, value in detections.items()}\r\n",
    "    detections['num_detections'] = num_detections\r\n",
    "\r\n",
    "    # detection_classes should be ints.\r\n",
    "    detections['detection_classes'] = detections['detection_classes'].astype(np.int64)\r\n",
    "    #print(detections)\r\n",
    "\r\n",
    "    image_np_with_detections = image_np.copy()\r\n",
    "\r\n",
    "    viz_utils.visualize_boxes_and_labels_on_image_array(\r\n",
    "          image_np_with_detections,\r\n",
    "          detections['detection_boxes'],\r\n",
    "          detections['detection_classes'],\r\n",
    "          detections['detection_scores'],\r\n",
    "          category_index,\r\n",
    "          use_normalized_coordinates=True,\r\n",
    "          max_boxes_to_draw=200,\r\n",
    "          min_score_thresh=.2,\r\n",
    "          agnostic_mode=False)\r\n",
    "\r\n",
    "    plt.figure()\r\n",
    "    plt.imshow(image_np_with_detections)\r\n",
    "    print('Done')\r\n",
    "plt.show()\r\n",
    "\r\n",
    "# sphinx_gallery_thumbnail_number = 2"
   ],
   "outputs": [],
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}