{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "HyperLabel Yolo v3 training.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RCRQLOGEoJsM",
        "colab_type": "text"
      },
      "source": [
        "# Prepration\n",
        "\n",
        "- In order to train the Yolo v3 model, we will need to compile `darknet` which is an open source neural network framework written in C and CUDA. \n",
        "- Once the `darknet` is compiled, we will download the exported data set from HyperLabel\n",
        "\n",
        "\n",
        "**Note**: Before starting to work, change the runtime to use the GPU.\n",
        "\n",
        "If you don't have any dataset or pretrained model (weight) for that dataset, you can use our demo data set and trained weight.\n",
        "\n",
        "Copy the following files to your Google drive (create a new folder called `ml` in your Google drive and put the files in there)\n",
        "\n",
        "- [Dataset](https://drive.google.com/file/d/1wxHNFR2xXx6wJMJR2ey2NP_rWFJNs8-I/view?usp=sharing)\n",
        "- [Pre trained weight](https://drive.google.com/file/d/1-5iEJhBgwRGBZ_ZBfO6Juo2tF6lGa4b6/view?usp=sharing)\n",
        "\n",
        "\n",
        "### Clone darknet"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qdQWft7loUTA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/AlexeyAB/darknet"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ThtCp4-EpQPT",
        "colab_type": "text"
      },
      "source": [
        "### Compile darknet"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PGiduF4bpUZ-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd darknet\n",
        "!sed -i 's/GPU=0/GPU=1/g' Makefile\n",
        "!cat Makefile\n",
        "!make"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "guZasLRGpgRD",
        "colab_type": "text"
      },
      "source": [
        "### Download the data set exported from HyperLabel\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uQ_Vl3RzpmPU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os.path\n",
        "import shutil\n",
        "from google.colab import drive\n",
        "\n",
        "if not os.path.exists('/content/gdrive'):\n",
        "  drive.mount('/content/gdrive')\n",
        "  \n",
        "DOWNLOAD_LOCATION = '/content/darknet/data/'\n",
        "DRIVE_DATASET_FILE = '/content/gdrive/My Drive/ml/yolo-dataset.zip' #adjust path/name of dataset which is in your G-drive\n",
        "\n",
        "shutil.copy(DRIVE_DATASET_FILE, DOWNLOAD_LOCATION)\n",
        "\n",
        "print('Successfully downloaded the dataset')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6mv4-cjHrjV0",
        "colab_type": "text"
      },
      "source": [
        "### Unzip the dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8MrNs0X2rixE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!unzip data/yolo-dataset.zip -d data/ # adjust the dataset filename which you have downloaded from Google drive"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CV_Y9nfbsATG",
        "colab_type": "text"
      },
      "source": [
        "##Training the model\n",
        "\n",
        "- Download the initial weight\n",
        "- Train the model\n",
        "\n",
        "### Download initial pre-trained weights for the convolutional layers (If you have already trained and saved the weights to your Google drive, you can skip this)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v9DIiCpAsNpK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget https://pjreddie.com/media/files/darknet53.conv.74"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BYioQTjxtuMr",
        "colab_type": "text"
      },
      "source": [
        "### Set Yolo configurations\n",
        "\n",
        "We need to set configurations for Yolo in order to properly train. There are few settings which we need to change in the default yolov3.cfg file.\n",
        "\n",
        "- batch\n",
        "- subdivisions (if you get memory out error, increase this 16, 32 or 64)\n",
        "- max_batches (it should be classes*2000)\n",
        "- steps (it should be 80%, 90% of max_batches)\n",
        "- classes (the number of classes which you are going to train)\n",
        "- filters (the value for filters can be calculated using (classes + 5)x3 )\n",
        "\n",
        "Change the values below as per your requirement"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ymbna8dntyTI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!sed -i 's/batch=1/batch=64/g' cfg/yolov3.cfg\n",
        "!sed -i 's/subdivisions=1/subdivisions=32/g' cfg/yolov3.cfg\n",
        "!sed -i 's/max_batches = 500200/max_batches = 10000/g' cfg/yolov3.cfg\n",
        "!sed -i 's/steps=400000,450000/steps=8000,9000/g' cfg/yolov3.cfg\n",
        "!sed -i 's/classes=80/classes=5/g' cfg/yolov3.cfg\n",
        "!sed -i 's/filters=255/filters=30/g' cfg/yolov3.cfg\n",
        "!cat cfg/yolov3.cfg"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "br8phhE-N1MW",
        "colab_type": "text"
      },
      "source": [
        "**Now we are moving forward with the training part. If you have already trained the model, you can fetch them from Google drive and skip the training part.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Dgv3LnvOAB3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Only run this cell, if you have already trained the model and have weights and backup files in your Google drive\n",
        "# (Optional) Download the pretrained weight from Google drive\n",
        "\n",
        "import os.path\n",
        "import shutil\n",
        "from google.colab import drive\n",
        "\n",
        "if not os.path.exists('/content/gdrive'):\n",
        "  drive.mount('/content/gdrive')\n",
        "  \n",
        "BACKUP_FOLDER = '/content/darknet/backup'\n",
        "DRIVE_YOLO_BACKUP = '/content/gdrive/My Drive/ml/yolov3_last.weights'\n",
        "\n",
        "shutil.copy(DRIVE_YOLO_BACKUP, BACKUP_FOLDER)\n",
        "\n",
        "print('Successfully fetched the pretrained files for Yolo from Google drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mRRffYhouYJp",
        "colab_type": "text"
      },
      "source": [
        "### Start the training\n",
        "\n",
        "It will take a long time to complete, so sit back and relax\n",
        "\n",
        "Note: If during training you see nan values for avg (loss) field - then training goes wrong, but if nan is in some other lines - then training goes well.\n",
        "\n",
        "- file yolo-obj_last.weights will be saved to the darknet/backup for each 100 iterations\n",
        "- file yolo-obj_xxxx.weights will be saved to the darknet/backup for each 1000 iterations\n",
        "- After each 100 iterations if you want, you can stop and later start training from this point. For example, after 2000 iterations you can stop training, and later just start training using: darknet detector train data/obj.data yolov3.cfg backup/yolo-obj_2000.weights"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vbfDfCnoua1i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# use the line below to train a fresh model\n",
        "!./darknet detector train data/obj.data cfg/yolov3.cfg darknet53.conv.74\n",
        "\n",
        "# use the line below to retrain your previous saved weight\n",
        "#!./darknet detector train data/obj.data cfg/yolov3.cfg backup/yolov3_last.weights"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TTliLPwYOk0d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Once you have trained your model, you can save them to your Google drive. So that next time, you don't need to retrain\n",
        "# This step is optional, you can skip it if you want\n",
        "import os.path\n",
        "import shutil\n",
        "from google.colab import drive\n",
        "\n",
        "if not os.path.exists('/content/gdrive'):\n",
        "  drive.mount('/content/gdrive')\n",
        "  \n",
        "YOLO_BACKUP = '/content/darknet/backup/yolov3_last.weights' #adjust the backup file name or keep it default\n",
        "DRIVE_DIR = '/content/gdrive/My Drive/ml/' #adjust path in your Google drive, or keep it default\n",
        "\n",
        "shutil.copy(YOLO_BACKUP, DRIVE_DIR)\n",
        "\n",
        "print('Saved training data to drive at: ' + DRIVE_DIR)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67xSMGMx2Qdo",
        "colab_type": "text"
      },
      "source": [
        "## Running predictions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o7RHTzYj2T98",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!./darknet detector test data/obj.data cfg/yolov3.cfg backup/yolov3_last.weights data/img/output-000000598.jpg"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_4jl4WkaanCY",
        "colab_type": "text"
      },
      "source": [
        "### Display result"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TeO5iiKNGUH-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def display_image(file_path = '/content/darknet/predictions.jpg'):\n",
        "    import cv2\n",
        "    import matplotlib.pyplot as plt\n",
        "    import os.path\n",
        "\n",
        "    if os.path.exists(file_path):\n",
        "      img = cv2.imread(file_path)\n",
        "      show_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) \n",
        "      plt.imshow(show_img)\n",
        "    else:\n",
        "      print('failed to open file')\n",
        "    \n",
        "display_image()"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}