{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resnet50 Inference\n",
    "\n",
    "## Description\n",
    "This example performs inference on a short wildlife video using a Resnet50 V2 model that has been pre-trained on imagenet data. The labels used for each class are simplified for readability, but still reflect the correct index-label pairs in official use. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import json\n",
    "import time\n",
    "import os.path\n",
    "from os import path \n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing MIGraphX Library\n",
    "Sometimes the PYTHONPATH variable is not set during installation of MIGraphX. If your receive a \"Module Not Found\" error when trying to `import migraphx` in your own application, try running:\n",
    "```\n",
    "$ export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH\n",
    "```\n",
    "For this example, the library will be added to the kernel's sys.path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "migx_lib_path = \"/opt/rocm/lib\"\n",
    "if migx_lib_path not in sys.path:\n",
    "    sys.path.append(migx_lib_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import migraphx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If this is your first time running this example, you will need to dowload the model and sample video.\n",
    "\n",
    "The following cell will ask you for your sudo password and then install/update the package `youtube-dl` if necessary. It will then use that tool to download the sample video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "········\n",
      "[youtube] TkqYmvH_XVs: Downloading webpage\n",
      "[youtube] TkqYmvH_XVs: Downloading MPD manifest\n",
      "[dashsegments] Total fragments: 34\n",
      "[download] Destination: sample_vid-TkqYmvH_XVs.f137.mp4\n",
      "\u001b[K[download] 100% of 70.35MiB in 00:06.31MiB/s ETA 00:000:11\n",
      "[dashsegments] Total fragments: 18\n",
      "[download] Destination: sample_vid-TkqYmvH_XVs.f140.m4a\n",
      "\u001b[K[download] 100% of 2.58MiB in 00:01.99MiB/s ETA 00:000102\n",
      "[ffmpeg] Merging formats into \"sample_vid-TkqYmvH_XVs.mp4\"\n",
      "Deleting original file sample_vid-TkqYmvH_XVs.f137.mp4 (pass -k to keep)\n",
      "Deleting original file sample_vid-TkqYmvH_XVs.f140.m4a (pass -k to keep)\n"
     ]
    }
   ],
   "source": [
    "if not path.exists(\"./sample_vid.mp4\"):\n",
    "    import getpass\n",
    "    import os\n",
    "    password = getpass.getpass()\n",
    "    command = \"sudo -H -S pip install --upgrade youtube-dl\"\n",
    "    os.system('echo %s | %s' % (password, command))\n",
    "    !youtube-dl https://youtu.be/TkqYmvH_XVs \n",
    "    !mv sample_vid-TkqYmvH_XVs.mp4 sample_vid.mp4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following will download the resnet50 v2 model from ONNX's model zoo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2021-01-13 20:24:16--  https://github.com/onnx/models/blob/master/vision/classification/resnet/model/resnet50-v2-7.onnx?raw=true\n",
      "Resolving github.com (github.com)... 140.82.112.3\n",
      "Connecting to github.com (github.com)|140.82.112.3|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://github.com/onnx/models/raw/master/vision/classification/resnet/model/resnet50-v2-7.onnx [following]\n",
      "--2021-01-13 20:24:16--  https://github.com/onnx/models/raw/master/vision/classification/resnet/model/resnet50-v2-7.onnx\n",
      "Reusing existing connection to github.com:443.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://media.githubusercontent.com/media/onnx/models/master/vision/classification/resnet/model/resnet50-v2-7.onnx [following]\n",
      "--2021-01-13 20:24:16--  https://media.githubusercontent.com/media/onnx/models/master/vision/classification/resnet/model/resnet50-v2-7.onnx\n",
      "Resolving media.githubusercontent.com (media.githubusercontent.com)... 151.101.48.133\n",
      "Connecting to media.githubusercontent.com (media.githubusercontent.com)|151.101.48.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 102442450 (98M) [application/octet-stream]\n",
      "Saving to: ‘resnet50-v2-7.onnx?raw=true’\n",
      "\n",
      "resnet50-v2-7.onnx? 100%[===================>]  97.70M  88.2MB/s    in 1.1s    \n",
      "\n",
      "2021-01-13 20:24:19 (88.2 MB/s) - ‘resnet50-v2-7.onnx?raw=true’ saved [102442450/102442450]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "if not path.exists(\"./resnet50.onnx\"):\n",
    "    !wget https://github.com/onnx/models/blob/master/vision/classification/resnet/model/resnet50-v2-7.onnx?raw=true\n",
    "    !mv 'resnet50-v2-7.onnx?raw=true' resnet50.onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the simplified imagenet labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('imagenet_simple_labels.json') as json_data:\n",
    "    labels = json.load(json_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model and Video Capture Setup\n",
    "\n",
    "The ONNX graph that is loaded by `parse_onnx()` is a generalized representation that must be compiled for a specific target before it can be executed. For this example, using the target \"gpu\" is recommended. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = migraphx.parse_onnx(\"resnet50.onnx\")\n",
    "model.compile(migraphx.get_target(\"gpu\"))\n",
    "model.print()     # Printed in terminal \n",
    "cap = cv2.VideoCapture(\"sample_vid.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pre-Processing Video Frames\n",
    "Resnet50 requires some preprocessing of video frames before it can run inference. \n",
    "\n",
    "The model will expect an NCHW tensor with the shape {1, 3, 224, 224} and the values loaded into a range of [0, 1] and then normalized using mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225]. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to square up the dimensions of the original image by cropping the longer of the two to the size of the shorter dimension. This will help to avoid any stretching or compressing of the input image.\n",
    "Then the image can be scaled up or down to the desired resolution of 224x224."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_nxn(image, n):\n",
    "    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "    if height > width:\n",
    "        dif = height - width\n",
    "        bar = dif // 2 \n",
    "        square = image[(bar + (dif % 2)):(height - bar),:]\n",
    "        return cv2.resize(square, (n, n))\n",
    "    elif width > height:\n",
    "        dif = width - height\n",
    "        bar = dif // 2\n",
    "        square = image[:,(bar + (dif % 2)):(width - bar)]\n",
    "        return cv2.resize(square, (n, n))\n",
    "    else:\n",
    "        return cv2.resize(image, (n, n))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the image data has the correct dimensions, the values can be normalized as described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(img_data):\n",
    "    mean_vec = np.array([0.485, 0.456, 0.406])\n",
    "    stddev_vec = np.array([0.229, 0.224, 0.225])\n",
    "    norm_img_data = np.zeros(img_data.shape).astype('float32')\n",
    "    for i in range(img_data.shape[0]):  \n",
    "        norm_img_data[i,:,:] = (img_data[i,:,:]/255 - mean_vec[i]) / stddev_vec[i]\n",
    "    return norm_img_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Inference on Single Frame\n",
    "\n",
    "The above pre-processing functions can now be applied to individual video frames and the data can be passed to the model for evaluation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_class(frame) -> int:\n",
    "    # Crop and resize original image\n",
    "    cropped = make_nxn(frame, 224)\n",
    "    # Convert from HWC to CHW\n",
    "    chw = cropped.transpose(2,0,1)\n",
    "    # Apply normalization\n",
    "    pp = preprocess(chw)\n",
    "    # Add singleton dimension (CHW to NCHW)\n",
    "    data = np.expand_dims(pp.astype('float32'),0)\n",
    "    # Run the model\n",
    "    results = model.run({'data':data})\n",
    "    # Extract the index of the top prediction\n",
    "    res_npa = np.array(results[0])\n",
    "    return np.argmax(res_npa)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference Loop over Full Video\n",
    "\n",
    "Now everything is in place so that we can run inference on each frame of the input video. The video will be played and the predicted label will be displayed on top of each frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "while (cap.isOpened()):\n",
    "    start = time.perf_counter()\n",
    "    ret, frame = cap.read()\n",
    "    if not ret: break\n",
    "    \n",
    "    top_prediction = predict_class(frame)\n",
    "    \n",
    "    end = time.perf_counter()\n",
    "    fps = 1 / (end - start)\n",
    "    fps_str = f\"Frames per second: {fps:0.1f}\"\n",
    "    label_str = \"Top prediction: {}\".format(labels[top_prediction])\n",
    "\n",
    "    labeled = cv2.putText(frame, \n",
    "                          label_str, \n",
    "                          (50, 50), \n",
    "                          cv2.FONT_HERSHEY_SIMPLEX, \n",
    "                          2, \n",
    "                          (255, 255, 255), \n",
    "                          3, \n",
    "                          cv2.LINE_AA)\n",
    "    labeled = cv2.putText(labeled, \n",
    "                          fps_str, \n",
    "                          (50, 1060), \n",
    "                          cv2.FONT_HERSHEY_SIMPLEX, \n",
    "                          2, \n",
    "                          (255, 255, 255), \n",
    "                          3, \n",
    "                          cv2.LINE_AA)\n",
    "    cv2.imshow(\"Resnet50 Inference\", labeled)\n",
    "\n",
    "    if cv2.waitKey(1) & 0xFF == ord('q'): # 'q' to quit\n",
    "        break\n",
    "\n",
    "cap.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}