dataset_yolo_onnx.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyOXQRKjhDYSPBf3u0NxKpvr",
      "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-colab/blob/main/dataset_yolo_onnx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# From Dataset to YOLO11 ONNX file\n",
        "\n",
        "1. **Prepare Files:** Download a dataset in yolo format or create your own.\n",
        "2. **Zip and Upload:** Zip your dataset and upload it to your Google Drive.\n",
        "3. **Colab Access:** Unzip the dataset in your Colab notebook environment.\n",
        "4. **Train YOLO11:** Use the provided code to train a YOLO11 model.\n",
        "5. **Convert to ONNX:** Convert the trained model `best.pt` to ONNX format `best.onnx` for Hailo's Dataflow Compiler.\n",
        "6. **Save model:** Zip training folders with the `best.pt` and `best.onnx` in the folder `weights`.\n",
        "\n",
        "### Before you start\n",
        "\n",
        "1. **Check internet** for the most recent information (for example on github) before you install Ultralytics. A miner was injected in ultralytics 8.3.41, 8.3.42, 8.3.45, and 8.3.46.\n",
        "2. **Select runtime type**\n",
        "Select a runtime type with a GPU, this is needed to train the YOLO model (T4 (free) or A100)."
      ],
      "metadata": {
        "id": "bY6YaSQWON74"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1. Prepair Files\n",
        "\n",
        "To train YOLO11 model with your data you need use images and labels in the correct folder structure and a yaml file.\n",
        "1. **Dataset:** a collection of images and their corresponding labels. Labels tell the model what objects are in the image and where they are located.You can either download an existing dataset in YOLO format from online sources like the [hornet3000+ dataset](https://www.kaggle.com/datasets/marcoryvandijk/vespa-velutina-v-crabro-vespulina-vulgaris) or create your own dataset using tools like CVAT or Roboflow to annotate images with bounding boxes around the objects you want the model to detect.\n",
        "\n",
        "2. **data.yaml:** a configuration file that tells the training script where to find your dataset and what it contains.\n",
        "\n",
        "### Steps to prepare your files\n",
        "\n",
        "1. Organize your dataset:\n",
        "  - If you created your own dataset, make sure it follows this structure. use these exact names of the folders (only the jpg en txt files in these folders have custom names)\n",
        "    ```\n",
        "    dataset\\\n",
        "      train\\\n",
        "        images\\\n",
        "        labels\\\n",
        "      valid\n",
        "        images\\\n",
        "        labels\\\n",
        "      test (optional)\n",
        "        images\\\n",
        "        labels\\\n",
        "    ```\n",
        "\n",
        "2. Create data.yaml:\n",
        "  - The example provided shows how to fill it out based on your dataset location. Remember to adjust the number of classes (nc) and class names (names) to match your specific dataset in the order of the class numbers (so 0 is Vespa_velutina in this example)\n",
        "\n",
        "    ```\n",
        "train: /content/dataset/train # path to train images\n",
        "val: /content/dataset/valid # path to val images\n",
        "nc: 3\n",
        "names: ['Vespa_velutina', 'Vespa_crabro', 'Vespula_vulgaris']\n",
        "    ```\n",
        "Add this file to your dataset folder."
      ],
      "metadata": {
        "id": "rqDdKd2X_Mns"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. Zip and Upload\n",
        "\n",
        "\n",
        "1. Zip your dataset folder:\n",
        "  - This step speeds up uploading your dataset to Google Drive. On MacOS use `ditto -c -k --norsrc --keepParent images dataset.zip` to exclude finderfiles from the zipped file.\n",
        "2. Upload the zipped folder to your Google Drive.\n"
      ],
      "metadata": {
        "id": "cuSaGAA8DyDq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3. Colab Access\n",
        "\n",
        "To get easy access to the dataset files needed to (re-)train the YOLO11 model run the following codeblocks for these steps:\n",
        "\n",
        "1. Mount google drive.\n",
        "2. Unzip the dataset in Colab folder and rename this folder (this exact name is needed so the YOLO training codeblocks can access the data).\n",
        "3. Optional: check presence of data.yaml and folder structure."
      ],
      "metadata": {
        "id": "0tb3TjiJWgVm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. Mount google drive."
      ],
      "metadata": {
        "id": "wfhxV1USS3mz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Mount google drive\n",
        "from google.colab import drive\n",
        "\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "SUPnZCbI96Al"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. Unzip the dataset.\n",
        "\n",
        "  Adjust the names in the boxes on the right side if needed.\n",
        "\n",
        "- **dataset_path:** This should be the path to your zipped dataset file on your Google Drive. Typically, it's starting with \"/content/drive/MyDrive/\" followed by the specific path to your file. You can copy this by clicking the folder icon on the left sidebar, navigating to \"gdrive\" -> \"MyDrive\", and right-clicking on the desired folder to copy the path.\n"
      ],
      "metadata": {
        "id": "HAWUsZ-6KZar"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Unzip the dataset and rename the folder\n",
        "import os\n",
        "\n",
        "# Set Paths to dataset and dataset filename\n",
        "dataset_path = \"/content/gdrive/MyDrive/hailo/vespA17000.zip\"  # @param {type:\"string\"}\n",
        "\n",
        "# Unzip the dataset\n",
        "!unzip {dataset_path} -d '/content/'\n"
      ],
      "metadata": {
        "id": "05bLQk36i-94"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "3. Rename the folder with the dataset.\n",
        "\n",
        "- **dataset_filename:** This should be the exact filename of your zipped dataset, including the file extension (e.g., dataset.zip or dataset_cats.zip). This name will be used to create a new folder in your Colab environment to extract the dataset."
      ],
      "metadata": {
        "id": "E_zMgp-PlBd9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Unzip the dataset and rename the folder\n",
        "import os\n",
        "\n",
        "# Set Paths to dataset and dataset filename\n",
        "dataset_foldername = \"vespA_combined2024\"  # @param {type:\"string\"}\n",
        "\n",
        "# Rename the extracted folder to dataset\n",
        "old_path = f'/content/{dataset_foldername}'\n",
        "new_path = '/content/dataset'\n",
        "os.rename(old_path, new_path)"
      ],
      "metadata": {
        "id": "QD5kCxm0YMrD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "4. Optional: Check is path is correct\n",
        "\n",
        "  The output should be:\n",
        "  ```\n",
        "  data.yaml train valid\n",
        "  images labels\n",
        "  images labels\n",
        "  ```\n"
      ],
      "metadata": {
        "id": "WGlnv5GXtaz5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Optional check dataset folder structure and data.yaml\n",
        "\n",
        "dataset_path = \"/content/gdrive/MyDrive/vespA/100_1_dataset.zip\"  # @param {type:\"string\"}\n",
        "\n",
        "!ls '/content/dataset/'\n",
        "!ls '/content/dataset/train/'\n",
        "!ls '/content/dataset/valid/'\n"
      ],
      "metadata": {
        "id": "bbYyE7zctW_4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 4. Train YOLO11\n",
        "\n",
        "Alternatively, you can train or download a model from Ultralytics Hub and use 'best.pt' to convert it in step five to 'best.onnx'.\n",
        "\n",
        "To train the YOLO11 model we need the dataset with the structure as described in step 1.\n",
        "\n",
        "\n",
        "1. Install ultralytics and wandb (weigths and biasis). The latter install is optional, you can remove these lines of code if you do not want to monitor your training metrics)\n",
        "2. Optional: Track your yolo training with Weight and Biases.\n",
        "3. Train your YOLO11 model.\n",
        "\n"
      ],
      "metadata": {
        "id": "QJ8PTAwjYv82"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. Install ultralytics package and wandb"
      ],
      "metadata": {
        "id": "EN0KKR_rY1Z8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Installing ultralytics and wandb\n",
        "!pip install -U ultralytics wandb"
      ],
      "metadata": {
        "id": "GFyr1lY49cBE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. Optional: Track with Weights and Biases\n",
        "\n",
        "  Ultralytics YOLO11 integration with Weights & Biases for enhanced experiment tracking, model-checkpointing, and visualization of model performance. For more information click [Weights and Biasis](https://docs.ultralytics.com/integrations/weights-biases/)\n",
        "  1. Make account and a project in [W&B](www.wandb.ai).\n",
        "  2. Store your API key in the Secrets in Google Colab\n",
        "In your Colab notebook, click the key icon in the left sidebar (this is the secrets manager).\n",
        "  3. Add a secret with the name WANDB_API_KEY and paste your WandB API key as its value. You can find your api key [here](https://wandb.ai/authorize)."
      ],
      "metadata": {
        "id": "1qWpahazk3Jw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import wandb\n",
        "from ultralytics import YOLO\n",
        "from google.colab import userdata\n",
        "\n",
        "# Retrieve the API key from Colab secrets\n",
        "wandb_api_key = userdata.get('WANDB_API_KEY')\n",
        "\n",
        "# Initialize your Weights & Biases environment\n",
        "wandb.login(key=wandb_api_key)"
      ],
      "metadata": {
        "id": "PnQRvg_NzQv-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "3. Train your YOLO11 model.\n",
        "\n",
        "In this setting it uses the default values. All hyperparameters can be found on the [YOLO11 documentation](https://docs.ultralytics.com/modes/train/#augmentation-settings-and-hyperparameters)."
      ],
      "metadata": {
        "id": "83pKvYwpFkc2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import wandb\n",
        "from ultralytics import YOLO\n",
        "from google.colab import userdata\n",
        "\n",
        "# Retrieve the API key from Colab secrets\n",
        "wandb_api_key = userdata.get('WANDB_API_KEY')\n",
        "\n",
        "# Initialize your Weights & Biases environment\n",
        "wandb.login(key=wandb_api_key)\n",
        "\n",
        "# Load a YOLO model\n",
        "model = YOLO(\"yolo11n.pt\")\n",
        "\n",
        "# Train and Fine-Tune the Model\n",
        "model.train(data=\"/content/dataset/data.yaml\",\n",
        "            epochs=100,\n",
        "            lr0=0.01,\n",
        "            imgsz=640,\n",
        "            batch=16, # Use batch=-1 to automatically adjust batch size for approximately 60% CUDA memory utilization.\n",
        "            patience=20,\n",
        "            augment=True,\n",
        "            warmup_epochs=3,\n",
        "            optimizer=\"auto\",\n",
        "            project=\"ultralytics\",\n",
        "            name=\"yolo11n\"\n",
        "            )"
      ],
      "metadata": {
        "id": "VQ5NgE46Fj2n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 5. Convert .pt file to ONNX\n",
        "\n",
        "1. Export to [onnx](https://onnx.ai/onnx/intro/) format (Open Neural Network Exchange). This is an input format the Hailo's DataFlow Compiler can handle.\n",
        "- Changing the opset is similar to upgrading a library. onnx and onnx runtimes must support backward compatibility. Default it is the latest version.\n",
        "- CHECK: change model_path to /content/ultralytics/..."
      ],
      "metadata": {
        "id": "JjZnAN98Owjy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the model path (outside a code cell)\n",
        "model_path = \"/content/ultralytics/yolo11n/weights/best.pt\"  #@param {type:\"string\"}\n",
        "\n",
        "# Import libraries (within a code cell)\n",
        "from ultralytics import YOLO\n",
        "\n",
        "# Load the YOLO11 model (within a code cell)\n",
        "model = YOLO(model_path)\n",
        "\n",
        "# Export the model to ONNX format (within a code cell)\n",
        "model.export(format=\"onnx\")  # creates 'yolo11n.onnx'"
      ],
      "metadata": {
        "id": "NvjSpDTjOu-3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. Verify ONNX model validity. Expected output (example):\n",
        "\n",
        "```\n",
        "  ONNX model is valid!\n",
        "  [[[5.85030460e+00 1.06851845e+01 1.97524452e+01 ... 2.14798126e+02\n",
        "   2.52929199e+02 2.87759094e+02]\n",
        "  ...\n",
        "  [2.34663486e-04 1.97619200e-04 1.53064728e-04 ... 2.67028809e-04\n",
        "   3.86625528e-04 4.52518463e-04]]]\n",
        "   ```"
      ],
      "metadata": {
        "id": "okUoGIyTtI67"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Check the .onnx file\n",
        "\n",
        "import onnx\n",
        "import onnxruntime as ort\n",
        "import torch\n",
        "\n",
        "# Set path to .onnx file\n",
        "onnx_model_path = '/content/ultralytics/yolo11n/weights/best.onnx'  # Path to your ONNX file\n",
        "\n",
        "# Load the ONNX model\n",
        "onnx_model = onnx.load(onnx_model_path)  # Load the ONNX model\n",
        "onnx.checker.check_model(onnx_model)  # Validate the model\n",
        "print(\"ONNX model is valid!\")\n",
        "\n",
        "# Test the ONNX model with ONNX Runtime\n",
        "dummy_input = torch.randn(1, 3, 640, 640).numpy()  # Adjust size to match model input (check Netron)\n",
        "ort_session = ort.InferenceSession(onnx_model_path)  # Pass file path instead of the loaded model\n",
        "outputs = ort_session.run(None, {\"images\": dummy_input})  # Match input name to ONNX model\n",
        "print(outputs[0])\n",
        "\n"
      ],
      "metadata": {
        "id": "8OkyENtCsuPp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 6. Zip the folder\n",
        "\n",
        "Zip the folder with yolo training session(s).\n",
        "\n",
        "When an **error** occurs select runtime -> restart session and run this codeblock again.\n",
        "\n",
        "This code downloads the zipped file to Google Drive and to your local computer.\n"
      ],
      "metadata": {
        "id": "oAZVq5QKKB2j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "from google.colab import files\n",
        "\n",
        "import os\n",
        "import shutil\n",
        "\n",
        "# Set foldername and filename\n",
        "train_folder_path = \"/content/ultralytics/\" #@param {type:\"string\"}\n",
        "download_file_name = \"/content/yolo11n.zip\" #@param {type:\"string\"}\n",
        "\n",
        "# mount google drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "# zip the ultralytics folder with training result(s) folders\n",
        "try:\n",
        "  # Zipping the folder\n",
        "  !zip -r {download_file_name} {train_folder_path}\n",
        "  # Downloading the zipped file\n",
        "  files.download(download_file_name)\n",
        "except Exception as e:\n",
        "  print(f\"An error occurred: {e}\")\n",
        "  print(\"Click 'Runtime' -> 'Restart session' and try running the code again.\")\n",
        "\n",
        "# Source file path in Colab\n",
        "source_file = 'download_file_name'\n",
        "\n",
        "# Destination directory in Google Drive\n",
        "drive_path = '/content/drive/MyDrive'\n",
        "dest_dir = os.path.join(drive_path, 'ultralytics')\n",
        "\n",
        "# Create the destination directory if it doesn't exist\n",
        "if not os.path.exists(dest_dir):\n",
        "    os.makedirs(dest_dir)\n",
        "\n",
        "# Copy the file to the destination directory\n",
        "shutil.copy(download_file_name, dest_dir)"
      ],
      "metadata": {
        "id": "S4HurBi1MXCh"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}