Lab 7 - YOLO (You Only Look Once).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Lab 7 - YOLO (You Only Look Once).ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sapjunior/chulacv2018/blob/master/Lab%207%20-%20YOLO%20(You%20Only%20Look%20Once).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "Qm3MMjj_9LgA",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# 2110443 - Computer Vision (2018/1)\n",
        "## Lab 7 - YOLO (You Only Look Once)\n",
        "In this lab, we will learn how to use a famouse regression based object detection algorithm name YOLO (You Only Look Once). In this class main repository, we already prepare a PyTorch version of YOLO  based on an [original Implementation](https://github.com/pjreddie/darknet), \n"
      ]
    },
    {
      "metadata": {
        "id": "9ZTnHne79GJ6",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Install based package\n",
        "!pip3 install http://download.pytorch.org/whl/cu92/torch-0.4.1-cp36-cp36m-linux_x86_64.whl\n",
        "!pip3 install torchvision\n",
        "\n",
        "# Clone YOLO based code\n",
        "!git clone https://github.com/sapjunior/objdetection.git\n",
        "  \n",
        "# Download trained weight on COCO dataset\n",
        "!wget https://cgci.cp.eng.chula.ac.th/yolo-darknet53-coco.pth -O yolo-darknet53-coco.pth"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TuBr83jHaYlI",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "The above command will clone [this repository](https://github.com/sapjunior/objdetection) in to objdetection directory. The internal folder is organized as follows:\n",
        "  - dataset   \n",
        "  => ListDataset : Read Object Detection Dataset from file. The input text file should be in following format (one line per image file)\n",
        "  \n",
        ">>`a.jpg xmin ymin xmax ymax label xmin ymin xmax ymax label`\n",
        "  \n",
        "  - utils ==> Utility functions, such as, calculate IOU between two boxes and NMS\n",
        "  - model  \n",
        "    ==> backbone : This directory contains a backbone network (feature extractor) for YOLO. In this example we use darknet-53 as a default network.\n",
        "    ==> yolo.py : A main network for YOLO Object Detection framework\n",
        "    ![alt text](https://cdn-images-1.medium.com/max/1750/1*d4Eg17IVJ0L41e7CTWLLSg.png/)"
      ]
    },
    {
      "metadata": {
        "id": "iAzpKRUsZ0Da",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Import OpenCV, PyTorch, Numpy and Matplotlib"
      ]
    },
    {
      "metadata": {
        "id": "4A3YziUK_fW0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import torch\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "from objdetection.models.yolo import YOLONet\n",
        "from objdetection.utils.yoloBoxCoder import YOLOBoxCoder, YOLOBoxPostProcess"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Brv8DuW1Z9fu",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Network Parameters ###"
      ]
    },
    {
      "metadata": {
        "id": "Bu7B0VDMDAYn",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "cocoNumClass = 80\n",
        "trainImageSize = (412,412)\n",
        "anchors = [ [[116, 90], [156, 198], [373, 326]],\n",
        "            [[30, 61], [62, 45], [59, 119]],\n",
        "            [[10, 13], [16, 30], [33, 23]]]\n",
        "with open('objdetection/coco.names') as f:\n",
        "    fileContent = f.readlines()\n",
        "cocoClassName = [x.strip() for x in fileContent]\n",
        "print(cocoClassName)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "VPKL472waCMq",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Decalare Network and Load Trained Weight"
      ]
    },
    {
      "metadata": {
        "id": "1YcDp2Z8Cckm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "net = YOLONet(numClass = cocoNumClass)\n",
        "checkpoint = torch.load('yolo-darknet53-coco.pth', map_location=lambda storage, location: storage)\n",
        "net.load_state_dict(checkpoint)\n",
        "if torch.cuda.is_available():\n",
        "  torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "  net.cuda()\n",
        "net.eval()\n",
        "boxCoder = YOLOBoxCoder(anchors=anchors, numClass= cocoNumClass, inputFeatMapSizes=net.outputFeatMapSizes)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "DG9VqC2zaJGm",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Input Image Preprocess and Warpped Classify Function"
      ]
    },
    {
      "metadata": {
        "id": "jy7RxM7MBRO5",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def preprocess(inputImage, trainImageSize):\n",
        "  preprocessImage = cv2.cvtColor(inputImage.copy(), cv2.COLOR_BGR2RGB)\n",
        "  preprocessImage = cv2.resize(preprocessImage, trainImageSize)\n",
        "  preprocessImage = (preprocessImage.astype(np.float32) / 255).transpose(2,0,1)\n",
        "  return preprocessImage"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_fOzFLtnDsZG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def classify(net, trainImageSize, inputImage):\n",
        "  preprocessImage = preprocess(inputImage, trainImageSize)\n",
        "  batchImage = torch.from_numpy(np.array([preprocessImage]))\n",
        "  if torch.cuda.is_available():\n",
        "    batchImage = batchImage.cuda()\n",
        "    \n",
        "  with torch.no_grad():\n",
        "    netOutputs = net(batchImage)\n",
        "    \n",
        "  # Output is in 412x412 image size => We need to multiple by inputImage resized ratio\n",
        "  outputBoxes = YOLOBoxPostProcess(boxCoder.decode(netOutputs), numClass = 80)\n",
        "  \n",
        "  outputImage = cv2.cvtColor(inputImage.copy(), cv2.COLOR_BGR2RGB)\n",
        "  resizeRatioH,resizeRatioW = np.array(inputImage.shape[0:2])/np.array(trainImageSize)\n",
        "  \n",
        "  outputBoxesRescale = []\n",
        "  \n",
        "  \n",
        "  for x1,y1,x2,y2,conf,cls_pred in outputBoxes[0]:\n",
        "    \n",
        "    # Rescale bbox by resized ratio\n",
        "    x1, y1, x2, y2 = int(x1*resizeRatioW), int(y1*resizeRatioH), int(x2*resizeRatioW), int(y2*resizeRatioH)\n",
        "    \n",
        "    # Append result to rescaled output\n",
        "    outputBoxesRescale.append([x1,y1,x2,y2,conf,cls_pred])\n",
        "    \n",
        "    # Draw Result\n",
        "    cv2.putText(outputImage, cocoClassName[int(cls_pred)]+' : '+str(round(conf,3))\n",
        "                ,(x1,y1+20), cv2.FONT_HERSHEY_SIMPLEX, 0.75,(255,0,0),2,cv2.LINE_AA)\n",
        "    cv2.rectangle(outputImage,(x1,y1),(x2,y2),(0,255,0),3)\n",
        "    \n",
        "  return outputImage, np.array(outputBoxesRescale)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "CSLpyCTSaP3p",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Read input image and classify"
      ]
    },
    {
      "metadata": {
        "id": "3g_WcXj-EI5U",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Put your image here\n",
        "inputImage = cv2.imread('objdetection/yolo-test.png')\n",
        "inputImageRGB = cv2.cvtColor(inputImage, cv2.COLOR_BGR2RGB)\n",
        "plt.imshow(inputImageRGB)\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "uKH3m2p9EyJ0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "outputImage, outputBoxes = classify(net, trainImageSize, inputImage)\n",
        "plt.imshow(outputImage)\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}