examples/HOWTO_train_YOLO_v9_with_Ikomia_API.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GyGbD_GxAOI0"
      },
      "source": [
        "<img src='https://github.com/Ikomia-dev/notebooks/blob/main/examples/img/banner_ikomia.png?raw=true'>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WKz9Zij7bN-t"
      },
      "source": [
        "# How to train YOLOv9"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CpxHrXhtbWTZ"
      },
      "source": [
        "## Overview of YOLOv9\n",
        "With the continuous evolution of computer vision technologies, YOLOv9 emerges as the latest advancement, developed by Chien-Yao Wang, I-Hau Yeh, and Hong-Yuan Mark Liao. This trio of researchers has a rich history in the field, having contributed to the development of preceding models such as YOLOv4, YOLOR, and YOLOv7.\n",
        "\n",
        "### 🌖 Release\n",
        "\n",
        "YOLOv9 not only continues the legacy of its predecessors but also introduces significant innovations that set new benchmarks in object detection capabilities.\n",
        "\n",
        "\n",
        "YOLOv9 is an advanced object detection model that represents a significant leap forward in computer vision technology. It is the latest iteration in the \"You Only Look Once\" (YOLO) series, known for its high speed and accuracy in detecting objects in images.\n",
        "\n",
        "\n",
        "### 🌆 Architecture and innovations:\n",
        "\n",
        "YOLOv9 stands out due to its incorporation of Programmable Gradient Information (PGI) and the introduction of the Generalized Efficient Layer Aggregation Network (GELAN), two groundbreaking innovations designed to enhance model performance and efficiency.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "### :rocket: Accuracy and performance:\n",
        "\n",
        "| Model | Test Size | AP<sup>val</sup> | AP<sub>50</sub><sup>val</sup> | AP<sub>75</sub><sup>val</sup> | Param. | FLOPs |\n",
        "| :-- | :-: | :-: | :-: | :-: | :-: | :-: |\n",
        "| **YOLOv9-S** | 640 | **46.8%** | **63.4%** | **50.7%** | **7.2M** | **26.7G** |\n",
        "| **YOLOv9-M** | 640 | **51.4%** | **68.1%** | **56.1%** | **20.1M** | **76.8G** |\n",
        "| **YOLOv9-C** | 640 | **53.0%** | **70.2%** | **57.8%** | **25.5M** | **102.8G** |\n",
        "| **YOLOv9-E** | 640 | **55.6%** | **72.8%** | **60.6%** | **58.1M** | **192.5G** |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x4CdI0J1ej5b"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5O31W412NRx"
      },
      "source": [
        "Please use a GPU for this tutorial.\n",
        "\n",
        "In the menu, select \"Runtime\" then \"Change runtime type\", choose GPU in \"Hardware accelerator\".\n",
        "\n",
        "Check your GPU with the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OJFMsi47Yrqj"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NBmJN2AaDmcI"
      },
      "source": [
        "First of all, you need to install Ikomia API pip package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8eSnQYJygrDy"
      },
      "outputs": [],
      "source": [
        "!pip install ikomia"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ktbA-VPOATgP"
      },
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "**-Google Colab ONLY- Restart runtime**\n",
        "\n",
        "Click on the \"RESTART RUNTIME\" button at the end the previous window.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2hS1T6ky1Wcw"
      },
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JJsRFzl9Au1c"
      },
      "source": [
        "Ikomia API has already more than 300 pre-integrated algorithms but the most interesting algorithms are in [Ikomia HUB](https://github.com/Ikomia-hub).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZgGi5tjosC8g"
      },
      "source": [
        "## How to train YOLOv9 on a custom dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pwJit3k190PN"
      },
      "source": [
        "Download your dataset from your preferred tool. In this example, we use a dataset from **Roboflow** which is a great annotation platform used by many developers and companies. The dataset is exported in COCO format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RpGiTWgeRnIq",
        "tags": []
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "import zipfile\n",
        "import os\n",
        "\n",
        "\n",
        "# Download the file\n",
        "url = \"https://universe.roboflow.com/ds/OG5R1BjHfr?key=U5vq9IGzOH\"\n",
        "response = requests.get(url, stream=True)\n",
        "with open(\"roboflow.zip\", \"wb\") as file:\n",
        "    for chunk in response.iter_content(chunk_size=8192):\n",
        "        file.write(chunk)\n",
        "\n",
        "# Unzip the file\n",
        "with zipfile.ZipFile(\"roboflow.zip\", 'r') as zip_ref:\n",
        "    zip_ref.extractall()\n",
        "\n",
        "# Remove the zip file\n",
        "os.remove(\"roboflow.zip\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BgYm8QNx-ekO"
      },
      "source": [
        "In order to train YOLOv9 on your custom dataset, please create a new workflow from scratch.\n",
        "\n",
        "Then you need 2 components:\n",
        "\n",
        "1.   A COCO dataset loader which loads dataset in COCO format and convert it to an Ikomia format\n",
        "2.   The YOLOv9 training algorithm which loads dataset in Ikomia format\n",
        "\n",
        "Add these 2 previous algorithms to your workflow and then it will automagically download all algorithms from Ikomia Hub and install all the Python dependencies (the 1st time, it can take a while, be patient ! )."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I1BKL7hpw15Z"
      },
      "source": [
        "Now, it's time to train your model !"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AW06ne0oBNtK",
        "tags": []
      },
      "outputs": [],
      "source": [
        "from ikomia.dataprocess.workflow import Workflow\n",
        "import os\n",
        "\n",
        "#----------------------------- Step 1 -----------------------------------#\n",
        "# Create a workflow which will take your dataset as input and\n",
        "# train a YOLOv9 model on it\n",
        "#------------------------------------------------------------------------#\n",
        "wf = Workflow()\n",
        "\n",
        "#----------------------------- Step 2 -----------------------------------#\n",
        "# First you need to convert the COCO format to IKOMIA format.\n",
        "# Add an Ikomia dataset converter to your workflow.\n",
        "#------------------------------------------------------------------------#\n",
        "\n",
        "dataset = wf.add_task(name = \"dataset_coco\")\n",
        "\n",
        "dataset.set_parameters({\n",
        "    \"json_file\": os.getcwd()+\"/train/_annotations.coco.json\",\n",
        "    \"image_folder\": os.getcwd()+\"/train\",\n",
        "    \"task\":\"detection\",\n",
        "    \"output_folder\": os.getcwd()+\"/dataset\"\n",
        "})\n",
        "\n",
        "\n",
        "#----------------------------- Step 3 -----------------------------------#\n",
        "# Then, you want to train a YOLOv9 model.\n",
        "# Add YOLOv9 training algorithm to your workflow\n",
        "#------------------------------------------------------------------------#\n",
        "\n",
        "train = wf.add_task(name=\"train_yolo_v9\", auto_connect=True)\n",
        "train.set_parameters({\n",
        "    \"model_name\": \"yolov9-c\",\n",
        "    \"epochs\": \"50\",\n",
        "    \"batch_size\": \"8\",\n",
        "    \"train_imgsz\": \"640\",\n",
        "    \"test_imgsz\": \"640\",\n",
        "    \"dataset_split_ratio\": \"0.8\",\n",
        "    \"output_folder\": os.getcwd(),\n",
        "})\n",
        "\n",
        "\n",
        "#----------------------------- Step 4 -----------------------------------#\n",
        "# Execute your workflow.\n",
        "# It automatically runs all your tasks sequentially.\n",
        "#------------------------------------------------------------------------#\n",
        "wf.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rpsaQoYSwma8"
      },
      "source": [
        "## Infer YOLOv9 object detection on images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-H2JCp2SOD2a"
      },
      "source": [
        "Once the training is finished, you may want to experiment the fresh model on new test images. Just use the following code to create a YOLO v8 instance\n",
        "segmentation inference workflow.\n",
        "Then run and test !"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sVH02tztrM4F",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# Set the output folder name\n",
        "TIMESTAMP = # \"27-02-2024T09h27m06s\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ecT5qPDywrEi",
        "tags": []
      },
      "outputs": [],
      "source": [
        "from ikomia.dataprocess.workflow import Workflow\n",
        "\n",
        "# Create your workflow for YOLO inference\n",
        "wf = Workflow()\n",
        "\n",
        "# Add YOLOv9 to your workflow\n",
        "yolov9 = wf.add_task(name=\"infer_yolo_v9\", auto_connect=True)\n",
        "\n",
        "yolov9.set_parameters({\n",
        "    \"model_weight_file\": os.getcwd()+ f'/{TIMESTAMP}/weights/best.pt',\n",
        "    \"class_file\": os.getcwd()+ f'/{TIMESTAMP}/classes.yaml',\n",
        "    \"conf_thres\": \"0.2\",\n",
        "    \"iou_thres\":\"0.25\",\n",
        "    \"input_size\":\"640\"\n",
        "})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pKPnRdUj6Vq8"
      },
      "source": [
        "## Run and display your results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YdXFRaDi6Vq8",
        "tags": []
      },
      "outputs": [],
      "source": [
        "from ikomia.utils.displayIO import display\n",
        "from PIL import ImageShow\n",
        "ImageShow.register(ImageShow.IPythonViewer(), 0)\n",
        "\n",
        "# Apply YOLOv9 object detection on your image\n",
        "# By default, YOLOv9 runs with a pre-trained model based on COCO\n",
        "# To use your custom model, set the parameters in the previous cell\n",
        "\n",
        "wf.run_on(url=\"https://pbs.twimg.com/ext_tw_video_thumb/1660454979298115585/pu/img/A_Jrl2uawkkDi_Kf.jpg\")\n",
        "# wf.run_on(path=os.getcwd()+\"/test/youtube-128_jpg.rf.2723e31eec77e1ff7b73c45c625082f6.jpg\")\n",
        "\n",
        "# Get YOLOv9 image result\n",
        "img_bbox = yolov9.get_image_with_graphics()\n",
        "\n",
        "# Display in Colab\n",
        "display(img_bbox)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oF5CAKdBrM4F"
      },
      "source": [
        "## Run on video\n",
        "\n",
        "This will work on local only, not on google colab"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g0F-qSDsrM4G"
      },
      "outputs": [],
      "source": [
        "from ikomia.dataprocess.workflow import Workflow\n",
        "from ikomia.utils.displayIO import display\n",
        "import cv2\n",
        "\n",
        "\n",
        "video_path = 'Path/to/your/video.mp4' # Example: https://www.youtube.com/watch?v=EAR5jTknVOw\n",
        "output_path = 'output.mp4'\n",
        "# Init your workflow\n",
        "wf = Workflow()\n",
        "\n",
        "# Add object detection algorithm\n",
        "detector = wf.add_task(name=\"infer_yolo_v9\", auto_connect=True)\n",
        "\n",
        "detector.set_parameters({\n",
        "    \"model_weight_file\": os.getcwd()+ f'/{TIMESTAMP}/weights/best.pt',\n",
        "    \"class_file\": os.getcwd()+ f'/{TIMESTAMP}/classes.yaml',\n",
        "    \"conf_thres\": \"0.2\",\n",
        "    \"iou_thres\":\"0.25\"\n",
        "})\n",
        "\n",
        "# Open the video file\n",
        "stream = cv2.VideoCapture(video_path)\n",
        "if not stream.isOpened():\n",
        "    print(\"Error: Could not open video.\")\n",
        "    exit()\n",
        "\n",
        "# Get video properties for the output\n",
        "frame_width = int(stream.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "frame_height = int(stream.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "frame_rate = stream.get(cv2.CAP_PROP_FPS)\n",
        "\n",
        "# Define the codec and create VideoWriter object\n",
        "# The 'XVID' codec is widely supported and provides good quality\n",
        "fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
        "out = cv2.VideoWriter(output_path, fourcc, frame_rate, (frame_width, frame_height))\n",
        "\n",
        "while True:\n",
        "    # Read image from stream\n",
        "    ret, frame = stream.read()\n",
        "\n",
        "    # Test if the video has ended or there is an error\n",
        "    if not ret:\n",
        "        print(\"Info: End of video or error.\")\n",
        "        break\n",
        "\n",
        "    # Run the workflow on current frame\n",
        "    wf.run_on(array=frame)\n",
        "\n",
        "    # Get results\n",
        "    image_out = detector.get_output(0)\n",
        "    obj_detect_out = detector.get_output(1)\n",
        "\n",
        "    # Convert the result to BGR color space for displaying\n",
        "    img_out = image_out.get_image_with_mask_and_graphics(obj_detect_out)\n",
        "    img_res = cv2.cvtColor(img_out, cv2.COLOR_RGB2BGR)\n",
        "\n",
        "    # Save the resulting frame\n",
        "    out.write(img_out)\n",
        "\n",
        "    # Display\n",
        "    display(img_res, title=\"YOLOv9 object detection\", viewer=\"opencv\")\n",
        "\n",
        "    # Press 'q' to quit the video processing\n",
        "    if cv2.waitKey(1) & 0xFF == ord('q'):\n",
        "        break\n",
        "\n",
        "# After the loop release everything\n",
        "stream.release()\n",
        "out.release()\n",
        "cv2.destroyAllWindows()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "environment": {
      "kernel": "venvsd1",
      "name": "common-gpu.m114",
      "type": "gcloud",
      "uri": "gcr.io/deeplearning-platform-release/base-gpu:m114"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}