{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference on pretrained SSD model using Tensor Cores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will show, how to run an inference with our SSD implementation.\n",
    "\n",
    "We will start with defining input pipeline, then we will see how to load the model and then we will run an inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading an image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets import libraries we will use to prepare an input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import torch\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From our examples we can import utility functions for inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dle.inference import load_image, rescale, crop_center, normalize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can load an example image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f4d70aad710>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 234600 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = load_image('http://images.cocodataset.org/val2017/000000397133.jpg')\n",
    "plt.imshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will rescale it, crop it and normalize it, so the model will get the expected input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = rescale(img, 300, 300)\n",
    "img = crop_center(img, 300, 300)\n",
    "img = normalize(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can present the image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f4d70419e10>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 49624 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I looks weird, because after normalization, data values are in range [-1..1]. Plotting lib expects values from [0..1] range. We can fix it for visualization purpose:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(300, 300, 3)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 150452 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "out = img/2+0.5\n",
    "plt.imshow(out)\n",
    "img.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building an predictor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have prepared our imput. Next thing is to load a SSD model.\n",
    "In our examples you can find some framework specific functions. Some of them will be explained here in detail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/resnet50-19c8e357.pth\" to /root/.torch/models/resnet50-19c8e357.pth\n",
      "102502400it [00:09, 10362562.63it/s]\n"
     ]
    }
   ],
   "source": [
    "from examples.SSD300_inference import load_checkpoint, build_predictor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can import the model. We need to set it in the evaluation mode also:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from apex.fp16_utils import network_to_half\n",
    "\n",
    "ssd300 = build_predictor('/checkpoints/SSD300v1.1.pt')\n",
    "ssd300 = ssd300.cuda()\n",
    "ssd300 = network_to_half(ssd300.cuda())\n",
    "ssd300 = ssd300.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model does not expect input as a `ndarray`. It prefers Pytorch Tensor data format. It also expects that input will be a batch of several images. What is more, it expects input in a bit different shape than usual.\n",
    "We can fulfill these requirements with following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 300, 300])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# change the shape\n",
    "HWC = img\n",
    "CHW = np.swapaxes(np.swapaxes(HWC, 0, 2), 1, 2)\n",
    "# make a batch of 1 image\n",
    "batch = np.expand_dims(CHW, axis=0)\n",
    "# turn input into tensor\n",
    "tensor = torch.from_numpy(batch)\n",
    "tensor = tensor.cuda()\n",
    "tensor = tensor.half()\n",
    "tensor.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can make some prediction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = ssd300(tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, the output from the model is not too easy to read. To present it in human-readable form there are a few steps missing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At first, more imports..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ssd.utils import dboxes300_coco, Encoder\n",
    "import matplotlib.patches as patches\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which allows us to decode the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "dboxes = dboxes300_coco()\n",
    "encoder = Encoder(dboxes)\n",
    "ploc, plabel = [val.float() for val in prediction]\n",
    "encoded = encoder.decode_batch(ploc, plabel, criteria=0.5, max_output=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `criteria` param allows to filter results with IoU not lower than the `criteria`.\n",
    "\n",
    "Encoder returns a batch of results in a form:\n",
    "```\n",
    "[ fst img: prediction,\n",
    "  snd_img: prediction,\n",
    "  ...\n",
    "]\n",
    "```\n",
    "While the prediction is:\n",
    "```\n",
    "( bounding boxes: [ fst detection: [x1, y1, x2, y2], snd detection ... ],\n",
    "  classes:        [ fst detextion: class idx,        snd detection ... ],\n",
    "  confidences:    [ fst detextion: confidence,       snd detection ... ]\n",
    ")\n",
    "```\n",
    "\n",
    "Now we cant take the result bact to the numpy world and put results on an image.\n",
    "We have single input image, then we will get rid of the batch also:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "bboxes, classes, confidences = [x.detach().cpu().numpy() for x in encoded[0]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can filter results with confidence lower than some treshold:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = np.argwhere(confidences > 0.3).squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To show labels on the detections we need to decode labels. Our model is trained on COCO 2017 dataset. Then we will use labels from COCO:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_file = '/datasets/coco2017/annotations/instances_val2017.json'\n",
    "with open(json_file,'r') as COCO:\n",
    "    js = json.loads(COCO.read())\n",
    "class_names = [ category['name'] for category in js['categories'] ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can start building the picture with results.\n",
    "Bounding boxes returned by the model are enclosed in a [0..1] range. We need to scale it to [0..300], as it is the size of the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 149828 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,ax = plt.subplots(1)\n",
    "ax.imshow(out)\n",
    "for idx in best:\n",
    "    left, top, right, bottom = bboxes[idx]\n",
    "    x, y, w, h = [val*300 for val in [left, top, right-left, bottom-top]]\n",
    "    rect = patches.Rectangle((x, y),w,h,linewidth=1,edgecolor='r',facecolor='none')\n",
    "    ax.add_patch(rect)\n",
    "    ax.text(x, y, class_names[classes[idx]-1], bbox=dict(facecolor='white', alpha=0.5))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
