Enhancing YOLOv8 with Attention Task Alignment Head for Prohibited Item Detection in Complex X-Ray Images
Zhihan Wang, Huiqian Du, Min Xie · 2024
Detecting prohibited items in X-ray images is challenging due to the complex backgrounds often encountered in security inspection scenarios. When prohibited items overlap with other objects, the inherent conflict between regression and classification tasks becomes more pronounced. To address this issue, we propose an Attention Task Alignment Head (ATAH) to enhance the YOLOv8 model. ATAH dynamically aligns regression and classification tasks by restructuring the network's extracted features and distributing them across the classification and regression branches. Each branch incorporates a Layer Attention Block (LAB) to adjust weights based on task-specific requirements. Additionally, the regression branch is designed to handle complex spatial variations in the images by utilizing Deformable Convolution (DCN). We also introduce Slide Loss to focus the model's learning on challenging samples. Experimental results on the PIDRay dataset demonstrate that our approach significantly outperforms the YOLOv8 benchmark.