How to Train a YOLO-V5 Model on Custom Dataset.

abdulrehman ajmal — Sun, 30 Jan 2022 18:22:47 +0000

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

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.

How We Going to Train Our Model.

First we need a data set.
Second we will tain our model through google colabs.
Third we will run it on our own machine.

Environment Setup

Here is link to notebook: Google Colab
You need a google account to use Google Colab. You can either use notebook which is provided and recomended by YOLO to train or you can create your own notebook and follow along.

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

Installing the YOLOv5 Environment

To run object detection model in your own machine you need to setup YOLO in your own machine following code will clone the code base.

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

Preparing the Dataset.

First we need to download the data set in your machine and go straight to Roboflow and create new project. Roboflow is a great tool also recomended by the YOlO readme.

After creating new project then we need to upload dataset. Roboflow is very fast and it will upload images and lables.

Roboflow will convert data set into Yolo-V5 format and divide dataset in to train validate and test according to your personal requirments. After that you only need to export by download code.

Now we have prepared our dataset lets move to next step of training the model.

Train the Model

For this part we will now move on to google coloabs notebook after saving a copy of notebook you only need to run all the code cells which are already written please run one cell at a time with sequence after cell is done move to next.

In the third block provide your snippet of code given by roboflow.

After that run the rest of the code cells it will take a lot of time training the model but after running the last cell you will have a weights file in your goolge drive.

Run it on Your Machine

Now remeber we cloned Yolo-V5 go in that directory. Downolad Best.pt file from google drive and place this file in Yolo-5v directory along with a test image.

python detect.py --weights best.pt --source image.jpg

Run this command and you done.