Yolov5 Object Detection with Your Own Dataset

Object Detection is a task in computer vision that focuses on detecting objects in images, videos and real time. There are various object detection algorithms such as Faster R-CNN, Single Shot Detector (SSD), YOLO(You Only Look Once) etc.

YOLO was initially introduced as the first object detection model that combined bounding box prediction and object classification into a single end to end differentiable network. It was written in Darknet. Besides, YOLOv5 is the first of the YOLO models to be written in the PyTorch framework. It got released by Glenn Jocher(Founder & CEO of Utralytics).

GitHub — ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

YOLOv5 Variants

YOLOv5 has four different variants such as s, m, l and xl. Each one of them offering different detection accuracy and performance as shown below.

PreTrained Checkpoints

source: https://github.com/ultralytics/yolov5#pretrained-checkpoints

In this article, you will discover how you train your model and real-time prediction with YOLOv5 🚀

So, let us get start step by step!

Steps that we will follow in this tutorial

• Preparing the dataset
• Environment Setup: Install YOLOv5 dependencies
• Configure the YAML files
• Training the model
• Visualize the training data
• Detection with our model

Preparing the dataset

In this tutorial, we are going to use our own dataset. You can use any annotation tools you want for your images. We will use LabelImg for this project because it allows us to label in txt format which is necessary for YOLOv5.

After labeling, you can split your data as train and test. In train and test folder, you should also seperate images and labels by creating images and labels folders.

Environment Setup: Install YOLOv5 dependencies

In Local

• Create new environment in anaconda prompt

Conda new env 
conda create --name yolo
conda activate yolo

• PyTorch Installation

conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge

You can check any version which suitable for your computer here

Now, check if you are on GPU or not

python
>>>import torch
>>>torch.cuda.is_available()
True 
>>> print('Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))
Using torch 1.9.0 _CudaDeviceProperties(name='NVIDIA GeForce GTX 1050', major=6, minor=1, total_memory=2048MB, multi_processor_count=5)

If you get similar outputs like above, you are on the right track! 🌟

• Clone Repo and Install Dependencies

git clone --depth 1 https://github.com/ultralytics/yolov5  
pip install -U -r yolov5/requirements.txt

In Google Colaboratory

Here’s the link to my Notebook: Google Colab

You need a google account to use Google Colab. 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.

I suggest that you use Google Colab for training and local computer for detection.

Configure the YAML files

In yolov5/data folder, there is a data.yaml file that you should configure it according to your data. Give the path of images which is in train and test folders, number of class and names of them.

In yolov5/models folder, there are yolov5.yaml files. Which variant of the yolov5.yaml file you want to use, just change the number of class value in that file.

Here is the file structure:

File Structure

Training the model

Here, we are able to pass a number of arguments:

• img: define input image size
• batch: determine batch size
• epochs: define the number of training epochs
• data: specify the path to our yaml file
• cfg: specify our model configuration
• weights: specify the path to weights
• name: result names
• cache: cache images for faster training

python train.py --img-size <your_image_size> --batch 16 --epochs 300 --data './data/data.yaml' --cfg ./models/yolov5s.yaml --weights './yolov5s.pt'  --name <name_of_result_file>  --cache

Visualize the training data

You can visualize the training process with tensorboard.

[COLAB]
%load_ext tensorboard
%tensorboard --logdir <path of result folder>

Detection with our model

For prediction, you can use detect.py file of yolov5 with this command line;

python detect.py --weights <path_of_best.pt> --img-size <image_size>  --conf <confidence_threshold_value> --source <path_of_images or videos> (0 for webcam)

⭐ BONUS ⭐

If you want to customize the real-time prediction for your business, here is the simplified detection code for you.

I hope this article helps you while you work with YOLOv5 and you discover how simple and fun YOLOv5 is. See you soon until the next YOLO release 🚀