Face Mask Detection using YOLOv4 + Darknet

Rather than me talking about the popular "YOLOv4" - Object Detection Framework, I'll point you to the most reliable resource, the Github page of AlexyAB who is the creator of the framework. In the later blog posts, I will try to put forth "What I understood about YOLOv4".

https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects

But for now, lets get started with training and testing a custom object detector.

That is right! As the title suggests, we will see how to train the popular 'Face Mask Detection' model (Thanks to COVID-19).

Since my machine has no GPU, I am using "Google Colab". If you don't know what it is, i suggest a quick learning session would help you. In the simlest terms, it is a F-R-E-E Cloud based environments which supports CPU/GPU and TPU processing **CONDITIONS APPLY :)

Once you get the basics of Colab, make sure the Runtime is set to GPU and proceed.

Go to Runtime -> Change runtime type -> Select hardware Accelerator as GPU.

I come from a Database Administration background. The most important rule we learn is to ensure we have "BACKUPS". Training a custom object detector can take HOURS. What good is the training if it fails just before reaching the final step and we haven't thought of creating a checkpoint based backup strategy to resume the training from the point of failure.

Well, I got you back there. Just follow the step shown below and map your Google Drive to keep all the backups in place. Thank me later when your training is interrupted after 10 hours and you have to resume it from that point. It will save a lot of time and pull you up from the abyss of frustration. Lots of talk, lets get back to work.

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

It will ask for an auth code. Just click the URL and follow along.

Run the following lines

%cd /content/drive/My\ Drive

!pwd

You should be in "/content/drive/My Drive" by now. If not, ensure your drive is mapped correctly.

Lets gitclone Alexy's Darknet repo

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

Lets compile DARKNET. Boooooo. I know it sounds like a hacker business. Alexy has a good taste, doesn't he?

%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
!sed -i 's/CUDNN_HALF=0/CUDNN_HALF=1/' Makefile

print("Building Darknet....Stay Tuned" )

!make &> build_log.txt

Time to download the dataset

!mkdir "/content/drive/My Drive/darknet/MASK-data"

%cd "/content/drive/My Drive/darknet/MASK-data"
!pwd

!wget "https://www.dropbox.com/s/6gewe947ake1g95/kaggle_and_no-mask_dataset.zip?dl=1" -O MASK-data.zip
!unzip MASK-data.zip &> /dev/null

Note:

You can get datasets from various sources. For instance check out Open Images Dataset adn customize it accordingly. Or if you need to use a custom dataset which is not present on one of the popular image datasets, collect enough positive and negative samples and find teh bounding boxes then use it for training.

You can label your dataset using Yolo Mark.

In our case however, we have a labelled dataset.

Phew..... We got the data set in place. Lets prep it for training.

%cd ..

import random
import os
import subprocess
import sys

image_dir = "./MASK-data"
f_val = open("data/train.txt", 'w')
f_train = open("data/test.txt", 'w')

path, dirs, files = next(os.walk(image_dir))
data_size = len(files)

ind = 0
data_test_size = int(0.2 * data_size)
test_array = random.sample(range(data_size), k=data_test_size)

for f in os.listdir(image_dir):
    if(f.split(".")[-1] == "jpg"):
        ind += 1
       
        if ind in test_array:
            f_val.write(image_dir+'/'+f+'\n')
        else:
            f_train.write(image_dir+'/'+f+'\n')

f_train.close()
f_val.close()

It basically creates 2 text files containing  the location of training and test images.

YOLOv4 Configuration Files:

We need 3 files to train our model.

1. obj.names
2. obj.data
3. yolo4-obj.cfg (Configuration file)

obj.names: Contains the names of classes. One in each line.

obf.data: Contains the number of classes, path to train and test data and backup location.

yolo4-obj-cfg (make a copy from cfg/yolov4-custom.cfg and modify it per your need): It has information about width, height, filters, steps, max_batches, burnout etc. Read the documentation for more info.

Ensure obj  files are copied to "data" directory and cfg file is copied to "cfg" folder directory "Darknet" director.

 

Were almost done with teh setp. All we need is one more important ingredient. Before we see what it is, let me ask you a question.

Would you rather re-use the optimized algorithm and make modification to suit your needs or would you build one from scratch? Well, if you are from Computer Science background, you sure know the concept of re-using teh code. That is what we will do now.

Alexy has already trained YOLOv4 for MS COCO/Imagenet dataset and we will use the pre-trained weights to train our own model on top of it. Thanks to the open-source community.

!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137

And teh final step in the training process is

!./darknet detector train data/obj.data cfg/yolov4-obj.cfg yolov4.conv.137 -dont_show -map 2> train_log.txt

we provide the following:

• Path to obj.data
• Path to config file
• Path to the pre-trained weights

Flags used:

dont_show: which wont display the graphs. This is required for Colab since it does not have a display and it will crash if it tries to display, not if you run the notebook on your local system.
mAP: this will calculate mAP - mean average precision for the test data which we have specified using the data_test.txt file which contains 20% of our data.

Grab a cup of coffee or go out with your friend or binge watch your favourite show or sleep off happily. Training will take many hours depending on your configuration.

Performing Inference:

Helper function to dislay the soure and predited images:

import cv2
import matplotlib.pyplot as plt
%matplotlib inline

def display_output(imagePath):
    src = cv2.imread(imagePath,1)
    output = cv2.imread("prediction.jpg")

    plt.figure(figsize=[20,8])
    plt.subplot(121)
    plt.imshow(src[:,:,::-1])
    plt.title("Original Image")
    plt.subplot(122)
    plt.imshow(output[:,:,::-1])
    plt.title("Predictions")
    plt.show()

And lets test if our model works fine

!./darknet detector test data/obj.data cfg/yolov4-obj.cfg backup/yolov4-obj_best.weights trump.jpg -thresh .6 2> /dev/null

Does it work on a video file? Yes it does.

!./darknet detector demo data/obj.data cfg/yolov4-obj.cfg backup/yolov4-obj_last.weights Team_Mask_Force.mp4 -thresh .6 -out_filename  Team_Mask_Force.avi  -dont_show

Output Video:

India Cricket Team Mask Force