This post is a step-by-step guide of using YOLOv5 with your custom dataset and Pytorch. The first challenge is, should we name it YOLOv5? There are several discussions on this topic, but we will wait for the community to agree. For now, we will wait for the machine learning community for better naming.
YOLO refers to the paper published in May 2016. The YOLOv5 is on Pytorch and all the previous models used the darknet implementation. If you are interested in comparing YOLOv5 and the previous one, you can check this link. We will compare the result of our model on the next blog post. Our goal is to use the YOLO for logo detection.
This blog provides the step by step guideline to run YOLOv5. The source code is available at GitHub. You can clone it from here.
First of all, let us check the Cuda version installed in our machine by using
Cuda compilation tools, release 10.0, V10.0.130. It is important to note that we need Cuda version 10.
Create an environment:
Let's create an environment to install all the requirements. If you do not have a virtual environment installed, please follow this link.
We created the yolov5env:
"python3 -m venv yolov5env".
To activate this environment:
To install the requirements mentioned in the GitHub:
"pip install -U -r requirements.txt"
In case of any error for the requirements, you can install NumPy separately.
Preparing the custom dataset:
It is essential to follow the structure of the data. Create an image and label folder containing the train, validation, and test data.
The labels contain the text file of the class of the object and the location of the anchor box.
Two files need to be modified for your data, data.yaml, model/yolov5s_customize.yaml.
The data.yaml contains the training and validation path, number of the classes (nc) and the labels of classes. When you create your classes use the meaningful name which appears on your logo detection.
You can edit the structure of your network and also update the number of classes in your dataset at model/yolov5s_customize.yaml. you can select any other model by modifying the model and number of the classes (nc).
for training the Yolov5 you need to change
1- the image size --img 416
2- batch size --batch 16
3- number of epochs 200
4- custom data --data "data.yaml"
5- config file --cfg ./models/ yolov5s_customize.yaml
python train.py --img 416 --batch 16 --epochs 200 --data 'data.yaml' --cfg ./models/yolov5s_customize.yaml
we used cache-images to have a faster training.
The training results are saved in the results.txt and result.png, here you can see the result for out dataset using pre-trained weights for 200 epochs.
We need the the test data directory as the source. The output will be in the inference folder.
python detect.py --img-size 416 --source ../../data/data_png/data_code/images/test/ --weights weights/best.pt --conf 0.4 --save-txt
All the text files are saved in the inference/output.
Here are some of the results on our dataset for logo-detection:
On the next post, we are going to cover the difference between yolov5 and yolov4.