multiple_files_inference_nb.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mask RCNN inference\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:55:32.677631Z",
     "start_time": "2018-09-01T12:55:25.898739Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/daniel/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:34: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "/Users/daniel/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/object_detection/utils/visualization_utils.py:25: UserWarning: \n",
      "This call to matplotlib.use() has no effect because the backend has already\n",
      "been chosen; matplotlib.use() must be called *before* pylab, matplotlib.pyplot,\n",
      "or matplotlib.backends is imported for the first time.\n",
      "\n",
      "The backend was *originally* set to 'module://ipykernel.pylab.backend_inline' by the following code:\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/runpy.py\", line 193, in _run_module_as_main\n",
      "    \"__main__\", mod_spec)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/runpy.py\", line 85, in _run_code\n",
      "    exec(code, run_globals)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py\", line 16, in <module>\n",
      "    app.launch_new_instance()\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/traitlets/config/application.py\", line 658, in launch_instance\n",
      "    app.start()\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel/kernelapp.py\", line 477, in start\n",
      "    ioloop.IOLoop.instance().start()\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/zmq/eventloop/ioloop.py\", line 177, in start\n",
      "    super(ZMQIOLoop, self).start()\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/tornado/ioloop.py\", line 888, in start\n",
      "    handler_func(fd_obj, events)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/tornado/stack_context.py\", line 277, in null_wrapper\n",
      "    return fn(*args, **kwargs)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/zmq/eventloop/zmqstream.py\", line 440, in _handle_events\n",
      "    self._handle_recv()\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/zmq/eventloop/zmqstream.py\", line 472, in _handle_recv\n",
      "    self._run_callback(callback, msg)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/zmq/eventloop/zmqstream.py\", line 414, in _run_callback\n",
      "    callback(*args, **kwargs)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/tornado/stack_context.py\", line 277, in null_wrapper\n",
      "    return fn(*args, **kwargs)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel/kernelbase.py\", line 283, in dispatcher\n",
      "    return self.dispatch_shell(stream, msg)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel/kernelbase.py\", line 235, in dispatch_shell\n",
      "    handler(stream, idents, msg)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel/kernelbase.py\", line 399, in execute_request\n",
      "    user_expressions, allow_stdin)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel/ipkernel.py\", line 196, in do_execute\n",
      "    res = shell.run_cell(code, store_history=store_history, silent=silent)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/ipykernel/zmqshell.py\", line 533, in run_cell\n",
      "    return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/interactiveshell.py\", line 2698, in run_cell\n",
      "    interactivity=interactivity, compiler=compiler, result=result)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/interactiveshell.py\", line 2802, in run_ast_nodes\n",
      "    if self.run_code(code, result):\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/interactiveshell.py\", line 2862, in run_code\n",
      "    exec(code_obj, self.user_global_ns, self.user_ns)\n",
      "  File \"<ipython-input-1-fd6d40acc226>\", line 30, in <module>\n",
      "    get_ipython().magic('matplotlib inline')\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/interactiveshell.py\", line 2146, in magic\n",
      "    return self.run_line_magic(magic_name, magic_arg_s)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/interactiveshell.py\", line 2067, in run_line_magic\n",
      "    result = fn(*args,**kwargs)\n",
      "  File \"<decorator-gen-107>\", line 2, in matplotlib\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/magic.py\", line 187, in <lambda>\n",
      "    call = lambda f, *a, **k: f(*a, **k)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/magics/pylab.py\", line 99, in matplotlib\n",
      "    gui, backend = self.shell.enable_matplotlib(args.gui)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/interactiveshell.py\", line 2930, in enable_matplotlib\n",
      "    pt.activate_matplotlib(backend)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/IPython/core/pylabtools.py\", line 307, in activate_matplotlib\n",
      "    matplotlib.pyplot.switch_backend(backend)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/matplotlib/pyplot.py\", line 229, in switch_backend\n",
      "    matplotlib.use(newbackend, warn=False, force=True)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/matplotlib/__init__.py\", line 1305, in use\n",
      "    reload(sys.modules['matplotlib.backends'])\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/importlib/__init__.py\", line 166, in reload\n",
      "    _bootstrap._exec(spec, module)\n",
      "  File \"/Users/daniel/anaconda3/lib/python3.6/site-packages/matplotlib/backends/__init__.py\", line 14, in <module>\n",
      "    line for line in traceback.format_stack()\n",
      "\n",
      "\n",
      "  import matplotlib; matplotlib.use('Agg')  # pylint: disable=multiple-statements\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import six.moves.urllib as urllib\n",
    "import sys\n",
    "import tarfile\n",
    "import tensorflow as tf\n",
    "import zipfile\n",
    "\n",
    "from collections import defaultdict\n",
    "from io import StringIO\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "\n",
    "# This is needed since the notebook is stored in the object_detection folder.\n",
    "# sys.path.append(\"..\")\n",
    "from object_detection.utils import ops as utils_ops\n",
    "\n",
    "if tf.__version__ < '1.4.0':\n",
    "  raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!')\n",
    "\n",
    "from copy import copy\n",
    "import cv2\n",
    "from skimage.morphology import skeletonize\n",
    "\n",
    "import time\n",
    "import datetime\n",
    "import pickle\n",
    "\n",
    "# This is needed to display the images.\n",
    "%matplotlib inline\n",
    "\n",
    "from tensorflow.models.research.object_detection.utils import label_map_util\n",
    "from tensorflow.models.research.object_detection.utils import visualization_utils as vis_util\n",
    "\n",
    "from utility_functions import load_image_into_numpy_array\n",
    "# from utility_functions import run_inference_for_single_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model preparation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:56:23.761011Z",
     "start_time": "2018-09-01T12:56:23.755485Z"
    }
   },
   "outputs": [],
   "source": [
    "# Path to frozen detection graph.\n",
    "MODEL_NAME = 'inception_v2/fine_tuned_model_100k_final_sets'\n",
    "PATH_TO_CKPT = os.path.join('./downloaded_models', MODEL_NAME, 'frozen_inference_graph.pb')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a (frozen) Tensorflow model into memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:56:25.357383Z",
     "start_time": "2018-09-01T12:56:24.477193Z"
    }
   },
   "outputs": [],
   "source": [
    "detection_graph = tf.Graph()\n",
    "with detection_graph.as_default():\n",
    "  od_graph_def = tf.GraphDef()\n",
    "  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:\n",
    "    serialized_graph = fid.read()\n",
    "    od_graph_def.ParseFromString(serialized_graph)\n",
    "    tf.import_graph_def(od_graph_def, name='')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading label map\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:56:25.968308Z",
     "start_time": "2018-09-01T12:56:25.957198Z"
    }
   },
   "outputs": [],
   "source": [
    "PATH_TO_LABELS = '/Users/daniel/Documents/UCL/Project/Code/tensorflow_Mask_RCNN/data/worm_label_map.pbtxt'\n",
    "NUM_CLASSES = 1\n",
    "\n",
    "label_map = label_map_util.load_labelmap(PATH_TO_LABELS)\n",
    "categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)\n",
    "category_index = label_map_util.create_category_index(categories)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load images for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:56:26.286194Z",
     "start_time": "2018-09-01T12:56:26.272226Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['JU2234_worms10_food1-10_Set1_Pos4_Ch3_20102017_125033', 'N2_worms10_CSCD068947_1_Set1_Pos4_Ch4_08082017_210418', 'CX11271_worms10_food1-10_Set1_Pos4_Ch4_19052017_113042', 'JU2587_worms10_food1-10_Set1_Pos4_Ch1_20102017_125044', 'NIC199_worms10_food1-10_Set7_Pos4_Ch4_19052017_153012', 'JU2578_worms10_food1-10_Set1_Pos4_Ch4_20102017_125033', 'VC2010_worms10_food1-10_Set1_Pos4_Ch6_20122017_150107', 'ED3049_worms10_food1-10_Set6_Pos5_Ch4_19052017_151021', 'N2_worms10_CSCD068947_10_Set2_Pos5_Ch1_08082017_212337', 'N2_worms10_food1-10_Set1_Pos4_Ch5_20102017_125024', 'JU360_worms10_food1-10_Set6_Pos5_Ch6_19052017_151012', 'CB4856_worms10_food1-10_Set1_Pos4_Ch2_20102017_125044']\n"
     ]
    }
   ],
   "source": [
    "NUM_IMAGES = 5\n",
    "\n",
    "DATASET_DIR = './data/fullsize_images/'\n",
    "datasets = [f for f in os.listdir(DATASET_DIR) if not f.startswith('.')]\n",
    "print(datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:56:26.452842Z",
     "start_time": "2018-09-01T12:56:26.435985Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running inference on dataset NIC199_worms10_food1-10_Set7_Pos4_Ch4_19052017_153012\n"
     ]
    }
   ],
   "source": [
    "dataset = datasets[4]\n",
    "print(\"Running inference on dataset {}\".format(dataset))\n",
    "\n",
    "PATH_TO_TEST_IMAGES_DIR = os.path.join(DATASET_DIR, dataset)\n",
    "# PATH_TO_TEST_IMAGES_DIR = '/Users/daniel/Documents/UCL/Project/Data/annotation-data/fullsize_collated_dataset/NIC199_worms10_food1-10_Set7_Pos4_Ch4_19052017_153012/'\n",
    "\n",
    "\n",
    "FNAMES = [f for f in os.listdir(PATH_TO_TEST_IMAGES_DIR) if not f.startswith('.')]\n",
    "FNAMES = sorted(FNAMES, key=int)[:NUM_IMAGES]\n",
    "# TEST_IMAGE_PATHS = [os.path.join(PATH_TO_TEST_IMAGES_DIR, '{}/image/image_{}.png'.format(i,i)) for i in FNAMES] # np.random.choice(len(FNAMES), size=NUM_IMAGES, replace=False)]\n",
    "\n",
    "# Size of output images.\n",
    "IMAGE_SIZE = (40,40)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:56:27.217485Z",
     "start_time": "2018-09-01T12:56:27.200948Z"
    }
   },
   "outputs": [],
   "source": [
    "# def run_inference_for_single_image(image, graph):\n",
    "#   with 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",
    "#       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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T13:12:50.738334Z",
     "start_time": "2018-09-01T13:12:31.169831Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'float' object has no attribute 'seconds'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-16-e8bd99537e14>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     79\u001b[0m             \u001b[0melapsed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0melapsed\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtotal_seconds\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 81\u001b[0;31m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Inference for one image: {}.{}s\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0melapsed\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mseconds\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mround\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0melapsed\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmicroseconds\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     82\u001b[0m             \u001b[0minference_times\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0melapsed\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtotal_seconds\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     83\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'float' object has no attribute 'seconds'"
     ]
    }
   ],
   "source": [
    "output_dicts_list = []\n",
    "inference_times = []\n",
    "\n",
    "visualise_outputs = True\n",
    "save_anns_to_file = False\n",
    "save_overlays_to_file = False\n",
    "\n",
    "\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",
    "        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, 2048, 2048)\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",
    "        \n",
    "        for fName in FNAMES:\n",
    "\n",
    "            image_path = os.path.join(PATH_TO_TEST_IMAGES_DIR, '{}/image/image_{}.png'.format(fName,fName))\n",
    "\n",
    "            image = Image.open(image_path)\n",
    "            # the array based representation of the image will be used later in order to prepare the\n",
    "            # result image with boxes and labels on it.\n",
    "            image_np = load_image_into_numpy_array(image)\n",
    "\n",
    "            start = datetime.datetime.now()\n",
    "\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",
    "\n",
    "\n",
    "\n",
    "            # Actual detection.\n",
    "        #     output_dict = run_inference_for_single_image(image_np, detection_graph)\n",
    "\n",
    "            # Run inference\n",
    "            output_dict = sess.run(tensor_dict,\n",
    "                                 feed_dict={image_tensor: np.expand_dims(image_np, 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",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "            end = datetime.datetime.now()\n",
    "            elapsed = end - start\n",
    "            elapsed = elapsed.total_seconds()\n",
    "            \n",
    "            print(\"Inference for one image: {}.{}s\".format(elapsed.seconds,round(elapsed.microseconds,2)))\n",
    "            inference_times.append(elapsed.total_seconds())\n",
    "\n",
    "            if visualise_outputs:\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=1)\n",
    "                plt.figure(figsize=IMAGE_SIZE)\n",
    "                plt.title(\"Image {}\".format(image_path[:-4]))\n",
    "                plt.imshow(image_np)\n",
    "                plt.show()\n",
    "                plt.close()\n",
    "\n",
    "            # Keeping only worms which scored > 0.5\n",
    "            found_worms = np.where(output_dict['detection_scores'] > 0.5)\n",
    "            output_dict['detection_boxes'] = output_dict['detection_boxes'][found_worms]\n",
    "            output_dict['detection_classes'] = output_dict['detection_classes'][found_worms]\n",
    "            output_dict['detection_scores'] = output_dict['detection_scores'][found_worms]\n",
    "            output_dict['detection_masks'] = output_dict['detection_masks'][found_worms]\n",
    "            output_dict['skeletons'] = []\n",
    "            output_dict['frame_num'] = fName\n",
    "\n",
    "            for m in output_dict['detection_masks']:\n",
    "                output_dict['skeletons'].append(skeletonize(m).astype(np.uint8))\n",
    "\n",
    "\n",
    "            OUTPUT_DIR_PATH = os.path.join('./data/inference_outputs', MODEL_NAME, dataset)\n",
    "\n",
    "\n",
    "            if save_anns_to_file:\n",
    "                #Save outputs to Pickle file\n",
    "                os.makedirs(os.path.join(OUTPUT_DIR_PATH, 'annotations'), exist_ok=True)\n",
    "                ANNS_OUTPUT_PATH = os.path.join(OUTPUT_DIR_PATH,'annotations', fName) + '.pickle'\n",
    "                with open(ANNS_OUTPUT_PATH, 'wb') as fp:\n",
    "                    pickle.dump(output_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "            if save_overlays_to_file:\n",
    "                #save image with annotations overlaid to file\n",
    "                os.makedirs(os.path.join(OUTPUT_DIR_PATH, 'images'), exist_ok=True)\n",
    "                IMG_OUTPUT_PATH = os.path.join(OUTPUT_DIR_PATH,'images', fName) + '.png'\n",
    "\n",
    "                plt.figure(figsize=IMAGE_SIZE)\n",
    "\n",
    "                # If the image ahsn't already been visualised, we need \n",
    "                # to add the masks and boxes now\n",
    "                if not visualise_outputs:\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=1)\n",
    "\n",
    "                plt.imshow(image_np)\n",
    "                plt.axis('off')\n",
    "                plt.savefig(fname=IMG_OUTPUT_PATH, bbox_inches='tight', pad_inches=0)\n",
    "                plt.close\n",
    "\n",
    "\n",
    "            output_dicts_list.append(output_dict)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T13:12:50.739008Z",
     "start_time": "2018-09-01T13:12:32.221Z"
    }
   },
   "outputs": [],
   "source": [
    "mean_time = round(sum(inference_times) / len(inference_times), 2)\n",
    "print(\"Average inference time for {} images: {}s\".format(NUM_IMAGES, mean_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Skeletonisation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualise the last image and its skeletons as a sense-check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:55:33.028822Z",
     "start_time": "2018-09-01T12:55:25.927Z"
    }
   },
   "outputs": [],
   "source": [
    "skeletons = np.array(output_dicts_list[-1]['skeletons'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:55:33.029782Z",
     "start_time": "2018-09-01T12:55:25.931Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(80,80))\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(image_np)\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(np.sum(skeletons, axis=0))\n",
    "\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tracking using SORT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:55:33.031101Z",
     "start_time": "2018-09-01T12:55:25.935Z"
    }
   },
   "outputs": [],
   "source": [
    "from sort import sort\n",
    "import matplotlib.patches as patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-09-01T12:55:33.032227Z",
     "start_time": "2018-09-01T12:55:25.939Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "display = True\n",
    "verbose = False\n",
    "colours = np.random.rand(32,3)\n",
    "\n",
    "\n",
    "\n",
    "#create instance of SORT\n",
    "worm_tracker = sort.Sort(max_age=5, min_hits=0) \n",
    "\n",
    "# get detections\n",
    "all_detections = [np.hstack((i['detection_boxes']*2048, np.expand_dims(i['detection_scores'],1))) for i in output_dicts_list]\n",
    "\n",
    "\n",
    "for frame in range(len(all_detections)):\n",
    "        \n",
    "    \n",
    "    # update SORT\n",
    "    detections = all_detections[frame]\n",
    "\n",
    "    \n",
    "    ids = worm_tracker.update(detections)\n",
    "    \n",
    "    if verbose:\n",
    "        print('Detections: {}'.format(detections))\n",
    "        print('ids shape: {}'.format(ids.shape))\n",
    "        print('ids: {}'.format(ids))\n",
    "    \n",
    "    print(\"Worm IDs: {}\".format(sorted(ids[:,-1], key=int)))\n",
    "    \n",
    "    if(display):\n",
    "        \n",
    "        fName = FNAMES[frame]\n",
    "    \n",
    "        image_path = os.path.join(PATH_TO_TEST_IMAGES_DIR, '{}/image/image_{}.png'.format(fName,fName))\n",
    "\n",
    "        img = Image.open(image_path)\n",
    "        img_np = load_image_into_numpy_array(img)\n",
    "        \n",
    "        fig = plt.figure(figsize=(IMAGE_SIZE))\n",
    "        \n",
    "        \n",
    "        ax = fig.add_subplot(121, aspect='equal')\n",
    "        ax.imshow(img_np)\n",
    "\n",
    "        for d in ids:\n",
    "            d = d.astype(np.int32)\n",
    "            ax.add_patch(patches.Rectangle((d[1],d[0]),d[3]-d[1],d[2]-d[0],\n",
    "                                           fill=False,\n",
    "                                           lw=3,\n",
    "                                           ec=colours[d[4]%32,:]))\n",
    "            ax.set_adjustable('box-forced')\n",
    "            plt.text(d[1],d[0],d[-1],color='w')\n",
    "        \n",
    "        ax_2 = fig.add_subplot(122, aspect='equal')\n",
    "        vis_util.visualize_boxes_and_labels_on_image_array(\n",
    "              img_np,\n",
    "              output_dicts_list[frame]['detection_boxes'],\n",
    "              output_dicts_list[frame]['detection_classes'],\n",
    "              output_dicts_list[frame]['detection_scores'],\n",
    "              category_index,\n",
    "              instance_masks=output_dicts_list[frame].get('detection_masks'),\n",
    "              use_normalized_coordinates=True,\n",
    "              line_thickness=2)\n",
    "        ax_2.imshow(img_np)\n",
    "        \n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "        plt.close\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-08-10T13:47:06.256080Z",
     "start_time": "2018-08-10T13:47:06.245922Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "version": "0.3.2"
  },
  "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}