计算机工程与应用 ›› 2023, Vol. 59 ›› Issue (9): 198-206.DOI: 10.3778/j.issn.1002-8331.2212-0336

• 图形图像处理 • 上一篇    下一篇

改进YOLOv5的无人机影像小目标检测算法

谢椿辉,吴金明,徐怀宇   

  1. 1.中国科学院 上海高等研究院,上海 201210
    2.上海科技大学 信息科学与技术学院,上海 201210
  • 出版日期:2023-05-01 发布日期:2023-05-01

Small Object Detection Algorithm Based on Improved YOLOv5 in UAV Image

XIE Chunhui, WU Jinming, XU Huaiyu   

  1. 1.Shanghai Advanced Research Institute, Chinese Academy of Sciences, Shanghai 201210, China
    2.School of Information Science and Technology, ShanghaiTech University, Shanghai 201210, China
  • Online:2023-05-01 Published:2023-05-01

摘要: 无人机航拍影像具有目标尺度变化大、背景复杂等诸多特性,导致现有的检测器难以检测出航拍影像中的小目标。针对无人机影像中小目标误检漏检的问题,提出了改进YOLOv5的算法模型Drone-YOLO。增加了检测分支以提高模型在多尺度下的检测能力。设计了多层次信息聚合的特征金字塔网络结构,实现跨层次信息的融合。设计了基于多尺度通道注意力机制的特征融合模块,提高对小目标的关注度。将预测头的分类任务与回归任务解耦,使用Alpha-IoU优化损失函数定义,提升模型检测的效果。通过无人机影像数据集VisDrone的实验结果表明,Drone-YOLO模型较YOLOv5模型在AP50指标上提高了4.91个百分点,推理延时仅需16.78?ms。对比其他主流模型对于小目标拥有更好的检测效果,能够有效完成无人机航拍影像的小目标检测任务。

关键词: 目标检测, 无人机, 小目标, 注意力机制, 特征融合, YOLO

Abstract: UAV aerial images have many characteristics, such as large-scale changes and complex backgrounds, so it is difficult for the existing detectors to detect small objects in aerial images. Aiming at the problem of mistake detection and omission, a small object detection algorithm model Drone-YOLO is proposed. A new detection branch is added to improve the detection capability at multiple scales, meanwhile the model contains a novel feature pyramid network with multi-level information aggregation, which realizes the fusion of cross-layers information. Then a feature fusion module based on multi-scale channel attention mechanism is designed to improve the focus on small objects. The classification task of the prediction head is decoupled from the regression task, and the loss function is optimized using Alpha-IoU to improve the accuracy of detection. The experimental results of VisDrone dataset show that the Drone-YOLO has improved the AP50 by 4.91?percentage points compared with the YOLOv5, and the inference time is only 16.78?ms. Compared with other mainstream models, it has a better detection effect for small targets, and can effectively complete the task of small target detection in UAV aerial images.

Key words: object detection, unmanned aerial vehicle(UAV), small object, attention mechanism, feature fusion, YOLO