Object Detection Model on Custom Dataset Using YOLO Pre-Trained Weights — The Training

Object Detection:

Data Preparation

Download LabelImg by running the following snippet in the terminal:

git clone https://github.com/tzutalin/labelImg.git

Now, open the directory and run “python3 labelImage.py”. It will open a window like this.

Click on “Open Dir” And choose your directory containing the data.

Now Click on Create on RectBox and make a box and label the name you want …

After labeling the data, we have to make a model (The Brain), that will make the boxes in the correct place, where the objects are.

Upload the folder containing the labels to your drive.

So here is a link to train the model: “https://github.com/nisargbhatt09/Object-Detection-Custom-Data/blob/main/Yolo_Training.ipynb”

Let’s Code

I would prefer that you use Google Colab if you don’t have a GPU on your device.

So, Firstly let’s check whether the GPU is enabled or not.

!nvidia-smi

If you get something like this, GPU is enabled.

+------------------------------------------------------------------+ | NVIDIA-SMI 460.56  Driver Version: 460.32.03  CUDA Version: 11.2 | |---------------------------+------------------+-------------------+ | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC | | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. | |                               |                      |               MIG M. | |===============================+======================+======================| |   0  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 | | N/A   37C    P0    26W / 250W |      0MiB / 16280MiB |      0%      Default | |                               |                      |                  N/A | +-------------------------------+----------------------+----------------------+                                                                                 +-----------------------------------------------------------------------------+ | Processes:                                                                  | |  GPU   GI   CI        PID   Type   Process name                  GPU Memory | |        ID   ID                                                   Usage      | |=============================================================================| |  No running processes found                                                 | +-----------------------------------------------------------------------------+

Access the Google Drive

from google.colab import drive
drive.mount('/content/gdrive')

List of items on your drive

!ln -s /content/gdrive/My\ Drive/ /mydrive
!ls /mydrive

Darknet: It is an open-source neural network framework. We will use it to train the pre-trained weights later.

Clone the darknet:

!git clone https://github.com/AlexeyAB/darknet

Set the OpenCV with the GPU

# change makefile to have GPU and OPENCV enable
%cd darknet
!sed -i 's/OPENCV=0/OPENCV=1/' Makefile
!sed -i 's/GPU=0/GPU=1/' Makefile
!sed -i 's/CUDNN=0/CUDNN=1/' Makefile
!make

Now let’s change the configuration file “yolov3-tiny.cfg” from cfg directory from the darknet. This is for one class, you may have more than one, just change 1 with the number of classes you want.

!cp cfg/yolov3-tiny.cfg cfg/yolov3_training.cfg
!sed -i 's/batch=1/batch=64/' cfg/yolov3_training.cfg
!sed -i 's/subdivisions=1/subdivisions=16/' cfg/yolov3_training.cfg
!sed -i 's/max_batches = 500200/max_batches = 4000/' cfg/yolov3_training.cfg
!sed -i '610 s@classes=80@classes=1@' cfg/yolov3_training.cfg
!sed -i '696 s@classes=80@classes=1@' cfg/yolov3_training.cfg
!sed -i '783 s@classes=80@classes=1@' cfg/yolov3_training.cfg
!sed -i '603 s@filters=255@filters=18@' cfg/yolov3_training.cfg
!sed -i '689 s@filters=255@filters=18@' cfg/yolov3_training.cfg
!sed -i '776 s@filters=255@filters=18@' cfg/yolov3_training.cfg

Create a folder on your drive to save the weights.

# Create folder on google drive so that we can save there the weights
!mkdir "/mydrive/yolov3_w1"

We will make a file called “obj.data” in data directory that will keep all the classes names, here I have only used one class, so bear with me 😌️

!echo "Class_Name" > data/obj.names
!echo -e 'classes= 1\ntrain  = data/train.txt\nvalid  = data/test.txt\nnames = data/obj.names\nbackup = /mydrive/yolov3_w1' > data/obj.data
!mkdir data/obj

Now we need pre-trained weights, you can use darknet53.conv.74. This weight is going to be retrained on given data.

!wget https://pjreddie.com/media/files/darknet53.conv.74

Extract the directory from the drive, unzip the directory to data/obj.

!unzip /mydrive/data.zip -d data/obj

We’re going to convert the class index on the .txt files. As we’re working with only one class, it’s supposed to be class 0. (The indentation is not correct, please refer to the code link.)

import glob
import os
import re
txt_file_paths = glob.glob(r"data/obj/new_dataset/*.txt")
for i, file_path in enumerate(txt_file_paths):
# get image size
with open(file_path, "r") as f_o:
lines = f_o.readlines()
text_converted = []
for line in lines:
print(line)
numbers = re.findall("[0-9.]+", line)
print(numbers)
if numbers:
# Define coordinates
text = "{} {} {} {} {}".format(0, numbers[1], numbers[2], numbers[3], numbers[4])
text_converted.append(text)
print(i, file_path)
print(text)
# Write file
with open(file_path, 'w') as fp:
for item in text_converted:
fp.writelines("%s\n" % item)

Gather the jpg files

import glob
images_list = glob.glob("data/obj/new_dataset/*.jpg")
## Perform Normalization!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
### I have not done it........
print(images_list)

Create the train.txt file in which there will be names of jpg files.

file = open("data/train.txt", "w")
file.write("\n".join(images_list))
file.close()

Start the Training, wait until the model converges.

Wait for eternity…

!./darknet detector train data/obj.data cfg/yolov3_training.cfg darknet53.conv.74 -dont_show

After this, the weights will be stored at “yolov3_w1” in your drive.

You Got It !!!

Thanks for reading, please do clap and share this article.