Training a Custom Object Detection Model With Yolo-V5

Prabhat Kumar Sahu
Jan 1 · 7 min read

Object Detection is a task in computer vision that focuses on detecting objects in images/videos.

There are various object detection algorithms out there like YOLO (You Only Look Once,) Single Shot Detector (SSD), Faster R-CNN, Histogram of Oriented Gradients (HOG), etc.

In this article, we are going to use Yolo-V5 to train our custom object detection model. YOLO is one of the most famous object detection models.

Version 1.0 of YOLO-V5 got released on 27 May 2020 by Glenn Jocher who is the founder & CEO of Utralytics. It’s written in PyTorch and it’s available in Github.

Prerequisites

It’s good to have a basic knowledge of deep learning computer vision. And how to work in a Google Colab environment.

Steps Covered in this Tutorial

To train our own custom object detector these are the steps to follow

  • Preparing the dataset
  • Environment Setup: Install YOLOv5 dependencies
  • Setup the data and the directories
  • Setup the YAML files for training
  • Training the model
  • Evaluate the model
  • Visualize the training data
  • Running inference on test images
  • Export the weight files for later use

Preparing the dataset

We are going to use the Aquarium Dataset which is available in Roboflow Public Object Detection datasets. You can check out more datasets that are there.

The aquarium dataset consists of 638 images. The images were already labeled by the Roboflow team. It has 7 classes such as fish, jellyfish, penguins, sharks, puffins, stingrays, and starfish, and most images contain multiple bounding boxes.

To download the dataset you need to create a roboflow account first. It’s very simple and easy.

Source: https://public.roboflow.com/object-detection/aquarium
Source: https://public.roboflow.com/object-detection/aquarium
Source: https://public.roboflow.com/object-detection/aquarium

If you want to use your own images you can use annotation tools such as LabelImg, CVAT or you can try any large-scale solutions like Scale or AWS Ground Truth.

When you are annotating yourself make sure to follow the best practices. Check this link out for more details.

After the dataset is prepared then we are all set to set up the environment and train the dataset.

Environment Setup

Here’s the link to my Notebook: Google Colab

You need a google account to use Google Colab. You can either use my notebook to train or you can create your own notebook and follow along.

In Google Colab, you will receive a free GPU for 12 hours. If you use a new notebook in Colab change the runtime session to GPU.

If you are planning to use my notebook then make sure to File → save a copy in your drive. Then you will be able to edit the code.

Installation of the dependencies

!git clone https://github.com/ultralytics/yolov5  # clone repo!pip install -U -r yolov5/requirements.txt  # install dependencies

Somehow the PyTorch version didn’t get compatible with the GPU so I installed another version of PyTorch by

#installing for google colab GPU use!pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

.

%cd /content/yolov5
!ls

We can import and take a look at our GPU Specification provided by Google Colab.

import torchfrom IPython.display import Image  # for displaying imagesfrom utils.google_utils import gdrive_download  # for downloading models/datasetsprint('Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

Here’s what I got

Using torch 1.6.0+cu101 _CudaDeviceProperties(name='Tesla T4', major=7, minor=5, total_memory=15079MB, multi_processor_count=40)

Google Colab comes preinstalled with Cuda and Torch and some other dependencies. If you are planning to train locally then you will have to setup Cuda and the dependencies on your own. I will surely make a tutorial about it later on.

Setup the data and the directories

After the environment set up is done. We can import the dataset into colab. As I am using the Roboflow dataset I will be downloading, if you plan to use your own you can import it using Google Drive.

# You need to sign up in roboflow to get the key and then you can use the dataset!curl -L “https://public.roboflow.com/ds/PUT YOUR OWN KEY HERE” > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

This will download the data, unzip and save it inside the yolov5 directory.

Project Folder structure

Setup the YAML files for training

To train a YOLO-V5 model, we need to have two YAML files.

The first YAML to specify:

  • where our training and validation data is
  • the number of classes that we want to detect
  • and the names corresponding to those classes

This YAML of ours looks like this:

train: ./train/images 
val: ./valid/images
nc: 7
names: ['fish', 'jellyfish', 'penguin', 'puffin', 'shark', 'starfish', 'stingray']

The second YAML is to specify the whole model configuration. You can change the network architecture in this step if you want but we will go with the default one.

The YAML which we term custom_yolov5s.yaml:

We can put the YAML file anywhere we want because we can reference the file path later on. But it’s a good idea to put it inside the YoloV5 directory.

Training the model

After the configuration is done we can begin our training.

There are multiple hyperparameters that we can specify which are:

  • img: define input image size
  • batch: determine batch size
  • epochs: define the number of training epochs.
  • data: set the path to our YAML file
  • cfg: specify our model configuration
  • weights: specify a custom path to weights
  • name: result names
  • nosave: only save the final checkpoint
  • cache: cache images for faster training

We need to specify the path of both YAML files which we created above.


%cd /content/yolov5/
!python train.py --img 416 --batch 80 --epochs 100 --data './data.yaml' --cfg ./models/custom_yolov5s.yaml --weights ''
Training

With 100 epochs the training got completed within 35 minutes.

Evaluate the model

Training losses and performance metrics are saved to Tensorboard and also to a logfile defined above with the — name flag when we train. In our case, we named this yolov5s_results. (If given no name, it defaults to results.txt.) The results file is plotted as a png after training completes.

Partially completed results.txt files can be plotted with from utils.utils import plot_results; plot_results().

# Start tensorboard
# Launch after you have started training to all the graphs needed for inspection
# logs save in the folder "runs"
%load_ext tensorboard
%tensorboard --logdir /content/yolov5/runs

Visualize the training data

After training starts, view train*.jpg images to see training images, labels, and augmentation effects. We can visualize both Ground Truth Training data, as well as Ground Truth, Augmented data.

# first, display our ground truth data
# The ground truth [Train data] is available in jpg file at location /content/yolov5/runs/train/exp2/test_batch0_labels.jpg
print("GROUND TRUTH TRAINING DATA:")
Image(filename='/content/yolov5/runs/train/exp2/test_batch0_labels.jpg', width=900)
# print out an augmented training example
# Below is the augmented training data.
# NOTE: The dataset already contains the augmented data with annotations, so that you dont have to do it.
print("GROUND TRUTH AUGMENTED TRAINING DATA:")
Image(filename='/content/yolov5/runs/train/exp2/train_batch0.jpg', width=900)

Running inference

Using the final trained weight which got saved after training we can run our inference

Weights folder

To run the model inference we can use the following command.

# use the best weights!
# Final weights will be by-default stored at /content/yolov5/runs/train/exp2/weights/best.pt
%cd /content/yolov5/!python detect.py --weights
/content/yolov5/runs/train/exp2/weights/best.pt --img 416 --conf 0.4 --source ./test/images
  • — source: input images directory or single image path or video path
  • — weights: trained model path
  • — conf: confidence threshold

This will process the input and store the output in our inference directory.

Here are some output images:

Export the weights for later use

Now that we have successfully trained our custom model. We can download the weight files and save them in our local directory or in Google Drive.

To do so we import a Google Drive module and send them out

from google.colab import drive
drive.mount('/content/gdrive')
%cp /content/yolov5/runs/train/exp2/weights/best.pt /content/gdrive/My\ Drive

Conclusion

I hope you were able to follow along and was able to train successfully.

I have uploaded the notebook, config files, and weight to my Github repository. You can check it out here.

If you have any questions, recommendations, or critiques, I can be reached via Twitter or via my mail. Feel free to reach out to me.

References:

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data…

Sign up for Analytics Vidhya News Bytes

By Analytics Vidhya

Latest news from Analytics Vidhya on our Hackathons and some of our best articles! Take a look.

By signing up, you will create a Medium account if you don’t already have one. Review our Privacy Policy for more information about our privacy practices.

Check your inbox
Medium sent you an email at to complete your subscription.

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data Science professionals. We are building the next-gen data science ecosystem https://www.analyticsvidhya.com

Prabhat Kumar Sahu

Written by

{ ML Engineer 💻 }

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data Science professionals. We are building the next-gen data science ecosystem https://www.analyticsvidhya.com

Medium is an open platform where 170 million readers come to find insightful and dynamic thinking. Here, expert and undiscovered voices alike dive into the heart of any topic and bring new ideas to the surface. Learn more

Follow the writers, publications, and topics that matter to you, and you’ll see them on your homepage and in your inbox. Explore

If you have a story to tell, knowledge to share, or a perspective to offer — welcome home. It’s easy and free to post your thinking on any topic. Write on Medium

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store