基于YOLOv5的遥感图像小目标检测方法研究

2021-07-09 13:49刘博宇
电子乐园·中旬刊 2021年8期

刘博宇

摘要:針对现有目标检测算法难以在遥感图像中对小目标进行高精度检测的问题,本文提出了一种基于YOLOv5s的遥感图像小目标检测改进算法。采用Mosaic mix对原有数据增强方式进行改进,丰富了样本数据集;在网络结构中引入CBAM注意力模块,引导模型对小目标信息的关注,提升了网络针对小目标的检测性能。结果表明,改进后的算法相较于原始YOLOv5s的MAP值(IOU0.5)提升了3.12%,相较于YOLOv3算法也涨幅了1.35%。

关键词:YOLOv5s算法;小目标;数据增强;CBAM注意力模块

现阶段目标检测算法可分为两大类,一种为以R-CNN[2-3]为代表的双阶段检测,另一种为以SSD[4]、YOLO[5]为代表的单阶段检测。本文以YOLOv5s模型作为基础针对小目标进行改进工作。

1 YOLOV5s算法

YOLOv5s网络结构主要由主干网络(Backbone)、颈部(Neck)、头部(Head)三部分组成。在主干网络Backbone部分,主要完成对输入图像的特征提取,其中包含卷积模块(CONV)、瓶颈模块(C3)以及金字塔池化(SPPF)。C3能够提升网络残差学习能力及网络运算速度。颈部Neck为特征融合网络,采用自顶向下与自底向上联合的特征融合方式,更好地融合了多尺度特征[6]。头部Head输出了包含类别概率、对象得分及输出框位置的向量,网络通过各检测层输出的向量生成目标的预测边界框及类别概率。

2 改进的YOLOv5算法

2.1改进数据增强

Mosaic数据增强是对图片进行随机抽取、剪切、排列的操作方法。任意选取数据集中的4张图片按顺序拼接,最后缩放裁切为固定尺寸作为一张新的数据集图片输入模型。由于原始Mosaic数据增强中拼接方式随机性较强、空间利用率不高,使得输入模型的数据多为背景信息。针对小目标在检测时精度低、易漏检的问题,本文将原有的4张图片扩充至9张。改进后一方面能够减少无用信息输入、提高模型训练效率;另一方面增加小目标出现的概率、丰富数据集。

在此基础上,我们对Mosaic数据增强进行更深层次的探究。在原先只利用一种Mosaic增强方法的基础上,提出了Mosaic mix数据增强方式。即按照不同比例对数据集分别进行Mosaic 1(不进行图片拼接,只进行旋转、缩放、色调增强等处理)、Mosaic 4及Mosaic 9的数据增强处理。此种方式一方面能够使得训练集的尺度变换特性更加多样,丰富原有数据集中的小目标样本;另一方面能够提升网络对小目标的特征学习能力,增强网络的鲁棒性。

2.2 CBAM注意力模块

CBAM模块中包含了通道注意力模块及空间注意力模块,并能将二者高效融合。相较于传统的单通道注意力机制,CBAM拥有更加良好的检测效果。本文在YOLOv5sSPPF操作之前,以及每一次尺度预测前分别添加CBAM模块。此时,注意力模块一方面能够在信息完备的检测图上实现特征增强;另一方面能够获取更准确的特征信息,提高检测准确性。

3 实验与分析

本文利用哈尔滨工业大学遥感数据集进行训练及测试,采用MAP(均值平均精度)作为评价指标,与原始YOLOv5网络及YOLOv3网络进行对比。

3.1 数据集及预处理

哈尔滨工业大学遥感图像数据集包含4500幅目标图像,目标种类包括汽车、飞机、舰船,共3类。对上述目标图像按4:1的比例分为训练集以及测试集。

3.2 网络训练环境

本实验采用的训练环境为:Pytorch深度学习框架;YOLOv5 6.0版本

训练时设置的参数:Batchsize为64;初始学习率为0.01%;训练总迭代次数为800。

3.3 结果分析

(1)数据增强改进

Mosaic mix最终MAP值可达83.26%,相对于针对小目标检测的YOLOv3算法(MAP值81.91%)也涨幅了1.35%。

(2)CBAM注意力模块

通过对比原始YOLOv5s与引入注意力机制后的模型测试结果,人眼都较难识别正确的小目标车辆,本文改进后的目标检测器却能正确识别并能够提升原有小目标的检测概率,说明本文引入CBAM注意力模块后,确实提升了模型的小目标检测能力。

4结束语

针对小目标在目标检测中存在的精度低、易漏检的问题,本文提出一种改进的YOLOv5模型。模型基于Mosaic mix数据增强及CBAM注意力模块。实验表明,本文提出的改进方法能够快速、准确地检测出遥感图像中的小目标,MAP值可达83.26%,同时可检测出诸多原始模型漏检的小目标。下一步将对复杂背景下的遥感图像小目标检测进行研究。

参考文献(References)

[1] Girshick R, Donahue J, Darrell T. Rich feature hierarchies for accurate object detection and semantic segmentation. 2014 IEEE Conference on Computer Vision and Pattern Recognition. Columbus: IEEE,2014, 580–587

[2] Girshick R. Fast R-CNN. 2015 IEEE International Conference on Computer Vision (ICCV). Santiago: IEEE, 2015, 1440–1448.

[3] Liu W, Anguelov D, Erhan D. SSD: Single shot multibox detector.European Conference on Computer Vision. Cham: Springer, 2016,21–37

[4] Redmon J, Divvala S, Girshick R. You only look once: Unified, realtime object detection. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). Las Vegas: IEEE, 2016, 779–788.

[5] ZHU L L,GENG X,LI Z,et al.Improving YOLOv5 with Attention Mechanism for Detecting Boulders from Planetary Images[J]. Remote Sensing,2021,13(18):152-161.