colab_yolov8s_create_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/marcory-hub/hailo/blob/main/colab_yolov8s_create_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. How to prepair you dataset\n",
        "\n",
        "1. In your Google Drive make a folder 'datasets'. Download the images and labels to that folder. If you do not have your own dataset, you can download the [hornet3000+ dataset](https://www.kaggle.com/datasets/marcoryvandijk/vespa-velutina-v-crabro-vespulina-vulgaris)\n",
        "\n",
        "2. Folder structure\n",
        "Only if you use your own dataset: make sure it has this folder structure with these exact names of the folders (class_1, class_2 can have your names of the classes):\n",
        "\n",
        "```\n",
        "data\n",
        "  train\n",
        "    images\n",
        "      class_1\n",
        "      class_2\n",
        "      etc...\n",
        "    labels\n",
        "      class_1\n",
        "      class_2\n",
        "      ect...\n",
        "   val\n",
        "    images\n",
        "      class_1\n",
        "      class_2\n",
        "      ...\n",
        "    labels\n",
        "      class_1\n",
        "      class_2\n",
        "      ...\n",
        "```\n",
        "\n",
        "3. Only if you use own dataset: zip the data folder to a file names `data.zip`.\n",
        "\n",
        "4. Copy the data.zip file (with the Kaggle dataset) to the folder datasets.\n",
        "\n",
        "5. Add dataset.yaml to folder datasets. This contains:\n",
        "- absolute path to train images (train)\n",
        "- absolute path validation (val)\n",
        "- the number of classes (nc)\n",
        "- and the class names (names)\n",
        "For example:\n",
        "```\n",
        "train: /content/data/train # train images\n",
        "val: /content/data/val # val images\n",
        "nc: 3\n",
        "names: ['Vespa_velutina', 'Vespa_crabro', 'Vespula_vulgaris']\n",
        "```\n",
        "If you use your own dataset make sure to adjust the nc and names accordingly."
      ],
      "metadata": {
        "id": "fOGSWAe4DjBI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. Mount your Google Drive"
      ],
      "metadata": {
        "id": "AHoqKf_x9L86"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tqo8tLtHAjiD"
      },
      "outputs": [],
      "source": [
        "# Mount Google Drive\n",
        "\n",
        "from google.colab import drive\n",
        "\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Check GPU (optional)"
      ],
      "metadata": {
        "id": "tHi8AHML9VhZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Check GPU (make sure runtimetype GPU (T4 or A100) is selected)\n",
        "# If 'command not found' then change the runtimetype (click 'Runtime' and select change runtimetype in the drop down menu)\n",
        "\n",
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "v10EkpTWAzcv",
        "collapsed": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3. Copy data.zip to Colab file system and unzip"
      ],
      "metadata": {
        "id": "_EpQpqQJ9b3i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Prepare data\n",
        "\n",
        "!scp '/content/gdrive/My Drive/datasets/data.zip' '/content/data.zip'\n",
        "\n",
        "!unzip '/content/data.zip'"
      ],
      "metadata": {
        "id": "dzksqg4LBEpa",
        "collapsed": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Check paths (optional)"
      ],
      "metadata": {
        "id": "4xmWC4k--UUp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Check is path is correct\n",
        "\n",
        "#  output should be\n",
        "#     train val\n",
        "#     images labels\n",
        "#     dataset.yaml data.zip\n",
        "\n",
        "# If not, change your datastructure as described under 1.\n",
        "\n",
        "!ls '/content/data/'\n",
        "\n",
        "!ls '/content/data/train/'\n",
        "\n",
        "!ls '/content/data/val/'\n",
        "\n",
        "!ls '/content/gdrive/My Drive/datasets/'"
      ],
      "metadata": {
        "id": "-IPvW8o5Ps2N"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 4. Install libraries and train the model"
      ],
      "metadata": {
        "id": "nRU0VsSi-k0q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Pip install method (recommended)\n",
        "\n",
        "!pip install ultralytics==8.0.196\n",
        "\n",
        "# Fix for \"TypeError: unhashable type: 'numpy.ndarray'\"\n",
        "\n",
        "!pip install albumentations==1.4"
      ],
      "metadata": {
        "id": "JuIVwb6gBqiq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Train the model\n",
        "\n",
        "# Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.\n",
        "# wandb: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)\n",
        "# wandb: You can find your API key in your browser here: https://wandb.ai/authorize\n",
        "\n",
        "import os\n",
        "\n",
        "from ultralytics import YOLO\n",
        "\n",
        "# Load a model\n",
        "model = YOLO(\"yolov8s.pt\")  # load a pretained model (choose yolov8n for the smallest model)\n",
        "\n",
        "# Use the model\n",
        "# Adjust epochs, batch, imgsz or other arguments if needed to optimize the model\n",
        "results = model.train(data='/content/gdrive/MyDrive/datasets/dataset.yaml', epochs=1, batch=16, imgsz=640)  # train the model"
      ],
      "metadata": {
        "id": "nAUMVQA1RuZ2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 5. Export the model in ONNX format"
      ],
      "metadata": {
        "id": "UIRlPuiKAKTd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Export the model in ONNX format\n",
        "\n",
        "model.export(format=\"onnx\", opset=9)\n"
      ],
      "metadata": {
        "id": "-Ra_K83jDJ-s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Export the model in ONNX format IR_version 9 (adjust IR version if needed)\n",
        "\n",
        "import onnx\n",
        "\n",
        "# Load the ONNX model\n",
        "model = onnx.load(\"/content/runs/detect/train/weights/best.onnx\")\n",
        "\n",
        "# Change the IR version to 9\n",
        "model.ir_version = 9\n",
        "\n",
        "# Save the modified model\n",
        "onnx.save(model, \"/content/runs/detect/train/weights/best9.onnx\")\n",
        "\n",
        "print(\"Model IR version changed to 9 and saved as best9.onnx\")\n"
      ],
      "metadata": {
        "id": "XP4pqUH-kp2w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 6. Zip and Download the results to your local computer\n",
        "\n",
        "Data is lost when closing the page!!!\n",
        "\n",
        "## If an error occurs:\n",
        "1. Click `Runtime` in the toolbar on the top\n",
        "2. Select `Restart session` / `Sessie opnieuw starten` from the drop down menu\n",
        "3. Run this codeblock again"
      ],
      "metadata": {
        "id": "E63Hr9A8CVft"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# DON'T FORGET TO DOWNLOAD THE RESULTS!!!\n",
        "\n",
        "# Data is lost when closing the page!!!\n",
        "\n",
        "# IN MOST CASES AN ERROR OCCURS:\n",
        "# Click `Runtime`, select `Restart session` / `Sessie opnieuw opstarten`\n",
        "# then run this codeblock again\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "try:\n",
        "  !zip -r /content/runs.zip /content/runs\n",
        "  files.download('/content/runs.zip')\n",
        "except Exception as e:\n",
        "  print(f\"An error occurred: {e}\")\n",
        "  print(\"Click 'Runtime' -> 'Restart session' and try running the code again.\")"
      ],
      "metadata": {
        "id": "GV9nv0PmCxgF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 7. [Compile the model in the Google Cloud Platform VM](https://github.com/marcory-hub/hailo/blob/main/compile-the-model-using-hailo-model-zoo.md).\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "DkdSG8W6IHsK"
      }
    }
  ]
}