Research on Improved Mask Detection Method Based on YOLOv5 Algorithm

doi:10.3778/j.issn.1002-8331.2304-0243

Abstract

Abstract: The existing mask detection models cannot balance detection accuracy and detection speed, and have a large number of parameters. In order to solve these problems, an improved YOLOv5 based mask detection algorithm is proposed. It mainly includes the following four improvements：Firstly, replace the C3 module in the YOLOv5s backbone network with the lightweight network GhostNetV2 to reduce the number of parameters; Secondly, replace the last C3 module of the YOLOv5s backbone extraction network and the C3 module of the last layer of Neck with a Swin Transformer structure to obtain more complete feature information and improve detection performance; Thirdly, introduce CBAM attention mechanism to better focus on key information, thereby improving detection efficiency and accuracy; Fourth, the loss function replaces GIoU with EIoU to improve positioning accuracy and speed up convergence. The experimental results on the AIZOO dataset show that the mAP value of the proposed impoved algorithm reaches 96.2%, Params is reduced to 6.6×106, and FPS is as high as 136. There is also a good improvement on the validation dataset, and compared to other algorithms, the performance is better. The proposed improved algorithm is more suitable for mask detection.

Key words: mask detection, GhostNetV2, Swin-Transformer, attention mechanism

摘要： 现有的口罩检测模型无法平衡检测精度和检测速度，参数量较大，为了解决这些问题，提出了一种基于改进YOLOv5的口罩检测算法。该算法主要包括以下四点改进：使用轻量化网络GhostNetV2替换YOLOv5s主干网络中的C3模块，以降低参数量；将YOLOv5s主干提取网络的最后一个C3模块和Neck最后一层的C3模块替换为Swin-Transformer结构，来获取更为完整的特征信息，提高检测效果；引入CBAM注意力机制以更好地聚焦于关键信息，从而提高检测效率和检测精度；损失函数使用EIoU替换掉GIoU来提高定位准确度，加快收敛速度。在AIZOO数据集上的实验结果表明，所提出的改进算法的mAP值达到了96.2%，Params降低为6.6×106，FPS高达136，验证数据集上的性能也有很好的提升，相比其他算法，改进算法的性能更优，更适用于口罩检测。

关键词: 口罩检测, GhostNetV2, Swin-Transformer, 注意力机制

DUAN Bichong, MA Mingtao. Research on Improved Mask Detection Method Based on YOLOv5 Algorithm[J]. Computer Engineering and Applications, 2023, 59(16): 223-231.

段必冲, 马明涛. 基于改进YOLOv5算法的口罩检测研究[J]. 计算机工程与应用, 2023, 59(16): 223-231.

References

[1] REDMON J，FARHADI A.YOLO9000：better，faster，stronger[C]//Proceedings of the 2017 IEEE Conference on Computer Vision and Pattern Recognition，2017：6517-6525.
[2] REDMON J，FARHADI A.YOLOv3：an incremental improvement[J].arXiv：1804.02767，2018.
[3] BOCHKOVSKIY A，WANG C Y，LIAO H Y M.YOLOv4：optimal speed and accuracy of object detection[J].arXiv：2004.10934，2020.
[4] 杨秋妹，陈淼彬，黄一桂，等.基于改进YOLOv5n的猪只盘点算法[J].农业机械学报，2023，54（1）：251-262.
YANG Q M，CHEN M B，HUANG Y G，et al.Pig counting algorithm based on improved YOLOv5n[J].Journal of Agricultural Machinery，2023，54（1）：251-262.
[5] 张烈平，李智浩，唐玉良.基于迁移学习的轻量化YOLOv2口罩佩戴检测方法[J].电子测量技术，2022，45（10）：112-117.
ZHANG L P，LI Z H，TANG Y L.Detection method of lightweight YOLOv2 mask wearing based on transfer learning[J].Electronic Measurement Technology，2022，45（10）：112-117.
[6] 徐泽昊，刘川.基于改进YOLOv3的口罩佩戴检测算法[J].信息记录材料，2022，23（7）：158-161.
XU Z H，LIU C.Mask wearing detection algorithm based on improved YOLOv3[J].Information Recording Materials，2022，23（7）：158-161.
[7] 陈志琳，齐华，王相超.基于改进YOLOv5算法的口罩佩戴检测研究[J].电子设计工程，2022，30（22）：67-72.
CHEN Z L，QI H，WANG X C.Research on mask wearing detection based on improved YOLOv5 algorithm[J].Electronic Design Engineering，2022，30（22）：67-72.
[8] 马驰，吴华瑞，于会山.基于YOLOX的穴盘甘蓝病害检测方法[J].江苏农业科学，2023，51（8）：193-202.
MA C，WU H R，YU H S.Detection method of cabbage disease based on YOLOX[J].Jiangsu Agricultural Science，2023，51（8）：193-202.
[9] HE K，ZHANG X，REN S，et al.Spatial pyramid pooling in deep convolutional networks for visual recognition[J].IEEE Transactions on Pattern Analysis and Machine Intelligence，2014，37（9）：1904-1916.
[10] WANG C Y，LIAO H，YEH I H，et al.CSPNet：a new backbone that can enhance learning capability of CNN[C]//Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops，Seattle，2019：1571-1580.
[11] LIN T Y，DOLLAR P，GIRSHICK R，et al.Feature pyramid networks for object detection[C]//Proceedings of the 2017 IEEE Conference on Computer Vision and Pattern Recognition，2017.
[12] LIU S，QI L，QIN H，et al.Path aggregation network for instance segmentation[C]//Proceedings of the 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition，2018.
[13] WOO S，Park J，LEE J Y，et al.CBAM：convolutional block attention module[J].arXiv：1807.06521，2018.
[14] 陈钊，鲁仕康，覃章健，等.SENet优化的Deeplabv3+滑坡识别[J].科学技术与工程，2022，22（33）：14635-14643.
CHEN Z，LU S K，QIN Z J，et al.SENet optimized Deeplabv3+landslide identification[J].Science and Technology and Engineering，2022，22（33）：14635-14643.
[15] 康烨.基于深度学习的大豆幼苗期苗情快速检测研究[D].大庆：黑龙江八一农垦大学，2022.
KANG Y.Research on rapid detection of soybean seedling stage based on deep learning[D].Daqing：Heilongjiang Bayi Agricultural University，2022.
[16] REZATOFIGHI H，TSOI N，GWAK J Y，et al.Generalized intersection over union：a metric and a loss for bounding box regression[C]//Proceedings of the 2019 IEEE Conference on Computer Vision and Pattern Recognition，2019：658-666.
[17] 陈小青.无人机遥感影像拼接技术的研究[D].贵州：贵州大学，2021.
CHEN X Q.Research on UAV remote sensing image mosaic technology[D].Guizhou：Guizhou University，2021.
[18] 王兵，乐红霞，李文璟，等.改进YOLO轻量化网络的口罩检测算法[J].计算机工程与用，2021，57（8）：62-69.
WANG B，LE H X，LI W J，et al.Mask detection algorithm based on improved YOLO lightweight network[J].Computer Engineering and Applications，2021，57（8）：62-69.
[19] 刘婷，罗佩琪，范云生.基于SSD的海面小目标检测综述[J].大连海事大学学报，2022，48（4）：65-75.
LIU T，LUO P Q，FAN Y S.Overview of SSD-based detection of small targets on the sea[J].Journal of Dalian Maritime University，2022，48（4）：65-75.
[20] LIN T Y，MAIRE M，BELONGIE S，et al.Microsoft COCO：common objects in context[C]//Proceedings of the 13th European Conference on Computer Vision，Zurich.Cham：Springer，2014：740-755.
[21] LIN T Y，GOYAL P，GIRSHICK R，et al.Focal loss for dense object detection[J].IEEE Transactions on Pattern Analysis and Machine Intelligence，2020，42（2）：318-327.
[22] 钱伍，王国中，李国平.改进YOLOv5的交通实时检测鲁棒算法[J].计算机科学与探索，2022，16（1）：231-241.
QIAN W，WANG G Z，LI G P.Improved YOLOv5 traffic light real-time detection robust algorithm[J].Journal of Frontiers of Computer Science and Technology，2022，16（1）：231-241.
[23] 朱杰，王建立，王斌.基于YOLOv4-tiny改进的轻量级口罩检测算法[J].液晶与显示，2021，36（11）：1525-1534.
ZHU J，WANG J L，WANG B.Lightweight mask detection algorithm based on improved YOLOv4-tiny[J].Chinese Journal of Liquid Crystals and Displays，2021，36（11）：1525-1534.