Implementing YOLO on a custom dataset

In this article we will learn step by step implementation of YOLO v2 using keras on a custom data set and some common issues and their solutions

Understanding YOLO — You look only once

From CNN to Mask R-CNN and Yolo Part 1

From CNN to Mask R-CNN and Yolo Part 2

Object detection using Yo lov3

Orignal paper on Yolo

I will use Kangaroo dataset as my custom dataset

Required libraries

• Python 3.5 or higher
• Tensorflow :I have used CPU version
• OpenCV
• VC++ build 14

Download Yolo code implemented using keras

We will be using DarkFlow repo which can be downloaded from here : https://github.com/thtrieu/darkflow

Down the pre-trained weights

Weight can be downloaded from https://pjreddie.com/darknet/yolov2/

https://pjreddie.com/darknet/yolov2/

I created a new folder called bin and placed the downloaded weights there as shown below

Tip: Configuration(cfg) file and weights should match

If the configuration(cfg) file and weights are not matching then we get below error

Darkflow loads the weights by reading the .cfg file layer by layer, reading corresponding chunk of bytes from .weights. When there is a mismatch between the layers between weights and configuartion file. Configuration file may still have layers to read while the weights parser has already reach the end of .weights. This generates the over read assertion error.

Build the code

There are three options to build the code.

Option 1:

python setup.py build_ext --inplace

Option 2:pip install darkflow globally in dev mode

pip install -e .

Option 3: Installpip globally

pip install .

For me option 2 and 3 worked well.

We are now ready to process images or video file.

Processing an image file using Yolo

Importing required libraries

import cv2
from darkflow.net.build import TFNet
import matplotlib.pyplot as plt%config InlineBackend.figure_format = 'svg'

To define the model we can use the following options

1. model: configuration file (*.cfg) contains the details of the model

2. load: pre-trained weight file

3. batch: number of data to train per a batch

4. epoch: number of iterations to train

5. gpu: set 1.0 if you want to fully utilize the GPU hardware. If you need to use cpu then exclude this option

6. train: use this option when training a dataset

7. annotation: directory where the annotation files are stored

8. dataset: directory where the image files are stored

options = {
 'model': 'cfg/yolo.cfg',
 'load': 'bin/yolov2.weights',
 'threshold': 0.3
}
tfnet = TFNet(options)

Loading the image and identifying the objects in the image using Yolo

# read the color image and covert to RGBimg = cv2.imread(‘sample_img\sample_dog.jpg’, cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)# use YOLO to predict the image
result = tfnet.return_predict(img)
result

predicting the class, confidence and bounding box in the image

Display the class and the bounding box on the image

# pull out some info from the results
for i in range(0, len(result)):
tl = (result[i]['topleft']['x'], result[i]['topleft']['y'])
    br = (result[i]['bottomright']['x'], result[i]['bottomright']['y'])
    label = result[i]['label']# add the box and label and display it
img = cv2.rectangle(img, tl, br, (0, 255, 0), 7)
    img = cv2.putText(img, label, tl, cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 2)
    plt.imshow(img)
plt.show()

Training on a new dataset using Yolov2

I am using Kangaroo dataset as my custom dataset.

Annotations should be Pascal VOC(Visual Object classification) compliant. This format provides standardized image data sets for object class recognition.

Create a folder under darkflow folder and store the images and annotations. I have created a new folder new_data and created images and annots folder under it as shown below

Update the labels.txt with the custom class names

Darkflow loads object classes from a custom labels file using labels option. When the lables flag is not set, darkflow will load from labels.txt by default.

Updating the model in .cfg file

We update the model in the .cfg file. We need to update the number of classes and the filters in the last convolution layer in the ,cfg file.

No. of filters= (5 + no. of classes)*5

In our case since we have only one class kangaroo, we will have 5*(5+1)=30 filters as highlighted below

cgf file for the custom dataset

Tip 2: Match the classes in the labels.txt and the classes in .cfg file

Darkflow tries to load the labels from labels.txt by default. The below error occurs when the the model finds wrong number of labels in labels.txt when compared to the classes in the .cfg file

Define the model options

options = {"model": "cfg/yolo.cfg", 
           "load": "bin/yolov2.weights",
           "batch": 2,
           "epoch": 10,
           "train": True,
           "annotation": "new_data/annots/",
           "dataset": "new_data/images/"}

Build an instance of Darkflow

from darkflow.net.build import TFNet
tfnet = TFNet(options)

Training the model

tfnet.train()

Training the model using command prompt

We can also train the custom dataset using command line with the following command

python flow --model cfg/yolo.cfg --load bin/yolov2.weights --train --annotation new_data\annots --dataset new_data\images --epoch 1

Predicting the images using the weights trained on the new dataset

options = {
 ‘model’: ‘cfg/yolo.cfg’,
 ‘load’: 10,
 ‘threshold’: 0.3,
 “backup”:”ckpt/”,
 
}tfnet2 = TFNet(options)

Load the checkpoint

This loads pre-trained parameters from the checkpoint that we just specified in options.

tfnet2.load_from_ckpt()

Predicting on an image from the custom dataset

We will take one of the images that has not been used for training to predict the class, bounding box and confidence

original_img = cv2.imread(“new_data/images/00001.jpg”)
original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)
results = tfnet2.return_predict(original_img)
print(results)

one error that I encountered during my experimentation was

Assertion error with four byte size difference when converting.

The solution was to change the file loader.py in in utils folder weights_walker.__init__() method

 change self.offset = 16 to self.offset = 20 

References:

https://pjreddie.com/darknet/yolov2/

Github