yolo_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a new model on YOLOv5\n",
    "\n",
    "YOLOv5 already provides some pre-trained models to be used for object detection. Nonetheless, you can also create train a new model to identify some labels based on your necessity.\n",
    "\n",
    "In our case, as we intend to identify some labels that are not covered by those pre-trained models, a new model is essential to achieve our goal.\n",
    "\n",
    "This tutorial was inspired by the article [How to Create an End to End Object Detector using Yolov5?](https://towardsdatascience.com/how-to-create-an-end-to-end-object-detector-using-yolov5-35fbb1a02810). Some code in this notebook is based on the content from that article."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Environment\n",
    "\n",
    "### Installing PyTorch\n",
    "\n",
    "Under the hood, YOLOv5 uses PyTorch as its framework. To use PyTorch in its maximum perfomance (using the GPU for training and prediction), we need to follow some steps when installing it.\n",
    "\n",
    "> This method should work on Nvidia GPUs. If you have a different GPU or want to use the CPU, visit [PyTorch Download](https://pytorch.org/get-started/locally/) page for instructions.\n",
    "\n",
    "First, you need to install the latest driver for the Nvidia GPU and CUDA Driver 11.x.\n",
    "After that, I recommend installing the cuDNN driver for CUDA 11.x.\n",
    "\n",
    "In your Python environment (I recommend creating a new venv for this project) run\n",
    "the following command:\n",
    "\n",
    "```sh\n",
    "pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113\n",
    "```\n",
    "\n",
    "If everything goes right, PyTorch should be installed and we can follow up with YOLOv5 installation.\n",
    "\n",
    "\n",
    "### Installing YOLOv5\n",
    "\n",
    "Now that PyTorch is installed, we're ready to setup YOLOv5. For that you'll need to download YOLOv5 repository and install its dependencies.\n",
    "\n",
    "```sh\n",
    "git clone https://github.com/ultralytics/yolov5.git\n",
    "cd yolov5\n",
    "pip3 install -r requirements.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Handling Dataset\n",
    "\n",
    "### Image annotation\n",
    "\n",
    "When creating a new dataset for our model, we need to add some images with labels indicating what is in the image and the location of those objects. For that we need an annotation tool to make our job easier.\n",
    "\n",
    "In this project, we used [**labelImg**](https://github.com/tzutalin/labelImg), as it is free and is licensed under the MIT License. All images should be added to a single folder and that folder should be opened on **labelImg**. After you create annotations, a new `<filename>.txt` is going to be created for each image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Organize dataset into training data and validation data\n",
    "\n",
    "To train the model, YOLOv5 needs a dataset to be organized into two sections:\n",
    "- Training data (used to train the model and estimate new weights)\n",
    "- Validation data (used to test the accuracy of the model)\n",
    "\n",
    "Now that we have a folder containing all the images and annotations we want to use for training, the script below can be used to randomly organize the data into those two sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Inside YOLOv5 directory\n",
    "\n",
    "import glob, os\n",
    "import random\n",
    "\n",
    "# put your own path here\n",
    "dataset_path = 'dataset'\n",
    "\n",
    "# Percentage of images to be used for the validation set\n",
    "percentage_test = 20\n",
    "\n",
    "# Populate the folders\n",
    "%mkdir training\n",
    "%cd training\n",
    "%mkdir data\n",
    "%mkdir data/images\n",
    "%mkdir data/labels\n",
    "%mkdir data/images/train\n",
    "%mkdir data/images/valid\n",
    "%mkdir data/labels/train\n",
    "%mkdir data/labels/valid\n",
    "\n",
    "p = percentage_test/100\n",
    "\n",
    "for pathAndFilename in glob.iglob(os.path.join(f\"training\\{dataset_path}\", \"*.png\")):  \n",
    "    title, ext = os.path.splitext(os.path.basename(pathAndFilename))\n",
    "    if random.random() <= p :\n",
    "        print(f\"cp training/{dataset_path}/{title}.png training/data/images/valid\")\n",
    "        print(f\"cp training/{dataset_path}/{title}.txt training/data/labels/valid\")\n",
    "        os.system(f\"cp training/{dataset_path}/{title}.png training/data/images/valid\")\n",
    "        os.system(f\"cp training/{dataset_path}/{title}.txt training/data/labels/valid\")\n",
    "    else:\n",
    "        print(f\"cp training/{dataset_path}/{title}.png training/data/images/train\")\n",
    "        print(f\"cp training/{dataset_path}/{title}.txt training/data/labels/train\")\n",
    "        os.system(f\"cp training/{dataset_path}/{title}.png training/data/images/train\")\n",
    "        os.system(f\"cp training/{dataset_path}/{title}.txt training/data/labels/train\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Model\n",
    "\n",
    "You should take a look on [description](https://pytorch.org/hub/ultralytics_yolov5/) for different models to be used.\n",
    "\n",
    "For the model, copy the file for the model you want (e.g. `models/yolov5s.yaml` -> `training/yolov5s.yaml`). Don't forget to change `nc: 3` in that file to the number of classes the model is training to identify.\n",
    "\n",
    "You should also create a new file `training/dataset.yaml` containing:\n",
    "\n",
    "```yaml\n",
    "# path to train and validation datasets\n",
    "train: training/data/images/train/\n",
    "val: training/data/images/valid/\n",
    "\n",
    "# number of classes\n",
    "nc: 3\n",
    "\n",
    "# class names\n",
    "names: ['cuphead', 'player-bullet', 'boss-carrot''boss-radish']\n",
    "```\n",
    "\n",
    "To run the training:\n",
    "\n",
    "```sh\n",
    "python train.py --img 640 --batch 16 --epochs 300 --data training/dataset.yaml --cfg training/yolov5s.yaml --weights \" \" --workers 2\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "47cae9dd482ebb91bc5ca54e3737c993e5cdc0924b728541a9d0433d9717019c"
  },
  "kernelspec": {
   "display_name": "Python 3.10.4 64-bit (windows store)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}