A novel attention-based enhancement framework for face mask detection in complicated scenarios^☆

2.1. Principle of YOLOv3 algorithm

YOLOv3 [20] is a masterpiece of one-stage detection algorithms, of which the main innovations include optimization of feature extraction network and implementation of multi-scale training, which achieves a trade-off between detection speed and accuracy. To be specific, YOLOv3 has adopted a feature extractor DarkNet-53, which borrows the idea of jump connection in residual network (ResNet) [26] and contains 53 layers with multiple residual blocks. Furthermore, DarkNet-53 has alleviated expensive computation and improved inference speed by reducing the number of both convolution layers in residual blocks and channels of feature maps. In addition, similar to the feature pyramid network (FPN) [27], up-sample operations have been employed to fuse deep and shallow features, which outputs feature maps with multi-scale resolutions including $13\times 13,26\times 26$ and 52 × 52 so as to achieve detection of targets in large, medium and small size, respectively.

As for the multi-scale training process in YOLOv3, firstly, the input images are divided into grids in size of $13\times 13,26\times 26$ and 52 × 52, each of which is corresponding to a multidimensional tensor. Secondly, for every vector three anchors are assigned, where $5+N$ parameters are included in each anchor. Specifically, the first five parameters depict the horizontal and vertical coordinate values of the center point, the width and height of anchor, and the confidence of bounding box, and $N$ is the number of categories. Finally, bounding box of each target is generated, and category probability is predicted as well.

It is noticeable that three parts are contained in the loss function of YOLOv3, including localization loss, confidence loss and classification loss. The last two are based on the cross-entropy loss function, and the first one $L_{loc}$ is defined as follows:

where $\sigma \left(\cdot \right)$ denotes the Sigmoid activation function and $\lambda$ is a weighting factor. For each divided grid, $\mathit{s}$ is the size and $\mathit{M}$ is the number of generated prior bounding boxes. ${\mathit{I}}_{ij}^{obj}$ takes value from $\{1,0\}$, which indicates whether the $j$th prior box of the $i$th grid is responsible for predicting the target; width and height ratios of the target to input image are denoted as $w_i^{\ast }$ and $h_i^{\ast }$, respectively. For a prior bounding box, $t_w$ and $t_h$ represent the width and height, $\sigma \left(t_x\right)$ and $\sigma \left(t_y\right)$ are horizontal and vertical coordinates of the center point. Notice that the above four parameters are prediction values, and accordingly, $t_w^{\ast }$, $t_h^{\ast }$, $\sigma \left(t_x^{\ast }\right)$, $\sigma \left(t_y^{\ast }\right)$ are the true values of corresponding parameter.

2.2. Existing mask detection methods

Recently, with the practical demands of pandemic control and the prosperous development of CV detection algorithms, plenty of methods have emerged for face mask detection tasks.

For one-stage frameworks, in [28] a YOLO-Mask detection algorithm has been proposed, which integrates attention mechanism into the feature extraction network to enhance the ability of representing salient features, which is designed to obtain enhanced features so as to achieve better detection performance and robustness. In [29], the authors have proposed an improved mask-wearing detection model based on YOLOv3, where spatial pyramidal pooling (SPP) structure is adopted to promote feature fusion at different levels, and experimental results have shown the strong robustness of the proposed model in complex cases. In [30], [31], fine-grained semantic features and detail information have been obtained through the strategies of broadening the feature pyramid structure and adding the feature fusion path, which is available in tiny face masks detection scenarios. In [32], a novel face mask detection Yolo (FMD-Yolo) model has been developed, which meticulously designs the input, backbone, neck, detection head and post-processing parts. Experimental results have demonstrated the strong generalization ability of FMD-Yolo in various kinds of complicated scenes with high detection precision.

For two-stage methods, a transfer learning technology based on fast RCNN has been designed in [33] to realize automatic location and type recognition of face masks. A multi-task learning framework based on mask RCNN and full convolution neural network has been proposed in [34], which compensates for the deficiency that single detection task cannot achieve pixel-wise face segmentation. In [35], a two-stage hybrid transfer learning approach has been developed, which completes the coarse extraction of face mask candidate regions based on RCNN and InceptionV2 at first, and then implements the mask wearing detection task by designing a broad learning system with a binary classification model. It has been proven that the hybrid model can present excellent performance even in complicated scenarios.

3.1. Overall structure of the proposed AI-Yolo

In Fig. 1, the framework of AI-Yolo is illustrated, whose backbone network adopts the DarkNet-53 in YOLOv3, which consists of one DBL (DarkNetConv2D BatchNormalization LeakyRelu) block and five Res-n (Residual block with n Res units) blocks. It is noticeable that via stacking a large number of 1 × 1 and 3 × 3 convolution layers, DarkNet-53 can effectively decline computational complexity and maintain feature extraction ability simultaneously.

According to Fig. 1, outputs of the last three stages in DarkNet-53 are fed into three SK modules independently, where dynamic selection of convolution kernels is completed. In particular, by adaptively adjusting weight distribution of convolution kernels on different feature maps, an effective information differentiation under the kernel attention mechanism can be realized in the SK module. Next, preliminary feature fusion is accomplished through up-sample, concatenation operations and DBL block, after which the outputs are sent into the SPP module that is committed to promoting sufficient fusion of multiple receptive fields and enhancing the expression ability of both global and local information. Finally, the FF module is designed to fully integrate the enhanced features so as to improve the robustness of model on detecting multi-scale objects, and more details of each module are presented in the following subsections.

3.2. Selective kernel (SK) attention mechanism

Since the face masks and other occlusions have different shapes and sizes in real-world scenes, how to obtain dynamic weight adjustment information so as to make the detector adapt to different complex environment deserves further attention. Consequently, SK attention mechanism [36] is employed in AI-Yolo, which is placed at the beginning of branches with different resolutions of inputs. The implementation of SK module is illustrated in Fig. 2, where “Split”, “Fuse” and “Select” operations are conducted, and the details are given as follows.

• •

  Split. An input feature map $X$ is split into two branches, where a group of 3 × 3 and 5 × 5 convolution kernels is set, respectively. Batch normalization is performed in both branches, and ReLU is adopted as the activation function, the obtained feature maps after transformation are denoted as $\tilde{U}$ and $\widehat{U}$, respectively.

• •
  Fuse. Global average pooling (GAP) and fully connected (FC) operations are carried out in sequence on $\tilde{U}$ and $\widehat{U}$ so as to achieve global information fusion of different kernel features, which can be depicted as:

  $$z=FC\left(GAP(\tilde{U}\oplus \widehat{U})\right)$$

  (2)

  where $z$ is the output of “Fuse” operation and $\oplus$ is the element-wise summation.
• •
  Select. Aiming at $z$, information with different spatial scales is adaptively selected as:

  $$V_c=\frac{1}{1+e^{z(B_c-A_c)}}\cdot \tilde{U}+\frac{1}{1+e^{z(A_c-B_c)}}\cdot \widehat{U}$$

  (3)

  where $V_c$ denotes the final feature map in $c$th channel, $A$ and $B$ represent the soft attention vectors of $\tilde{U}$ and $\widehat{U}$, respectively.

3.3. Spatial pyramid pooling (SPP)

In the face mask occlusion detection task, the excessive differences among target sizes make the accurate small object (i.e., the mask) detection an extremely challenging work. To handle this issue, spatial pyramid pooling module [37] is applied in AI-Yolo to realize the fusion of local and global information, which can effectively avoid the distortion of feature maps caused by down-sample and other operations. As is shown in Fig. 3, three groups of parallel max pooling operations constitute the main framework of the SPP module.

3.4. Feature fusion (FF) module

In general, the shallow feature maps have a small receptive field with rich detail information, while the deep feature maps have a large receptive field with abstract semantic information. Therefore, to integrate the multi-scale features will no doubt benefit the model performance, while it is noticeable that a simple concatenation may cause performance degradation in some complicated situations. In this regard, advantages of individual scale features are fully exploited in the designed FF module in AI-Yolo, of which the structure is displayed in Fig. 4.

As is shown, convolution, up- and down-sample operations are mainly adopted in the FF module to promote the effective multi-scale feature fusion without increasing too much complexity, which can be formulated as:

where $S_m$ and $P_n(m,n=1,2,3)$ represent the input and output feature maps, $C_o$ denotes a 1 × 1 convolution with stride equaling one, $U_i$ and $D_j(i,j=2,4)$ refer to the $i\times$ up-sample and $j\times$ down-sample operations, respectively. As a result, FF module can effectively avoid the performance degradation caused by the lack of semantic information or the missing of detailed features.

3.5. Loss function in AI-Yolo

It is worth mentioning that in addition to the above mentioned SK, SPP and FF modules, which aim at model optimization from the perspective of structural configuration, the loss function also plays an important role in determining the overall model performance. Considering that the original loss function of YOLOv3 is susceptible to the influence of target scales, in the proposed AI-Yolo, complete intersection over union (CIoU) loss [38] is utilized, which takes three important geometric factors of the bounding box into account, including overlap region, centroid distance and aspect ratio [39]. Calculation of CIoU loss $L_{CIoU}$ is given as:

where $A$ and $B$ are the center point coordinates of the prediction box and the ground-truth one, respectively, ${\rho }^2\left(\cdot \right)$ measures the Euclidean distance. $c$ denotes the diagonal length of the smallest outer rectangle of prediction box and the ground truth. $\alpha$ is a weighting factor and $\nu$ refers to similarity of the aspect ratio metric, which are calculated as:

where $w$, $w^{gt}$ and $h$, $h^{gt}$ are width and height of the prediction box and the ground-truth box, respectively. Degree of overlap between above two boxes is measured by $I_{IoU}$ as:

where $S$ and $S_{gt}$ represent the area of prediction box and the ground-truth box, respectively.

4.1. Experimental settings

To evaluate performance of proposed AI-Yolo, two public face mask detection datasets are used in this paper, which are denoted as WMD-1 and WMD-2 for convenience. The former is a combination of WIDER FACE [40] and MAFA [41] datasets, and the latter is accessible on website https://www.kaggle.com/datasets/andrewmvd/face-mask-detection, more details of the above two datasets are provided in Table 1.

To be specific, WMD-1 includes two categories, which are labeled with no_mask ($a_1$) and mask ($a_2$); whereas WMD-2 contains three classes of without_mask ($b_1$), with_mask ($b_2$) and mask_weared_incorrect ($b_3$). It is worth mentioning that accurate classification on both of the two datasets are arduous due to the great environment interference such as small size and dense distribution of objects, etc. In addition, there are 6120, 1530 samples in training and validation set of WMD-1, and 683, 170 samples in those of WMD-2, respectively.

Furthermore, during the training stage, the Adam optimizer is utilized to update the network weights, and 200 epochs are set for training iteration, where the parameters of backbone are frozen in the first 100 epochs. The initial learning rate is set to $1\times 10^{-3}$ with 0.1 decay factor for the last 100 epochs, and the initial batch size is defined as 8, which is adjusted to 4 with the decay of learning rate.

All experiments in this study are run under the deep learning framework PyTorch 1.2.0 with python 3.6. The hardware configuration is Nvidia GeForce GTX 1080Ti graphics card with 11 GB of video memory and Inter Xeon(R) E5-2682 processor on the Windows 10 operating system.

4.2. Evaluation metrics

In an object detection task, it is common to evaluate performance of algorithms from two aspects, which are classification accuracy and localization precision. In this paper, mean average precision ($mAP$) and F1 score are selected as evaluation metrics to measure the performance of the proposed AI-Yolo. Both of the two indicators are based on precision ($P$) and recall ($R$), where the former refers to that in all samples predicted to be a certain category, how many predictions are correct; whereas the latter refers to that in all samples belonging to a certain category, how many of them are correctly classified. Calculation of $mAP$ is given as:

where $N$ is the number of target categories. $P_i$ and $R_i$ stand for the precision and recall values of the $i$th category, respectively.

Another metric F1 score is the harmonic average of $P$ and $R$, which is regarded as a criterion for evaluation generalization ability and can be obtained as:

In addition, frames per second (FPS) is employed to measure the model inference speed, and the model complexity is evaluated by model size and the giga floating-point operations per second (GFLOPs) in this study.

4.3. Results and analysis

First of all, change curves of loss function of the proposed AI-Yolo and original YOLOv3 before and after smooth operation are displayed in Fig. 5, which can reflect the convergence tendency when training the model. As is shown, the proposed AI-Yolo can reach a lower stable loss value as compared with YOLOv3, which means that the prediction of the developed AI-Yolo is more accurate and the convergence ability is stronger than those of YOLOv3.

To further verify the superiority of AI-Yolo, some other state-of-the-art one- and two-stage algorithms are employed for comparisons on the same datasets, including Faster R-CNN(ResNet50) [13], Faster R-CNN(VGG16) [13], deconvolutional single shot detector (DSSD) [42], EfficientDet-D0 [43], YOLOv3 [20], YOLOv4 [21] and RetinaNet [44]. Comparison results in terms of $mAP$ and FPS on the two datasets are reported in Table 2, Table 3, respectively, where the IoU detection threshold of the target and prediction frame is set to 0.5.

According to Table 2, Table 3, on dataset WMD-1, the proposed AI-Yolo achieves the best $mAP$ value of 94.1%, and as compared with the sub-optimal model YOLOv4, the performance is improved by 2%. On dataset WMD-2, the proposed AI-Yolo also achieves the best $mAP$ of 90.7% that outperforms the sub-optimal model by nearly 5.1%. In terms of detection speed, our AI-Yolo ranks third out of the eight models, which is still a satisfactory results. Therefore, it can be concluded that the proposed AI-Yolo is a competitive object detection model with excellent overall performance, which may due to that the designed three modules effectively highlights the important features and reduces redundant information. Moreover, the improved loss function may also promote the stable target frame regression and accurate localization.

According to Fig. 6(a), the proposed AI-Yolo can also well balance the precision and recall, which keeps relatively high F1 scores on each category. To be specific, on both categories of the WMD-1 dataset, the F1 score obtained by AI-Yolo reaches over 90%, which performs the best among all of the eight algorithms. On dataset WMD-2, as compared with the advanced YOLOv4 model, our AI-Yolo obtains equivalent F1 score on category $b_2$, while on the other two categories, AI-Yolo yields better results, which show the superiority of the proposed model. In addition, the average F1 scores for all categories on the two datasets are illustrated in Fig. 6(b), and it can be clearly observed that the proposed AI-Yolo obtains the best mean F1 score of 90% and 79% on dataset WMD-1 and WMD-2, which improves the performance by 2% and 7% in comparison to the second-ranked YOLOv4 model, respectively.

It is noticeable that the three-category classification task in dataset WMD-2 is more difficult than that in WMD-1, while the improvement in terms of F1 score on WMD-2 is even more significant than that on WMD-1. The reasons for this phenomenon may lie in the applications of the designed SK, SPP and FF modules, where sufficient multi-scale enhanced feature fusion is achieved, which further promotes the model performance in dealing with interference of similar occlusions and handling the small object detection tasks. Consequently, the proposed AI-Yolo is proven to have strong robustness that can well adapt to complicated face mask detection scenarios.

Furthermore, the precision–recall (P–R) curves of AI-Yolo on WMD-1 and WMD-2 datasets are presented in Fig. 7, which provide a category-wise evaluation including the listed five classes in Table 1. It is worth mentioning that the P–R curve of a certain class can comprehensively reflect the model performance, where the area enclosed by the curve and coordinate axes represents the average precision ($AP$) of corresponding class. In Fig. 7, for a better view, aforementioned area has been filled with blue shadows, and the $AP$ values of each category are 92.73%, 95.56%, 86.15%, 94.37% and 91.66%, respectively. In general, the value of precision tends to decrease with the increasing threshold of recall, but it is observed that the developed AI-Yolo is able to maintain a high level of $P$ on each category as $R$ increases, which further demonstrates the practicality of the developed AI-Yolo model in complex and uncertain detection scenes with various kinds of mask occlusions. In Fig. 8, the P–R curves of other models are illustrated as well, and it is found that the P–R curve obtained by the proposed AI-Yolo encloses the largest area along with the coordinate axes on both datasets, which indicates that our AI-Yolo has strong robustness and satisfactory generalization ability.

In addition, the model complexity of the proposed AI-Yolo is also compared with other seven advanced models, which is characterized by model size and GFLOPs, and the results are displayed in Table 4.

As can be seen from Table 4, AI-Yolo has the model size of 80MB, which is 150MB fewer than that of the classical YOLOv3 model. Combining with the results in Table 2, Table 3, the $mAP$ of AI-Yolo is even 6% and 8% higher than that of YOLOv3 on both two datasets, which shows that the developed AI-Yolo can effectively balance the computational costs and detection accuracy. It may owe to the proposed AI-Yolo has effectively reduced the parameters by employing SK attention mechanism, which enriches the perceptual field information so as to enhance the representation of both local and global features. Although the proposed AI-Yolo is not the lightest model, the size is still 27.5%, 41.8%, 67.3%, 42.6% smaller than that of Faster R-CNN(ResNet50), Faster R-CNN(VGG16), YOLOv3, YOLOv4, RetinaNet, respectively, and simultaneously the proposed AI-Yolo yields better accuracy than above models. Nevertheless, the proposed AI-Yolo model consumes relatively large GLOPs (92.69 G), which is probably due to the large number of applied up- and down-sample operations in the FF module. Although sufficient multi-scale information fusion is realized, extra computational costs are introduced as well, thereby leading to large GLOPs. While according to Fig. 5, it is found that the proposed AI-Yolo model presents considerable convergence and the inference speed is even slightly faster than the original YOLOv3 model, which indicates an acceptable balance between the model complexity and detection accuracy of the proposed AI-Yolo model.

In Fig. 9, some visualization results of the developed AI-Yolo on the two datasets are displayed, where the first and second row present results on dataset WMD-1 and WMD-2, respectively. It should be highlighted that many mask detection scenes are complicated with the characteristics of high density, multiple disturbances and target confusion. Take Fig. 9(e) as an example, which contains various face targets with multiple occlusion types. Hence, the challenges of detection simultaneously include small object sizes, different face orientation and incomplete target display. According to the results, all targets in Fig. 9(e) are correctly identified with high confidence, which demonstrates that AI-Yolo can present excellent detection ability and strong robustness in extremely complicated situations.

4.4. Ablation study

In this subsection, ablation studies are carried out to further validate the effectiveness of each module in AI-Yolo, where the original YOLOv3 model is selected as the baseline. In particular, three variants of AI-Yolo are designed, which are YOLOv3-1 (YOLOv3$+$SK), YOLOv3-2 (YOLOv3$+$SK$+$SPP) and YOLOv3-3 (YOLOv3$+$SK$+$SPP$+$FF), respectively. It is noticeable that the distinction between YOLOv3-3 and AI-Yolo is the applied loss function (see Section 3.5). Other basic experimental settings remain unchanged, and results are reported in Table 5.

In Table 5, symbol $\uparrow$ denotes the performance improvement. As can be seen, after adding the designed SK, SPP and FF module successively, there is a steady increase of both $mAP$ and F1 score on the two datasets as compared to the baseline model, which implies that above three modules do contribute to enhancing the model performance. To be specific, on WMD-1 dataset, SPP module has significant influences on both $mAP$ and F1 score; whereas on WMD-2 database, SPP module has shown the strongest positive effects on F1 score, while $mAP$ is mostly promoted by SK module. With the designed FF module and CIoU loss function are further introduced, the proposed AI-Yolo model obtains the best results on the two datasets in terms of both evaluation metrics, which verifies that the strategies adopted in AI-Yolo have played important roles in extracting local and global key features, which obtains rich receptive field information and promotes enhanced feature fusion.

Based on above discussions, the developed AI-Yolo is proven to have considerable performance in terms of both accurate recognition and precise localization, which is competent to detect face masks under complex circumstances. In future, we aim to (1) optimize the model structure from the aspect of system science [45] to pursue the balance between the computational complexity and detection accuracy; (2) employ some heuristic optimization algorithms [46] to realize an automatic hyper-parameter adjustment; (3) apply the proposed AI-Yolo model in other detection scenarios to further verify the practicality such as industrial defect inspection [47].