AstroNet_YOLOv4.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "colab": {
      "name": "AstroNet-YOLOv4.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Training a custom YOLOv4 implementation on the SKA data"
      ],
      "metadata": {
        "id": "M0Xia5f1QwEO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, clone the repository"
      ],
      "metadata": {
        "id": "yZqGpmLZQqGx"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!git clone https://github.com/vairodp/AstroNet.git"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bqqzvQzcQgJY",
        "outputId": "ccb8f6ab-1a69-4b9a-ed31-060f851a83d7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "%cd AstroNet"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6R1gfhqvILXY",
        "outputId": "e8e26d2e-3eb1-4e43-8c35-89042974a971"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then install the missing libraries that the code requires:"
      ],
      "metadata": {
        "id": "I3Eg9axERUTJ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!pip install tensorflow_addons\r\n",
        "!pip install tensorflow-datasets==4.3.0\r\n",
        "!pip install imgaug==0.4.0"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "prrtr9i8IQZX",
        "outputId": "648e2d8d-2173-4c46-dc01-2509a632cb19"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make sure you're using a GPU in order to get fast train and inference."
      ],
      "metadata": {
        "id": "lgaLTHmXRcZl"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "import tensorflow as tf\r\n",
        "device_name = tf.test.gpu_device_name()\r\n",
        "if device_name != '/device:GPU:0':\r\n",
        "  raise SystemError('GPU device not found')\r\n",
        "print('Found GPU at: {}'.format(device_name))"
      ],
      "outputs": [],
      "metadata": {
        "id": "dR0-0zDuIUCz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "75299616-2e44-4ffe-e3e4-3480e4d0bd91"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "device_name = tf.test.gpu_device_name()\r\n",
        "if device_name != '/device:GPU:0':\r\n",
        "  print(\r\n",
        "      '\\n\\nThis error most likely means that this notebook is not '\r\n",
        "      'configured to use a GPU.  Change this in Notebook Settings via the '\r\n",
        "      'command palette (cmd/ctrl-shift-P) or the Edit menu.\\n\\n')\r\n",
        "  raise SystemError('GPU device not found')"
      ],
      "outputs": [],
      "metadata": {
        "id": "DFEB9YbgIY1R"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Finally, download the pretrained YOLOv4 weights in `checkpoints` and run the code\n",
        "\n",
        "**Note:** in case you want to modify some parameters of the network, you need to change the `configs/train_config.py` file"
      ],
      "metadata": {
        "id": "Q-NoHJEJRmIR"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights -P checkpoints/"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rF8UAuTMUWOL",
        "outputId": "f4f02ae7-c6a8-4a51-a53d-5390e316b0a0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "%cd src"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JGU_M-5uIdc7",
        "outputId": "b90e09b8-1e96-4852-b7fb-ece3be684cc6"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!python train.py"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8mjnKso0IjDi",
        "outputId": "d52c771f-442a-48ff-8203-cba3a4c43646"
      }
    }
  ]
}