{
 "cells": [
	{
		"cell_type": "code",
		"metadata": {},
		"source": [
			"#  The MIT License (MIT)",
			"#",
			"#  Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.",
			"#",
			"#  Permission is hereby granted, free of charge, to any person obtaining a copy",
			"#  of this software and associated documentation files (the 'Software'), to deal",
			"#  in the Software without restriction, including without limitation the rights",
			"#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell",
			"#  copies of the Software, and to permit persons to whom the Software is",
			"#  furnished to do so, subject to the following conditions:",
			"#",
			"#  The above copyright notice and this permission notice shall be included in",
			"#  all copies or substantial portions of the Software.",
			"#",
			"#  THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR",
			"#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,",
			"#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE",
			"#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER",
			"#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,",
			"#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN",
			"#  THE SOFTWARE."
		]
	},

  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade pip\n",
    "!pip install opencv-python==4.1.2.30\n",
    "!pip install matplotlib\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt \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. \n",
    "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.\n",
    "\n",
    "If you receive \"cannot open shared object file: No such file or directory\" , please make sure `/opt/rocm/lib` is included in $LD_LIBRARY_PATH\n",
    "```\n",
    " cannot open shared object file: No such file or directory\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not path.exists(\"./resnet50.onnx\"):\n",
    "    !wget https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx",
    "    !mv resnet50-v2-7.onnx resnet50.onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the simplified imagenet labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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. If you are working on headless server, please execute the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": "markdown",
   "metadata": {},
   "source": [
    "If script is run on a headless server where .imshow() experiences problems, the following cell for histogram can be run to verify functionalty:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_labels = []\n",
    "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",
    "    output_labels.append(labels[top_prediction])\n",
    "\n",
    "cap.release()\n",
    "output_labels = np.array(output_labels)\n",
    "plt.hist(output_labels) \n",
    "plt.xticks(rotation = 90)\n",
    "plt.show()"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}