{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Compile YOLO-V2 and YOLO-V3 in DarkNet Models\n**Author**: [Siju Samuel](https://siju-samuel.github.io/)\n\nThis article is an introductory tutorial to deploy darknet models with TVM.\nAll the required models and libraries will be downloaded from the internet by the script.\nThis script runs the YOLO-V2 and YOLO-V3 Model with the bounding boxes\nDarknet parsing have dependancy with CFFI and CV2 library\nPlease install CFFI and CV2 before executing this script\n\n```bash\npip install cffi\npip install opencv-python\n```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# numpy and matplotlib\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport sys\n\n# tvm, relay\nimport tvm\nfrom tvm import te\nfrom tvm import relay\nfrom ctypes import *\nfrom tvm.contrib.download import download_testdata\nfrom tvm.relay.testing.darknet import __darknetffi__\nimport tvm.relay.testing.yolo_detection\nimport tvm.relay.testing.darknet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Choose the model\nModels are: 'yolov2', 'yolov3' or 'yolov3-tiny'\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Model name\nMODEL_NAME = \"yolov3\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Download required files\nDownload cfg and weights file if first time.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "CFG_NAME = MODEL_NAME + \".cfg\"\nWEIGHTS_NAME = MODEL_NAME + \".weights\"\nREPO_URL = \"https://github.com/dmlc/web-data/blob/main/darknet/\"\nCFG_URL = REPO_URL + \"cfg/\" + CFG_NAME + \"?raw=true\"\nWEIGHTS_URL = \"https://pjreddie.com/media/files/\" + WEIGHTS_NAME\n\ncfg_path = download_testdata(CFG_URL, CFG_NAME, module=\"darknet\")\nweights_path = download_testdata(WEIGHTS_URL, WEIGHTS_NAME, module=\"darknet\")\n\n# Download and Load darknet library\nif sys.platform in [\"linux\", \"linux2\"]:\n    DARKNET_LIB = \"libdarknet2.0.so\"\n    DARKNET_URL = REPO_URL + \"lib/\" + DARKNET_LIB + \"?raw=true\"\nelif sys.platform == \"darwin\":\n    DARKNET_LIB = \"libdarknet_mac2.0.so\"\n    DARKNET_URL = REPO_URL + \"lib_osx/\" + DARKNET_LIB + \"?raw=true\"\nelse:\n    err = \"Darknet lib is not supported on {} platform\".format(sys.platform)\n    raise NotImplementedError(err)\n\nlib_path = download_testdata(DARKNET_URL, DARKNET_LIB, module=\"darknet\")\n\nDARKNET_LIB = __darknetffi__.dlopen(lib_path)\nnet = DARKNET_LIB.load_network(cfg_path.encode(\"utf-8\"), weights_path.encode(\"utf-8\"), 0)\ndtype = \"float32\"\nbatch_size = 1\n\ndata = np.empty([batch_size, net.c, net.h, net.w], dtype)\nshape_dict = {\"data\": data.shape}\nprint(\"Converting darknet to relay functions...\")\nmod, params = relay.frontend.from_darknet(net, dtype=dtype, shape=data.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import the graph to Relay\ncompile the model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "target = tvm.target.Target(\"llvm\", host=\"llvm\")\ndev = tvm.cpu(0)\ndata = np.empty([batch_size, net.c, net.h, net.w], dtype)\nshape = {\"data\": data.shape}\nprint(\"Compiling the model...\")\nwith tvm.transform.PassContext(opt_level=3):\n    lib = relay.build(mod, target=target, params=params)\n\n[neth, netw] = shape[\"data\"][2:]  # Current image shape is 608x608"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load a test image\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "test_image = \"dog.jpg\"\nprint(\"Loading the test image...\")\nimg_url = REPO_URL + \"data/\" + test_image + \"?raw=true\"\nimg_path = download_testdata(img_url, test_image, \"data\")\n\ndata = tvm.relay.testing.darknet.load_image(img_path, netw, neth)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Execute on TVM Runtime\nThe process is no different from other examples.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import graph_executor\n\nm = graph_executor.GraphModule(lib[\"default\"](dev))\n\n# set inputs\nm.set_input(\"data\", tvm.nd.array(data.astype(dtype)))\n# execute\nprint(\"Running the test image...\")\n\n# detection\n# thresholds\nthresh = 0.5\nnms_thresh = 0.45\n\nm.run()\n# get outputs\ntvm_out = []\nif MODEL_NAME == \"yolov2\":\n    layer_out = {}\n    layer_out[\"type\"] = \"Region\"\n    # Get the region layer attributes (n, out_c, out_h, out_w, classes, coords, background)\n    layer_attr = m.get_output(2).numpy()\n    layer_out[\"biases\"] = m.get_output(1).numpy()\n    out_shape = (layer_attr[0], layer_attr[1] // layer_attr[0], layer_attr[2], layer_attr[3])\n    layer_out[\"output\"] = m.get_output(0).numpy().reshape(out_shape)\n    layer_out[\"classes\"] = layer_attr[4]\n    layer_out[\"coords\"] = layer_attr[5]\n    layer_out[\"background\"] = layer_attr[6]\n    tvm_out.append(layer_out)\n\nelif MODEL_NAME == \"yolov3\":\n    for i in range(3):\n        layer_out = {}\n        layer_out[\"type\"] = \"Yolo\"\n        # Get the yolo layer attributes (n, out_c, out_h, out_w, classes, total)\n        layer_attr = m.get_output(i * 4 + 3).numpy()\n        layer_out[\"biases\"] = m.get_output(i * 4 + 2).numpy()\n        layer_out[\"mask\"] = m.get_output(i * 4 + 1).numpy()\n        out_shape = (layer_attr[0], layer_attr[1] // layer_attr[0], layer_attr[2], layer_attr[3])\n        layer_out[\"output\"] = m.get_output(i * 4).numpy().reshape(out_shape)\n        layer_out[\"classes\"] = layer_attr[4]\n        tvm_out.append(layer_out)\n\nelif MODEL_NAME == \"yolov3-tiny\":\n    for i in range(2):\n        layer_out = {}\n        layer_out[\"type\"] = \"Yolo\"\n        # Get the yolo layer attributes (n, out_c, out_h, out_w, classes, total)\n        layer_attr = m.get_output(i * 4 + 3).numpy()\n        layer_out[\"biases\"] = m.get_output(i * 4 + 2).numpy()\n        layer_out[\"mask\"] = m.get_output(i * 4 + 1).numpy()\n        out_shape = (layer_attr[0], layer_attr[1] // layer_attr[0], layer_attr[2], layer_attr[3])\n        layer_out[\"output\"] = m.get_output(i * 4).numpy().reshape(out_shape)\n        layer_out[\"classes\"] = layer_attr[4]\n        tvm_out.append(layer_out)\n        thresh = 0.560\n\n# do the detection and bring up the bounding boxes\nimg = tvm.relay.testing.darknet.load_image_color(img_path)\n_, im_h, im_w = img.shape\ndets = tvm.relay.testing.yolo_detection.fill_network_boxes(\n    (netw, neth), (im_w, im_h), thresh, 1, tvm_out\n)\nlast_layer = net.layers[net.n - 1]\ntvm.relay.testing.yolo_detection.do_nms_sort(dets, last_layer.classes, nms_thresh)\n\ncoco_name = \"coco.names\"\ncoco_url = REPO_URL + \"data/\" + coco_name + \"?raw=true\"\nfont_name = \"arial.ttf\"\nfont_url = REPO_URL + \"data/\" + font_name + \"?raw=true\"\ncoco_path = download_testdata(coco_url, coco_name, module=\"data\")\nfont_path = download_testdata(font_url, font_name, module=\"data\")\n\nwith open(coco_path) as f:\n    content = f.readlines()\n\nnames = [x.strip() for x in content]\n\ntvm.relay.testing.yolo_detection.show_detections(img, dets, thresh, names, last_layer.classes)\ntvm.relay.testing.yolo_detection.draw_detections(\n    font_path, img, dets, thresh, names, last_layer.classes\n)\nplt.imshow(img.transpose(1, 2, 0))\nplt.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",
      "version": "3.7.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}