Computer Engineering and Applications ›› 2023, Vol. 59 ›› Issue (16): 170-176.DOI: 10.3778/j.issn.1002-8331.2210-0151

• Graphics and Image Processing • Previous Articles     Next Articles

Algorithm for Detecting Prohibited Items in X-Ray Images Based on Improved YOLOv5

LI Wenqiang, CHEN Li, XIE Xu, HAO Xingxing, LI Haobin   

  1. School of Information Science & Technology, Northwest University, Xi’an 710127, China
  • Online:2023-08-15 Published:2023-08-15

改进YOLOv5的X光图像违禁品检测算法

李文强,陈莉,谢旭,郝星星,李豪斌   

  1. 西北大学 信息科学与技术学院,西安 710127

Abstract: Aiming at the problems of complex background, missing detection, and imbalance of positive and negative categories in X-ray image contraband detection, an X-ray contraband detection algorithm based on YOLOv5 is proposed. Firstly, the algorithm introduces the Swin Transformer into the YOLOv5s backbone network, and uses its local self-attention and Shifted Window to improve the algorithm’s ability to extract global features of X-ray images, the spatial attention mechanism and channel attention mechanism are added after the backbone network to improve the algorithm’s ability to extract key features of contraband. Secondly, an adaptive spatial feature fusion structure is introduced to alleviate the interference of the conflict between feature maps at different levels in the feature pyramid on the model gradient. Finally, the Focal Loss is introduced to improve the background prediction loss function and classification loss function of YOLOv5s, and improve the detection ability of the algorithm in the case of imbalance between positive and negative samples and difficult and easy samples. The average detection accuracy of the algorithm in the public dataset SIXray100 reaches 57.4%, which is 4.5 percentage points higher than that of YOLOv5s; the average detection accuracy in the SIXray positive sample dataset is 90.4%, which is 2.4 percentage points higher than that of YOLOv5s. The experimental results show that the improved algorithm has a great improvement in detection accuracy compared with the original YOLOv5s algorithm, which proves the effectiveness of the algorithm.

Key words: deep learning, object detection, prohibited items detection, YOLOv5, attention mechanism

摘要: 针对X光图像违禁品检测中的复杂背景、正负类别不平衡和漏检等问题,提出一种基于YOLOv5的X光违禁品检测算法。该算法通过在YOLOv5s骨干网络中引入Swin Transformer模块,利用局部自注意力与Shifted Window机制提升模型对X光图像全局特征的提取能力,并且在主干网络后增加空间注意力机制与通道注意力机制,以提升算法对违禁品关键特征的提取能力。引入一种自适应空间特征融合结构,缓解特征金字塔中不同层级特征图之间冲突对模型梯度的干扰。引入Focal Loss函数用于改进YOLOv5s的背景预测损失函数和分类损失函数,提升算法在正负样本与难易样本失衡情况下的检测能力。该算法在公开数据集SIXray100上的平均检测精度达到57.4%,相比YOLOv5s提高了4.5个百分点;在SIXray正样本数据集上的平均检测精度达到90.4%,相比YOLOv5s提高了2.4个百分点。实验结果表明,改进后的算法相比原始YOLOv5s算法检测精度有较大提升,证明了算法的有效性。

关键词: 深度学习, 目标检测, 违禁品检测, YOLOv5, 注意力机制