Face Mask Detection using YOLOv5. COVID-19: Guide to build a face mask... | by Àlex Escolà Nixon | Towards Data Science

In this post, we’ll be going through a step-by-step guide on how to train a YOLOv5 model to detect whether people are wearing a mask or not on a video stream.
We’ll start by going through some basic concepts behind object detection models and motivate the use of YOLOv5 for this problem.
From there, we’ll review the dataset we’ll be using to train the model, and see how it can be adapted so that it conforms with the darknet format.
I’ll then show you how you can train a YOLOv5 model using the downloaded dataset and run inference on images or video files.
Find a notebook with the code used for this section here.
This project was done alongside with Fran Pérez, also a contributor of this community. I’d recommend you to read some of his excellent posts!
Also, you might find it helpful to start with a more general overview of object detection and its main algorithms. If that is the case, I’d recommend reading this blog post for a good starting point.
“YOLO”, refering to “You Only Look Once”, is a family of object detection models introduced by Joseph Redmon with a 2016 publication “You Only Look Once: Unified, Real-Time Object Detection”.
Since then, several newer versions have been released, of which, the first three were released by Joseph Redmon. On June 29, Glenn Jocher released the latest version YOLOv5, claiming significant improvements with respect to its predecessor.
The most interesting improvement, is its “blazingly fast inference”. As posted in this article by Roboflow, running in a Tesla P100, YOLOv5 achieves inference times of up to 0.007 seconds per image, meaning 140 FPS!
Using an object detection model such as YOLOv5 is most likely the simplest and most reasonable approach to this problem. This is because we’re limiting the computer vision pipeline to a single step, since object detectors are trained to detect a:
Bounding box and a
Corresponding label
This is precisely what we’re trying to achieve for this problem. In our case, the bounding boxes will be the detected faces, and the corresponding labels will indicate whether the person is wearing a mask or not.
Alternatively, if we wanted to build our own deep learning model, it would be more complex, since it would have to be 2-fold: we’d need a model to detect faces in an image, and a second model to detect the presence or absence of face mask in the found bounding boxes.
A drawback of doing so, apart from the complexity, is that the inference time would be much slower, especially in images with many faces.
We now know which model we can use for this problem. The next natural, and probably most important, aspect is... what about the data??
Luckily, there is a publicly available dataset in Kaggle named Face Mask Detection, which will make our life way easier.
The dataset contains 853 images and their corresponding annotation files, indicating whether a person is wearing a mask correctly, incorrectly or not wearing it.
Bellow is a sample from the dataset:
In this case, we’ll simplify the above to detect if a person is wearing the mask or not (we’ll see how in the Roboflow section).
We now know everything we need to get started, so its time to get hands-on!
The first thing we need to do is clone the repository from ultralytics/yolov5, and install all required dependencies:
!git clone https://github.com/ultralytics/yolov5 # clone repo!pip install -U -r yolov5/requirements.txt # install dependencies
The main files that we’ll need from the repository are structured as follows:
yolov5           # project root├── models       # yolov5 models├── train.py     # training script└── detect.py    # inference script
The 📁models folder contains several .yml files with the different proposed models. In it we can find 4 different models, ordered from smaller to larger (in terms of the amount of parameters): yolov5-s, yolov5-m, yolov5-l and yolov5-x. For a detailed comparison, see here.
train.py and detect.py will be the scripts that we’ll be calling to train the model and predict on new images/videos respectively.
In order to train the model, a necessary step will be to change the format of the .xml annotation files so that they conform with the darknet format. In the linked github thread, we’ll see that each image has to have a .txt file associated with it, with rows with the format:
<object-class> <x> <y> <width> <height>
Each line will represent the annotation for each object in the image, where <x> <y> are the coordinates of the centre of the bounding box, and <width> <height> the respective width and height.
For example an img1.jpg must have an associated img1.txt containing:
1 0.427234 0.123172 0.191749 0.1712390 0.183523 0.431238 0.241231 0.1741211 0.542341 0.321253 0.191289 0.219217
The good news, is that this step is made really simple thanks to Roboflow. Roboflow enables to easily change between annotation formats, as well as as to augment our image data and split it into training and validation sets, which will be extremely handy!
This can be done in 5 simple steps:
Upload the images and annotations
Choose the train, validation and test proportions you want (train and validation will be enough)
Add an augmentation step choosing from among the existing filters, for instance blur, brightness, rotation etc.
And finally, generate the new images and export to YOLO Darknet format
We should now have a separate folder for each split, train and validation (and test if included), where each of these should contain the .jpg augmented images, the corresponding .txt annotation files and a ._darknet.labels file with the labels in their corresponding order:
mask_weared_incorrect # label 0with_mask             # label 1without_mask          # label 2
This will depend on how you want your model to behave, but as mentioned above, I’ve decided to simplify the problem by limiting the labels to mask or no mask. Also, the label mask_weared_incorrect appears very few times in the dataset, so the model would most likely do a poor job at classifying it.
This simple script will do:
Now, the numerical labels in the annotation files, will map to:
without_mask          # label 0with_mask             # label 1
The dataset is slightly unbalanced, having more with_mask labels, hence something we can do is augment with images of people not wearing a mask. This will also help improve significantly the performance of the model, since the dataset we’ll be using is quite small.
Find a notebook with the code used for this section here.
One way we could do this is by downloading images from COCO dataset and then annotating them ourselves. The COCO dataset has an official API, pycocotools, which we can use to download images with person labels:
The above will give us a list of dictionaries with the details and url of the images, which we can then use to download locally with:
The next step will be to obtain the face bounding boxes for these images. For this a simple way is to use a pre-trained model to detect the faces, and label them in the appropriate format.
For that you need to first install pytorch and then facenet-pytorch :
conda install pytorch torchvision torchaudio cpuonly -c pytorchpip install facenet-pytorch
I also use cv2 to check that the annotations are correct. To have the annotations in darknet format, we’ll need to convert the detections made by facenet-pytorch. We can do this with the function bellow:
And you can then obtain the annotations for the downloaded images with the following script:
Which also prints the detected annotations, to verify that they are correct. Here are some examples:
Before training the model, we will need to create a data.yml, specifying where the training and validation images are and the amount of labels as well as the label names of out training data. The file should be structured as:
train: train/imagesval: valid/imagesnc: 2names: ['bad', 'good']
To train the model we have to run train.py, which takes the following arguments:
img: input image size
batch: batch size
epochs: number of epochs
data: path to the data.yml file
cfg: model to choose among the preexisting in 📁models
weights: initial weights path, defaults to yolov5s.pt
name: renames output folder
device: Wheter to train on cpu or gpu.
In my case, I trained the model with a GPU, but otherwise you’ll have to add to the arguments --device cpu to run locally on a CPU. For the rest, you just need to take care of the specified routes:
python ~/github/yolov5/train.py --img 416 --batch 16 --epochs 150 --data data.yaml --cfg yolov5s.yaml --weights ‘’ --name robo4_epoch150_s --adam
Once the training completes, you can see how your model has performed by checking the generated logs, as well as the loss plots saved as a .png file.
Training losses and performance metrics are logged to Tensorboard and also to a runs/exp0/results.txt logfile, which is plotted as results.png after the training completes.
Here we see yolov5s.pt trained to 150 epochs (blue), and 300 epochs (orange):
As soon as the first epoch is complete we will have a mosaic of images showing both the ground truth and prediction results on test images, which will look like:
There’s a detailed tutorial on a colab notebook by ultralytics where you’ll find the above explained in more detail and available for you to run yourself!
Now that we’ve trained the model on a dataset of face mask images, it is ready to run inference both on single images and video streams!
To run inference, we have to call detect.py, and adjust the following parameters:
weights: weights of the trained model
source: input file/folder to run inference on, 0 for webcam
output: directory to save results
iou-thres: IOU threshold for NMS, defaults to 0.45
conf-thres: object confidence threshold
The two last you’ll probably have to tweak a little bit based on the generated results:
iou-thres defines a threshold on the Intersection Over Union metric obtained for a given object. It defaults to 0.45 , generally a IOU higher than 0.5 is considered a good prediction. See this great post by Adrian Rosebrock for more details on this metric.
conf-thres is a measure that determines how confident the model is that a given bounding box contains an object. It equals to the IoU times a factor representing the probability of an object. This metric mostly prevents the model from predicting backgrounds. The higher the threshold, the lesser false positives. However, if we set it too high the model will miss many correct detections on which it has a lower confidence.
After a few trial and errors, the following setting seemed to yield the best results for this problem:
python ~/github/yolov5/detect.py --weights raw/roboflow4/weights/last_robo4_epoch300_s.pt --source ./office_covid.mp4 --output raw/roboflow4/inference/output --iou-thres 0.3 --conf-thres 0.6
To test the model, I used a video taken with my phone. These are the results:
Here you can see that the model is capable of detecting face masks accurately.
More data training diversity. The model struggles to detect face masks under certain conditions, for instance, it tends to confuse a long beard with a mask. This can be mitigated by adding more diversity in the training images and perhaps running additional data augmentation techniques.
Comparing the different YOLOv5 models to check which performs best based on the metrics and also taking the training time into account. The bigger model YOLOv5x will take significantly longer to train.
Thanks a lot for taking the time to read this post and I hope you enjoyed it :)