Mask Detection using YOLOv5. Explanation of key concepts followed by... | by Sean Yap | Towards Data Science

I recently completed DeepLearningAI’s Convolutional Neural Network course taught by Andrew Ng on Coursera and a portion of the course dwelled into the field of Computer Vision. I was fascinated by the capability of computer vision and how it can be integrated into our daily lives which motivated me to dive deeper into computer vision and initiate projects to enforce my learning.
Before implementing a project, it's best to understand a few fundamental concepts on Object Detection and how it works together. Let's start by defining Object Detection:
Image classification is where an algorithm is applied to an image to predict the class of one object eg: Car. Object localization not only predicts the class of objects but also figures out the location of the object by drawing a bounding box around the object. Object detection involves both classification and localization and detects more than one object & more than one class even.
A standard classification task would involve an image running through a Convnet with multiple layers in which vector features are fed into a softmax unit for example that outputs the predicted class (Object categories that the algorithm is trying to detect i.e cars, trees, pedestrians). Object detection models such as YOLO works by splitting images into a grid of cells where each grid cell is responsible for predicting a bounding box if the centre of the bounding box lies in the cell. It will then output the predicted class, coordinates of the bounding box as shown below:
To learn more about convolutional neural networks, look up Convolutional Neural Networks, Explained by Mayank Mishra. He does a pretty good job at explaining how Convnet works
When the algorithm outputs bounding boxes localizing objects detected, how do you tell if the algorithm is working well? This is where Intersection over Union(IoU) comes into play. Generally, IoU is a measure of overlap between two bounding boxes: algorithm predicted bounding box and ground truth bounding box.
The formula for IoU is the size of the intersection divided by the size of the union of both bounding boxes. The threshold for IoU is around 0.5. IoUs with values ≥ 0.5 are deemed “correct” predictions.
With reference to the bounding box image below, cells labeled 1–33 all predict that the centre of the bounding box lies in their cell.
This will lead to the algorithm detecting the object multiple times instead of just once. This is where Non-max suppression comes in where it ensures the algorithm detects each object only once. As was previously stated, each cell outputs y = (Pc, bx, bγ, bw, bh, c) with Pc being the probability there is an object. What Non-max suppression does is it takes the bounding box with the highest Pc, discards any bounding boxes with Pc ≤ 0.6and “silences” the other bounding boxes with IoU ≥ 0.5.
What if multiple objects lie in the same grid cell? How would the bounding boxes be like? Here’s where the idea of anchor boxes can be used. Looking at the image below, notice how both the person and car midpoints lie within the grid cell. (For simplicity sake I split the image into a 3 by 3 grid)
The current cell output y = (Pc, bx, bγ, bw, bh, c), the cell can only choose one of the two objects detected. But with anchor boxes (usually predefined beforehand using a k-means analysis on the training dataset), the cost label y becomes (Pc, bx, bγ, bw, bh, c, Pc1, bx1, bγ1, bw1, bh1, c, Pc, ....) basically repeated depending on how man anchor boxes are there, with the first output encoded with anchor box 1 and second output encoded with anchor box 2 and so on. Each output unit detects for one object class; anchor box 1 is similar to the car so the output c will be for car, the next output c will be for person as it is encoded to anchor box 2.
Note that each output unit c representing the object class is influenced by high IoU with the object’s ground truth bounding box.
For this project, I will be using the YOLOv5 to train an object detection model. YOLO is an acronym for “You Only Look Once”. A popular architecture due to:
Speed (Base model — 45 frames per second, Fast model — 155 frames per second, 1000x faster than R-CNN, )
The architecture comprises only a single neural network (Can be optimized end to end directly on detection performance)
Able to learn the general representation of objects (Predictions are informed by the global context of image)
Open-source
To learn more about the model visit their repository: Ultralytics YOLOv5 Github repository.
I found this Face Mask Detection dataset on Kaggle, comprising of 853 images with 3 classes: With mask, without mask, and mask worn incorrectly. Each image comes with an XML file in PASCAL VOC format that contains annotations of its bounding boxes. Here’s an example:
<annotation>   <folder>images</folder>   <filename>maksssksksss4.png</filename>   <size>       <width>301</width>       <height>400</height>       <depth>3</depth>   </size>   <segmented>0</segmented>   <object>       <name>with_mask</name>       <pose>Unspecified</pose>       <truncated>0</truncated>       <occluded>0</occluded>       <difficult>0</difficult>       <bndbox>            <xmin>70</xmin>            <ymin>185</ymin>            <xmax>176</xmax>            <ymax>321</ymax>       </bndbox>   </object></annotation>
The key information needed for this project are:
<width>and <height> : Dimension of the image in pixelsThe entirety of <bndbox>: xmin, ymin indicates the location of the top left-hand corner of the bounding box while xmax, ymax indicates the bottom right-hand corner of the bounding box in pixels
<width>and <height> : Dimension of the image in pixels
The entirety of <bndbox>: xmin, ymin indicates the location of the top left-hand corner of the bounding box while xmax, ymax indicates the bottom right-hand corner of the bounding box in pixels
Before training the data with the model, the annotation data in PASCAL VOC XML format have to be changed to YOLO format, with one *.txt file per image with the following specifications: (Also illustrated with example image maksssksksss4.png below)
One row per object
Each row is class x_center y_center width height format.
Box coordinates must be in normalized xywh format (from 0–1). If your boxes are in pixels, divide x_center and width by image width, and y_center and height by image height.
Class numbers are zero-indexed (start from 0).
I wrote a function that extracts the required information from the XML files using the XML.etree library and calculates both x_centre and y_centre. As the annotation data is in image pixels, I normalized the final value to match the requirements.
The outcome would look like this:
It is also required for the directories to be formatted in a particular way as shown below where training and validation images and labels are separated into each independent folder
I wrote another simple function using the Pathlib library.
Do note that the entire formatting mentioned above can be done using Roboflow, a simple, code, and hassle-free alternative. It is my personal preference to manually do it by code
Before training the model, we will need to create a projectdata. yaml file, specifying where the training and validation images are and the number of labels as well as the label names of our training data. The file should be structured as shown:
# specify pathway which the val and training data is at# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]path: ../mask_detection/projectdatatrain: images/trainval: images/val# Classesnc: 3names: ['no mask', 'mask worn incorrectly', 'mask on']
To train the model with the custom dataset, I ran train.py in the local terminal with the following arguments:
img: input image size
batch: batch size
epochs: number of epochs
data: path to the projectdata.yaml file
cfg: model to choose among the preexisting in 📁models
weights: initial weights path, defaults to yolov5s.pt
cache: cache images for faster training
I chose the smallest and fastest model yolov5s. I also used the model’s pre-trained weights for transfer learning instead of training new weights all over again, which is time-consuming and impossible to train on a laptop due to high processing power requirements. I use a batch size of ‘1’ and trained the model for 10 epochs.
python mask_detection\yolov5\train.py --img 640 --batch 1 --epochs 10 --data projectdata.yaml--weights yolov5s.pt --cfg mask_detection\yolov5\models\yolov5s.yaml --cache
The command will output the following and start training if all the steps were done correctly. Keep an eye out for mAP@.5 to see how to model is performing. Once training begins, under the “run” folder, Yolov5 training pipeline inputs both ground truth and prediction results on test images as shown below.
Once the training is completed, the trained model will be saved in your “weights” folder/directory and the validation metrics will be logged onto Tensorboard. (I opted for the data be logged to wandb instead as it was recommended)
Popular challenges and competitions such as:
The PASCAL VOC Challenge (Everingham et al. 2010)
The COCO Object Detection Challenge (Lin et al. 2014)
The Open Images Challenge (Kuznetsova 2018)
All 3 challenges use mean average precision(mAP) as the main metric for evaluating object detectors. What exactly is mAP? First, let's go through some foundational concepts.
Confidence score: Probability that an anchor box contains an object predicted by the classifier
Intersection over Union (IoU): Area of intersection of bounding boxes divided by Area of union of the predicted bounding box
Precision: The number of True Positive(TP) divided by the sum of True Positive(TP) & False Positive(FP)
Recall: The number of True Positive divided by the sum of True Positive(TP) & False Negative (FP)
Both Confidence score and IoU are used to determine whether the predicted detection is TP or FP only if it satisfies the following conditions(Note: Any violation of the latter 2 conditions makes it a FP)
Confidence score > threshold (If < threshold, detection counts as a False Negative (FN))Predicted bounding box has a higher IoU than the thresholdPredicted class matches the class of ground truth
Confidence score > threshold (If < threshold, detection counts as a False Negative (FN))
Predicted bounding box has a higher IoU than the threshold
Predicted class matches the class of ground truth
As the confidence score increases, recall decreases monotonically while precision can go up and down, but for this project, all classes (except mask worn incorrectly) increases.
Although precision-recall metric can be used to evaluate the performance of object detector, it is not easy to compare among different detectors. To view the whole thing for comparison, average precision(AP) which is based on the precision-recall curve, comes into play. By definition, AP is finding the area under the precision-recall curve above. Mean average precision (mAP) is the mean of AP.
Note that there are variations in definitions and implementations of mAP.
In The PASCAL VOC Challenge, AP for one object class is calculated for an IoU threshold of 0.5. So the mAP is the averaged of all object classes.
For The COCO Object Detection challenge, the mAP is averaged over all object categories and 10 IoU thresholds.
Now that the model is trained, it's time for the fun part: running inference on either images or videos! Call detect.py in the local terminal to run inference with the following parameters: (view detect.py parser for full list of parameters)
weights: weights of the trained model
source: input file/folder to run inference on, 0 for webcam
iou-thres: IOU threshold for Non-max suppression, default value: 0.45
python yolov5\detect.py --source vid.mp4 --weights runs\train\exp\weights\best.pt --img 640 --iou-thres 0.5
As you can see, the model is able to detect the presence of face masks!
Far from perfect, the model still has room for improvement
The model is unable to detect incorrectly worn masks which is one of the classes. This is most probably due to the large imbalance of data. Only a small fraction of the data consists of incorrectly worn masks. A possible solution would be to apply data balancing so that the model will be better at identifying incorrectly worn masks.
A larger framework could have been used such as Yolov5x instead of Yolov5s to achieve higher mAP. But a possible drawback would be the extra time taken to train such a large model.
With that, I end my article here! This project has been really fun for me, I enjoyed learning new and interesting concepts especially when in the beginning, learning computer vision seemed like such a daunting task but I am glad to pull through it! Cheers!
LinkedIn Profile: Sean Yap
[1]Joseph Redmon, Santosh Divvala, Ross Girshick, Ali Farhadi, You Only Look Once: Unified, Real-Time Object Detection(2015)
[2]Andrew Ng, Convolutional Neural Network, DeepLearning.AI Coursera
[3]The PASCAL Visual Object Classes (VOC) Challenge, by Mark Everingham, Luc Van Gool, Christopher K. I. Williams, John Winn, and Andrew Zisserman