diff --git a/Train_YOLO_Models.ipynb b/Train_YOLO_Models.ipynb new file mode 100644 index 0000000000..644e16a12f --- /dev/null +++ b/Train_YOLO_Models.ipynb @@ -0,0 +1,937 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1sUfcA8ZgR2t" + }, + "source": [ + "# Train YOLO Models in Google Colab\n", + "**Author:** Evan Juras, [EJ Technology Consultants](https://ejtech.io)\n", + "\n", + "**Last updated:** January 3, 2025\n", + "\n", + "**GitHub:** [Train and Deploy YOLO Models](https://github.com/EdjeElectronics/Train-and-Deploy-YOLO-Models)\n", + "\n", + "# Introduction\n", + "\n", + "This notebook uses [Ultralytics](https://docs.ultralytics.com/) to train YOLO11, YOLOv8, or YOLOv5 object detection models with a custom dataset. At the end of this Colab, you'll have a custom YOLO model that you can run on your PC, phone, or edge device like the Raspberry Pi.\n", + "\n", + "

\n", + "
\n", + "Custom YOLO candy detection model in action!\n", + "

\n", + "\n", + "I created a YouTube video that walks through this guide step by step. I recommend following along with the video while working through this notebook.\n", + "\n", + "

\n", + "
\n", + "Click here to go to the video!\n", + "

\n", + "\n", + "**Important note: This notebook will be continuously updated to make sure it works with newer versions of Ultralytics and YOLO. If you see any differences between the YouTube video and this notebook, always follow the notebook!**\n", + "\n", + "### Working in Colab\n", + "Colab provides a virtual machine in your browser complete with a Linux OS, filesystem, Python environment, and best of all, a free GPU. We'll install PyTorch and Ultralytics in this environment and use it to train our model. Simply click the Play button on sections of code in this notebook to execute them on the virtual machine.\n", + "\n", + "### Navigation\n", + "To navigate this notebook, use the table of contents in the left sidebar to jump from section to section.\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "source": [ + "**Verify NVIDIA GPU Availability**\n", + "\n", + "Make sure you're using a GPU-equipped machine by going to \"Runtime\" -> \"Change runtime type\" in the top menu bar, and then selecting one of the GPU options in the Hardware accelerator section. Click Play on the following code block to verify that the NVIDIA GPU is present and ready for training." + ], + "metadata": { + "id": "3NW7LLv_QPOO" + } + }, + { + "cell_type": "code", + "source": [ + "!nvidia-smi" + ], + "metadata": { + "id": "cfaWho47RGDf", + "outputId": "4bc2b452-6c03-4b5e-8be1-69315e497f35", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 1, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Sun Jun 22 02:11:18 2025 \n", + "+-----------------------------------------------------------------------------------------+\n", + "| NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: 12.4 |\n", + "|-----------------------------------------+------------------------+----------------------+\n", + "| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\n", + "| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\n", + "| | | MIG M. |\n", + "|=========================================+========================+======================|\n", + "| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\n", + "| N/A 41C P8 9W / 70W | 0MiB / 15360MiB | 0% Default |\n", + "| | | N/A |\n", + "+-----------------------------------------+------------------------+----------------------+\n", + " \n", + "+-----------------------------------------------------------------------------------------+\n", + "| Processes: |\n", + "| GPU GI CI PID Type Process name GPU Memory |\n", + "| ID ID Usage |\n", + "|=========================================================================================|\n", + "| No running processes found |\n", + "+-----------------------------------------------------------------------------------------+\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TIHu25pnjjJ1" + }, + "source": [ + "#1. Gather and Label Training Images" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "O6Y1vBiRjpcq" + }, + "source": [ + "Before we start training, we need to gather and label images that will be used for training the object detection model. A good starting point for a proof-of-concept model is 200 images. The training images should have random objects in the image along with the desired objects, and should have a variety of backgrounds and lighting conditions.\n", + "\n", + "There are a couple options for gathering images:\n", + "\n", + "\n", + "* Build a custom dataset by taking your own pictures of the objects and labeling them (this typically results in the best performance)\n", + "* Find a pre-made dataset from sources like [Roboflow Universe](), [Kaggle](), or [Google Images V7]()\n", + "\n", + "\n", + "If you want to build your own dataset, there are several tools available for labeling images. One good option is [Label Studio](https://labelstud.io/?utm_source=youtube&utm_medium=video&utm_campaign=edjeelectronics), a free and open-source labeling tool that has a simple workflow while providing capabilities for more advanced features. My YouTube video that walks through this notebook (link to be added soon) shows how to label images with Label Studio.\n", + "\n", + "

\n", + "
\n", + "Example of a candy image labeled with Label Studio.\n", + "

\n", + "\n", + "If you used Label Studio to label and export the images, they'll be exported in a `project.zip` file that contains the following:\n", + "\n", + "- An `images` folder containing the images\n", + "- A `labels` folder containing the labels in YOLO annotation format\n", + "- A `classes.txt` labelmap file that contains all the classes\n", + "- A `notes.json` file that contains info specific to Label Studio (this file can be ignored)\n", + "\n", + "If you obtained your dataset from another source (like Roboflow Universe) or used another tool to label your dataset, make sure the files are organized in the same folder structure.\n", + "\n", + "

\n", + "
\n", + "Organize your data in the folders shown here. See my Candy Detection Dataset for an example.\n", + "

\n", + "\n", + "Once you've got your dataset built, put into the file structure shown above, and zipped into `data.zip`, you're ready to move on to the next step." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8eDhuvzDfIFS" + }, + "source": [ + "# 2. Upload Image Dataset and Prepare Training Data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ZW_0c110fOiz" + }, + "source": [ + "Next, we'll upload our dataset and prepare it for training with YOLO. We'll split the dataset into train and validation folders, and we'll automatically generate the configuration file for training the model." + ] + }, + { + "cell_type": "markdown", + "source": [ + "## 2.1 Upload images\n", + "\n", + "First, we need to upload the dataset to Colab. Here are a few options for moving the `data.zip` folder into this Colab instance." + ], + "metadata": { + "id": "FwKAqFIQSBpn" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Option 1. Upload through Google Colab**\n", + "\n", + "Upload the `data.zip` file to the Google Colab instance by clicking the \"Files\" icon on the left hand side of the browser, and then the \"Upload to session storage\" icon. Select the zip folder to upload it.\n", + "\n", + "

\n", + "
\n", + "\n", + "

\n" + ], + "metadata": { + "id": "ZPZEM27IOh79" + } + }, + { + "cell_type": "markdown", + "source": [ + "\n", + "**Option 2. Copy from Google Drive**\n", + "\n", + "You can also upload your images to your personal Google Drive, mount the drive on this Colab session, and copy them over to the Colab filesystem. This option works well if you want to upload the images beforehand so you don't have to wait for them to upload each time you restart this Colab. If you have more than 50MB worth of images, I recommend using this option.\n", + "\n", + "First, upload the `data.zip` file to your Google Drive, and make note of the folder you uploaded them to. Replace `MyDrive/path/to/data.zip` with the path to your zip file. (For example, I uploaded the zip file to folder called \"candy-dataset1\", so I would use `MyDrive/candy-dataset1/data.zip` for the path). Then, run the following block of code to mount your Google Drive to this Colab session and copy the folder to this filesystem." + ], + "metadata": { + "id": "TC4bZM1UWRdY" + } + }, + { + "cell_type": "code", + "source": [ + "from google.colab import drive\n", + "drive.mount('/content/gdrive')\n", + "\n", + "\n", + "!cp /content/gdrive/MyDrive/yolo.zip /content" + ], + "metadata": { + "id": "ZfQBSwDdWoWp", + "outputId": "a6ed27dd-8211-4cf7-c44f-f9abde49bde9", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 4, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Drive already mounted at /content/gdrive; to attempt to forcibly remount, call drive.mount(\"/content/gdrive\", force_remount=True).\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "**Option 3. Use my candy detection or coin detection dataset**\n", + "\n", + "If you just want to test the process on a pre-made dataset, you can use one of my datasets:\n", + "\n", + "* [Candy image dataset](https://s3.us-west-1.amazonaws.com/evanjuras.com/resources/candy_data_14DEC24.zip), which contains 162 pictures of popular candies (Skittles, Snickers, etc)\n", + "* [Coin image dataset](https://s3.us-west-1.amazonaws.com/evanjuras.com/resources/YOLO_coin_data_12DEC30.zip), which contains 750 pictures of US coins (pennies, dimes, nickels, and quarters)\n", + "\n", + "Download one of the datasets by running the following code block. I'll use the candy detection dataset as the example for the rest of the notebook." + ], + "metadata": { + "id": "q43_b9-sWsdB" + } + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ZQXLBvL5grDl" + }, + "outputs": [], + "source": [ + "# To use my one of pre-made dataset instead of your own custom dataset, download it here (control which dataset is downloaded by commenting out either line)\n", + "!wget -O /content/data.zip https://s3.us-west-1.amazonaws.com/evanjuras.com/resources/candy_data_06JAN25.zip # Candy dataset\n", + "#!wget -O /content/data.zip https://s3.us-west-1.amazonaws.com/evanjuras.com/resources/YOLO_coin_data_12DEC30.zip # Coin dataset" + ] + }, + { + "cell_type": "markdown", + "source": [ + "## 2.2 Split images into train and validation folders" + ], + "metadata": { + "id": "m7Iz9eBzW5zm" + } + }, + { + "cell_type": "markdown", + "source": [ + "At this point, whether you used Option 1, 2, or 3, you should be able to click the folder icon on the left and see your `data.zip` file in the list of files. Next, we'll unzip `data.zip` and create some folders to hold the images. Run the following code block to unzip the data." + ], + "metadata": { + "id": "58JuFGc2PatU" + } + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "id": "z8O6z-wVcPEF", + "outputId": "4ef10a41-2117-4be2-fa58-0990dda25c33", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "replace /content/custom_data/yolo/classes.txt? [y]es, [n]o, [A]ll, [N]one, [r]ename: y\n", + "replace /content/custom_data/yolo/desktop.ini? [y]es, [n]o, [A]ll, [N]one, [r]ename: y\n", + "replace /content/custom_data/yolo/images/1 (10).jpg? [y]es, [n]o, [A]ll, [N]one, [r]ename: y\n", + "replace /content/custom_data/yolo/images/1 (11).jpg? [y]es, [n]o, [A]ll, [N]one, [r]ename: y\n", + "replace /content/custom_data/yolo/images/1 (12).jpg? [y]es, [n]o, [A]ll, [N]one, [r]ename: A\n" + ] + } + ], + "source": [ + "# Unzip images to a custom data folder\n", + "# Unzip images to a custom data folder\n", + "!unzip -q /content/gdrive/MyDrive/yolo.zip -d /content/custom_data" + ] + }, + { + "cell_type": "markdown", + "source": [ + "Ultralytics requires a particular folder structure to store training data for models. Ultralytics requires a particular folder structure to store training data for models. The root folder is named “data”. Inside, there are two main folders:\n", + "\n", + "* **Train**: These are the actual images used to train the model. In one epoch of training, every image in the train set is passed into the neural network. The training algorithm adjusts the network weights to fit the data in the images.\n", + "\n", + "\n", + "* **Validation**: These images are used to check the model's performance at the end of each training epoch.\n", + "\n", + "In each of these folders is a “images” folder and a “labels” folder, which hold the image files and annotation files respectively." + ], + "metadata": { + "id": "eoPjqW6AYebn" + } + }, + { + "cell_type": "markdown", + "source": [ + "\n", + "I wrote a Python script that will automatically create the required folder structure and randomly move 90% of dataset to the \"train\" folder and 10% to the \"validation\" folder. Run the following code block to download and execute the scrpt." + ], + "metadata": { + "id": "f2ohNAhWj41n" + } + }, + { + "cell_type": "code", + "source": [ + "!wget -O /content/train_val_split.py https://raw.githubusercontent.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/refs/heads/main/utils/train_val_split.py\n", + "\n", + "# TO DO: Improve robustness of train_val_split.py script so it can handle nested data folders, etc\n", + "!python train_val_split.py --datapath=\"/content/custom_data/yolo\" --train_pct=0.9" + ], + "metadata": { + "id": "8X62eFTugosf", + "outputId": "e334cf3e-9a24-410e-b88d-46ff82846b13", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 10, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "--2025-06-22 02:55:39-- https://raw.githubusercontent.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/refs/heads/main/utils/train_val_split.py\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 3203 (3.1K) [text/plain]\n", + "Saving to: ‘/content/train_val_split.py’\n", + "\n", + "\r /content/ 0%[ ] 0 --.-KB/s \r/content/train_val_ 100%[===================>] 3.13K --.-KB/s in 0s \n", + "\n", + "2025-06-22 02:55:39 (66.5 MB/s) - ‘/content/train_val_split.py’ saved [3203/3203]\n", + "\n", + "Number of image files: 596\n", + "Number of annotation files: 596\n", + "Images moving to train: 536\n", + "Images moving to validation: 60\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "# 3. Install Requirements (Ultralytics)\n", + "\n", + "Next, we'll install the Ultralytics library in this Google Colab instance. This Python library will be used to train the YOLO model." + ], + "metadata": { + "id": "B2L2qGCJzwY9" + } + }, + { + "cell_type": "code", + "source": [ + "!pip install ultralytics" + ], + "metadata": { + "id": "EMEDk5byzxY5", + "outputId": "b00fbc62-524e-4c2b-e202-5b3ded779b97", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 11, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Requirement already satisfied: ultralytics in /usr/local/lib/python3.11/dist-packages (8.3.158)\n", + "Requirement already satisfied: numpy>=1.23.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (2.0.2)\n", + "Requirement already satisfied: matplotlib>=3.3.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (3.10.0)\n", + "Requirement already satisfied: opencv-python>=4.6.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (4.11.0.86)\n", + "Requirement already satisfied: pillow>=7.1.2 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (11.2.1)\n", + "Requirement already satisfied: pyyaml>=5.3.1 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (6.0.2)\n", + "Requirement already satisfied: requests>=2.23.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (2.32.3)\n", + "Requirement already satisfied: scipy>=1.4.1 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (1.15.3)\n", + "Requirement already satisfied: torch>=1.8.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (2.6.0+cu124)\n", + "Requirement already satisfied: torchvision>=0.9.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (0.21.0+cu124)\n", + "Requirement already satisfied: tqdm>=4.64.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (4.67.1)\n", + "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from ultralytics) (5.9.5)\n", + "Requirement already satisfied: py-cpuinfo in /usr/local/lib/python3.11/dist-packages (from ultralytics) (9.0.0)\n", + "Requirement already satisfied: pandas>=1.1.4 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (2.2.2)\n", + "Requirement already satisfied: ultralytics-thop>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from ultralytics) (2.0.14)\n", + "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (1.3.2)\n", + "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (0.12.1)\n", + "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (4.58.4)\n", + "Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (1.4.8)\n", + "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (24.2)\n", + "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (3.2.3)\n", + "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.3.0->ultralytics) (2.9.0.post0)\n", + "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas>=1.1.4->ultralytics) (2025.2)\n", + "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas>=1.1.4->ultralytics) (2025.2)\n", + "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests>=2.23.0->ultralytics) (3.4.2)\n", + "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests>=2.23.0->ultralytics) (3.10)\n", + "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests>=2.23.0->ultralytics) (2.4.0)\n", + "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests>=2.23.0->ultralytics) (2025.6.15)\n", + "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (3.18.0)\n", + "Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (4.14.0)\n", + "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (3.5)\n", + "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (3.1.6)\n", + "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (2025.3.2)\n", + "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.4.127)\n", + "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.4.127)\n", + "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.4.127)\n", + "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (9.1.0.70)\n", + "Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.4.5.8)\n", + "Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (11.2.1.3)\n", + "Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (10.3.5.147)\n", + "Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (11.6.1.9)\n", + "Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.3.1.170)\n", + "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.2 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (0.6.2)\n", + "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (2.21.5)\n", + "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.4.127)\n", + "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (12.4.127)\n", + "Requirement already satisfied: triton==3.2.0 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (3.2.0)\n", + "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch>=1.8.0->ultralytics) (1.13.1)\n", + "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch>=1.8.0->ultralytics) (1.3.0)\n", + "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/dist-packages (from python-dateutil>=2.7->matplotlib>=3.3.0->ultralytics) (1.17.0)\n", + "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch>=1.8.0->ultralytics) (3.0.2)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "cuZoMkSFN9XG" + }, + "source": [ + "# 4. Configure Training\n" + ] + }, + { + "cell_type": "markdown", + "source": [ + "There's one last step before we can run training: we need to create the Ultralytics training configuration YAML file. This file specifies the location of your train and validation data, and it also defines the model's classes. An example configuration file model is available [here](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco128.yaml).\n", + "\n", + "Run the code block below to automatically generate a `data.yaml` configuration file. Make sure you have a labelmap file located at `custom_data/classes.txt`. If you used Label Studio or one of my pre-made datasets, it should already be present. If you assembled the dataset another way, you may have to manually create the `classes.txt` file (see [here](https://github.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/blob/main/doc/classes.txt) for an example of how it's formatted)." + ], + "metadata": { + "id": "0c5Kdh0GmQHS" + } + }, + { + "cell_type": "code", + "source": [ + "# Python function to automatically create data.yaml config file\n", + "# 1. Reads \"classes.txt\" file to get list of class names\n", + "# 2. Creates data dictionary with correct paths to folders, number of classes, and names of classes\n", + "# 3. Writes data in YAML format to data.yaml\n", + "\n", + "import yaml\n", + "import os\n", + "\n", + "def create_data_yaml(path_to_classes_txt, path_to_data_yaml):\n", + "\n", + " # Read class.txt to get class names\n", + " if not os.path.exists(path_to_classes_txt):\n", + " print(f'classes.txt file not found! Please create a classes.txt labelmap and move it to {path_to_classes_txt}')\n", + " return\n", + " with open(path_to_classes_txt, 'r') as f:\n", + " classes = []\n", + " for line in f.readlines():\n", + " if len(line.strip()) == 0: continue\n", + " classes.append(line.strip())\n", + " number_of_classes = len(classes)\n", + "\n", + " # Create data dictionary\n", + " data = {\n", + " 'path': '/content/custom_data/yolo',\n", + " 'train': 'train/images',\n", + " 'val': 'validation/images',\n", + " 'nc': number_of_classes,\n", + " 'names': classes\n", + " }\n", + "\n", + " # Write data to YAML file\n", + " with open(path_to_data_yaml, 'w') as f:\n", + " yaml.dump(data, f, sort_keys=False)\n", + " print(f'Created config file at {path_to_data_yaml}')\n", + "\n", + " return\n", + "\n", + "# Define path to classes.txt and run function\n", + "path_to_classes_txt = '/content/custom_data/yolo/classes.txt'\n", + "path_to_data_yaml = '/content/data.yaml'\n", + "\n", + "create_data_yaml(path_to_classes_txt, path_to_data_yaml)\n", + "\n", + "print('\\nFile contents:\\n')\n", + "!cat /content/data.yaml" + ], + "metadata": { + "id": "4letvP7X12ji", + "outputId": "612b7e19-4810-439d-8154-63c02ce5b1bf", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 12, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Created config file at /content/data.yaml\n", + "\n", + "File contents:\n", + "\n", + "path: /content/custom_data/yolo\n", + "train: train/images\n", + "val: validation/images\n", + "nc: 4\n", + "names:\n", + "- car\n", + "- omnibus\n", + "- bus\n", + "- truck\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "myP80_bnTNMi" + }, + "source": [ + "# 5. Train Model" + ] + }, + { + "cell_type": "markdown", + "source": [ + "## 5.1 Training Parameters\n", + "Now that the data is organized and the config file is created, we're ready to start training! First, there are a few important parameters to decide on. Visit my article on [Training YOLO Models Locally](https://www.ejtech.io/learn/train-yolo-models) to learn more about these parameters and how to choose them.\n", + "\n", + "**Model architecture & size (`model`):**\n", + "\n", + "There are several YOLO11 models sizes available to train, including `yolo11n.pt`, `yolo11s.pt`, `yolo11m.pt`, `yolo11l.pt`, and `yolo11xl.pt`. Larger models run slower but have higher accuracy, while smaller models run faster but have lower accuracy. I made a brief YouTube video that compares performance of different YOLO models on a Raspberry Pi 5 and a laptop with a RTX 4050 GPU, [check it out here to get a sense of their speed accuracy](https://youtu.be/_WKS4E9SmkA). If you aren't sure which model size to use, `yolo11s.pt` is a good starting point.\n", + "\n", + "You can also train YOLOv8 or YOLOv5 models by substituting `yolo11` for `yolov8` or `yolov5`.\n", + "\n", + "\n", + "**Number of epochs (`epochs`)**\n", + "\n", + "In machine learning, one “epoch” is one single pass through the full training dataset. Setting the number of epochs dictates how long the model will train for. The best amount of epochs to use depends on the size of the dataset and the model architecture. If your dataset has less than 200 images, a good starting point is 60 epochs. If your dataset has more than 200 images, a good starting point is 40 epochs.\n", + "\n", + "\n", + "**Resolution (`imgsz`)**\n", + "\n", + "Resolution has a large impact on the speed and accuracy of the model: a lower resolution model will have higher speed but less accuracy. YOLO models are typically trained and inferenced at a 640x640 resolution. However, if you want your model to run faster or know you will be working with low-resolution images, try using a lower resolution like 480x480.\n" + ], + "metadata": { + "id": "DfKspYasCzC8" + } + }, + { + "cell_type": "markdown", + "source": [ + "## 5.2 Run Training!" + ], + "metadata": { + "id": "V17UjYU5ZQdR" + } + }, + { + "cell_type": "markdown", + "source": [ + "In this section i will fine tune my model\n" + ], + "metadata": { + "id": "I9YpgnHhogbk" + } + }, + { + "cell_type": "markdown", + "source": [ + "Run the following code block to begin training. If you want to use a different model, number of epochs, or resolution, change `model`, `epochs`, or `imgsz`." + ], + "metadata": { + "id": "nQi_hXnUVPr-" + } + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8bbpob1gTPlo" + }, + "outputs": [], + "source": [ + "!yolo detect train data=/content/data.yaml model=yolo11s.pt epochs=60 imgsz=640" + ] + }, + { + "cell_type": "markdown", + "source": [ + "The training algorithm will parse the images in the training and validation directories and then start training the model. At the end of each training epoch, the program runs the model on the validation dataset and reports the resulting mAP, precision, and recall. As training continues, the mAP should generally increase with each epoch. Training will end once it goes through the number of epochs specified by `epochs`.\n", + "\n", + "> **NOTE:** Make sure to allow training to run to completion, because an optimizer runs at the end of training that strips out unneeded layers from the model.\n", + "\n", + "The best trained model weights will be saved in `content/runs/detect/train/weights/best.pt`. Additional information about training is saved in the `content/runs/detect/train` folder, including a `results.png` file that shows how loss, precision, recall, and mAP progressed over each epoch." + ], + "metadata": { + "id": "vv0EYWJ5V6mC" + } + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vo8BJRXeg0Ap" + }, + "source": [ + "#6. Test Model" + ] + }, + { + "cell_type": "markdown", + "source": [ + "The model has been trained; now it's time to test it! The commands below run the model on the images in the validation folder and then display the results for the first 10 images. This is a good way to confirm your model is working as expected. Click Play on the blocks below to see how your model performs." + ], + "metadata": { + "id": "BX3PTrEPacGY" + } + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PooP5Vjsg2Jn" + }, + "outputs": [], + "source": [ + "!yolo detect predict model=runs/detect/train/weights/best.pt source=data/validation/images save=True" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "zEEObQqoiGrs" + }, + "outputs": [], + "source": [ + "import glob\n", + "from IPython.display import Image, display\n", + "for image_path in glob.glob(f'/content/runs/detect/predict/*.jpg')[:10]:\n", + " display(Image(filename=image_path, height=400))\n", + " print('\\n')\n" + ] + }, + { + "cell_type": "markdown", + "source": [ + "The model should draw a box around each object of interest in each image. If it isn't doing a good job of detecting objects, here are a few tips:\n", + "\n", + "1. Double-check your dataset to make sure there are no labeling errors or conflicting examples.\n", + "2. Increase the number of epochs used for training.\n", + "3. Use a larger model size (e.g. `yolo11l.pt`).\n", + "4. Add more images to the training dataset. See my [dataset video](https://www.youtube.com/watch?v=v0ssiOY6cfg) for tips on how to capture good training images and improve accuracy.\n", + "\n", + "You can also run the model on video files or other images images by uploading them to this notebook and using the above `!yolo detect predict` command, where `source` points to the location of the video file, image, or folder of images. The results will be saved in `runs/detect/predict`.\n", + "\n", + "Drawing boxes on images is great, but it isn't very useful in itself. It's also not very helpful to just run this models inside a Colab notebook: it's easier if we can just run it on a local computer. Continue to the next section to see how to download your newly trained model and run it on a local device." + ], + "metadata": { + "id": "EGiQw_gWbSBa" + } + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W7yrFRViVczX" + }, + "source": [ + "#7. Deploy Model" + ] + }, + { + "cell_type": "markdown", + "source": [ + "Now that your custom model has been trained, it's ready to be downloaded and deployed in an application! YOLO models can run on a wide variety of hardware, including PCs, embedded systems, and phones. Ultralytics makes it easy to convert the YOLO models to various formats (`tflite`, `onnx`, etc.) and deploy them in a variety of environments.\n", + "\n", + "This section shows how to download the model and provides links to instructions for deploying it on your PC and edge devices like the Raspberry Pi." + ], + "metadata": { + "id": "FEtybPmB_ERi" + } + }, + { + "cell_type": "markdown", + "source": [ + "## 7.1 Download YOLO Model\n", + "\n", + "First, zip and download the trained model by running the code blocks below.\n", + "\n", + "The code creates a folder named `my_model`, moves the model weights into it, and renames them from `best.pt` to `my_model.pt`. It also adds the training results in case you want to reference them later. It then zips the folder as `my_model.zip`." + ], + "metadata": { + "id": "IcoBAeHXa86W" + } + }, + { + "cell_type": "code", + "source": [ + "# Create \"my_model\" folder to store model weights and train results\n", + "!mkdir /content/my_model\n", + "!cp /content/runs/detect/train/weights/best.pt /content/my_model/my_model.pt\n", + "!cp -r /content/runs/detect/train /content/my_model\n", + "\n", + "# Zip into \"my_model.zip\"\n", + "%cd my_model\n", + "!zip /content/my_model.zip my_model.pt\n", + "!zip -r /content/my_model.zip train\n", + "%cd /content" + ], + "metadata": { + "id": "qcBdnOA9v85S" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "43ypwonynLVu" + }, + "outputs": [], + "source": [ + "# This takes forever for some reason, you can also just download the model from the sidebar\n", + "from google.colab import files\n", + "\n", + "files.download('/content/my_model.zip')" + ] + }, + { + "cell_type": "markdown", + "source": [ + "## 7.2 Deploy YOLO Model on Local Devices\n", + "\n", + "Next, we'll take our downloaded model and run it on a local device. This section provides instructions showing how to deploy YOLO models on various devices.\n", + "\n", + "I wrote a basic Python script, `yolo_detect.py`, that shows how to load a model, run inference on an image source, parse the inference results, and display boxes around each detected class in the image. The [script](https://github.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/blob/main/yolo_detect.py) gives an example of how to work with Ultralytics YOLO models in Python, and it can be used as a starting point for more advanced applications." + ], + "metadata": { + "id": "YL06c6pb_UqZ" + } + }, + { + "cell_type": "markdown", + "source": [ + "### 7.2.1 Deploy on PC (Windows, Linux, or macOS)\n", + "\n", + "The easiest way to run Ultralytics models on a PC is using Anaconda. Anaconda sets up a virtual Python environment and allows you to easily install Ultralytics and PyTorch. It automatically installs CUDA and cuDNN, which allows you to speed up model inference with your NVIDIA GPU.\n", + "\n", + "> **NOTE:** My YouTube video (link to be added) shows how to deploy your model on a PC. It walks through the following steps, so watch the video if you prefer having visual instructions.\n", + "\n", + "**1. Download and Install Anaconda**\n", + "\n", + "Go to the Anaconda download page at https://anaconda.com/download, click the “skip registration” button, and then download the package for your OS. When it's finished downloading, run the installer and click through the installation steps. You can use the default options for installation.\n", + "\n", + "**2. Set up virtual environment**\n", + "\n", + "Once it's installed, run Anaconda Prompt from the Start Bar. (If you're on macOS or Linux, just open a command terminal).\n", + "\n", + "Issue the following commands to create a new Python environment and activate it:\n", + "\n", + "```\n", + "conda create --name yolo-env1 python=3.12 -y\n", + "conda activate yolo-env1\n", + "```\n", + "\n", + "Install Ultralytics (which also installs import libraries like OpenCV-Python, Numpy, and PyTorch) by issuing the following command:\n", + "\n", + "```\n", + "pip install ultralytics\n", + "```\n", + "\n", + "If you have an NVIDIA GPU, you can install the GPU-enabled version of PyTorch by issuing the following command:\n", + "\n", + "```\n", + "pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124\n", + "```\n", + "\n", + "**3. Extract downloaded model**\n", + "Take the `my_model.zip` file you downloaded in Step 7.1 and unzip it to a folder on your PC. In the Anaconda Prompt terminal, move into the unzipped folder using:\n", + "\n", + "```\n", + "cd path/to/folder\n", + "```\n", + "\n", + "**4. Download and run yolo_detect.py**\n", + "\n", + "Download the `yolo_detect.py` script into the `my_model` folder using:\n", + "\n", + "```\n", + "curl -o yolo_detect.py https://raw.githubusercontent.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/refs/heads/main/yolo_detect.py\n", + "```\n", + "\n", + "Alright! We're ready to run the script. To run inference with a yolov8s model on a USB camera at 1280x720 resolution, issue:\n", + "\n", + "```\n", + "python yolo_detect.py --model my_model.pt --source usb0 --resolution 1280x720\n", + "```\n", + "\n", + "A window will appear showing a live feed from your webcam with boxes drawn around detected objects in each frame.\n", + "\n", + "You can also run the model on an video file, image, or folder of images. To see a full list of arguments for `yolo_detect.py`, issue `python yolo_detect.py --help` or see the [README file](https://github.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/blob/main/README.md).\n", + "\n", + "\n" + ], + "metadata": { + "id": "gzaJQ2sGEPhP" + } + }, + { + "cell_type": "markdown", + "source": [ + "### 7.2.2 Deploy on Raspberry Pi\n", + "\n", + "Keep an eye out for an article showing how to convert YOLO models to NCNN format and run them on the Raspberry Pi!" + ], + "metadata": { + "id": "GelkpRLPEYmJ" + } + }, + { + "cell_type": "markdown", + "source": [ + "# 8. Conclusion" + ], + "metadata": { + "id": "y8fOJ4g8Q5x0" + } + }, + { + "cell_type": "markdown", + "source": [ + "Congratulations! You've successfully trained and deployed a YOLO object detection model. 😀\n", + "\n", + "Next, you can extend your application beyond just drawing boxes and counting objects. Add functionality like logging the number of objects detected over time or taking a picture when certain objects are detected. Check out some example applications at our GitHub repository: https://github.com/EdjeElectronics/Train-and-Deploy-YOLO-Models\n", + "\n", + "Thanks for working through this notebook, and good luck with your projects!" + ], + "metadata": { + "id": "DEZGuG1-Peg5" + } + }, + { + "cell_type": "markdown", + "source": [ + "# Appendix: Common Errors" + ], + "metadata": { + "id": "fXeDs6SaQBRy" + } + }, + { + "cell_type": "markdown", + "source": [ + "If you run into any errors working through this notebook, please do the following:\n", + "\n", + "\n", + "- Double-check that the dataset files are set up in the correct folder structure\n", + "- Make sure there are no typos or errors in your labelmap file\n", + "- Google search the error to look for solutions\n", + "\n", + "If none of those help, please submit an [Issue](https://github.com/EdjeElectronics/Train-and-Deploy-YOLO-Models/issues) on the GitHub page. In this section, I will add resolutions to common errors as they come up." + ], + "metadata": { + "id": "Q19ENCHRQOCH" + } + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file