Description:
Ensuring both the accuracy of vehicle target detection and meeting real-time requirements is crucial in traffic videos. The YOLOv5s target detection framework, known for its accuracy and efficiency, has attracted attention in academic circles. However, there are still some features that can be optimized. First of all, the detection subnet of the YOLOv5s framework cannot smoothly convert complex feature maps into relatively sparse target prediction boxes. To solve this, we integrate a self-attention-based gating mechanism into the detection subnet, forming the YOLOv5s-SAG network. Secondly, the loss function of CIoU used by YOLOv5s pays insufficient attention to the overlapping area of the detection frame, which can be used as metric for measuring target detection effectiveness. We add the loss term of area ratio to CIoU to further improve the modeling ability. Finally, the current multi-class Non-Maximum Suppression algorithm can cause high overlap of multi-class detection frames. To improve it, we propose a multi-class CS-NMS algorithm based on category suppression. Experimental results show an approximately 8% improvement in the mAP50 index on the UA-DETRAC dataset compared with YOLOv5s. The proposed algorithm also achieves better detection results compared to mainstream target detection algorithms and meets the real-time requirements of traffic video analysis.