基于YOLO v5融合通道注意力机制的海面小目标识别

潘昱辰,徐 浩,钱 夔,陆保国

(1. 南京工程学院机械工程学院, 江苏 南京 211167;2. 中国电子科技集团第二十八研究所, 江苏 南京 210007)

海面目标识别算法在渔业管理、海上救援以及海战局势判断等众多场景都有广泛应用.由于海域辽阔,海面摄像头捕捉的画面中目标面积与图像面积之比的中位数在0.08%～0.58%时定义为海面小目标识别[1].

目标识别算法具有悠久的发展历史[2],传统的目标检测算法有Viola-Jones检测器、HOG(histogram of oriented gradient)检测器等[3-4],多以滑动窗口和人工特征提取为基础, 可达到一定的目标识别效果,但计算难度高、在复杂场景下鲁棒性差.

早期的海面船舶目标识别算法多以船舶目标的特有属性为切入点进行特征提取.文献[5]利用船舶的形状特点(船头形状接近于V字、船身接近于两条平行线)提出一种基于加权的船舶检测方法,该方法适用于特定类型船舶的检测;文献[6]提出采用支持向量机(support vector machine,SVM)进行候选区训练的检测算法,通过预先对候选区的特征进行训练得到候选区的分类数据,提高了海面目标的筛查能力;文献[7]提出一种特征级神经网络算法,将BP神经网络与D-S理论(Dempster-Shafer envidence theory)结合,从多维度学习训练多传感器获取的数据,在保证识别正确率的前提下该算法具有较快的收敛速度.基于手工设计特征的海面目标检测算法对于特征种类的海面目标可行性较好,但模型泛化能力有待提高.

现有的目标识别算法多基于深度学习技术学习图像中目标的特征,对目标进行高精度的识别[8-9],对于小目标识别也颇有成效[10].典型算法有R-CNN、YOLO、SSD等[11-13],这些算法在目标检测中取得了较好的成果[14-15].

文献[16]提出一种基于YOLO v3-tiny的实时目标检测框架,采用K-means聚类方法对目标进行聚类,获得最优锚框,该方法检测速度快,但对小目标物体有漏检情况;文献[17]提出了一种基于Mask R-CNN的船舶检测方法,利用空间注意力机制来调整图像每个像素的权重,能够更好地描述目标的特征;文献[18]利用Haar级联分类器训练算法建立目标物的机器学习库,通过多组训练目标物对比试验,得出参数值的最适设定范围,为海面目标识别算法研究提供了技术支撑.

YOLO v5是深度学习的一个典型算法,其检测速度快、检测精度高,在目标检测领域有着广泛的应用[19-20],但该算法对小目标识别精度不够,在视角变换的情况下漏检率高.

本文提出一种基于深度学习的海面小目标识别算法YOLO v5-SE,以YOLO v5算法的网络架构为依托,通过对数据集以及网络架构进行优化,提高算法系统对目标尺度变化和视角变化的适应性.

1 YOLO v5-SE神经网络

YOLO v5-SE神经网络模型如图1所示.本文算法以YOLO v5网络模型为基础,首先进行数据增强,对数据集进行图像处理,包括马赛克数据增强[21]、比度增强、高斯噪声处理以及随机遮挡等;然后针对海面小目标对模型进行优化,在网络中加入通道注意力机制[22],使网络重点关注需要关注的特征,从而增加小目标稀少特征的利用率,提高模型的检测精度.

图1 YOLO v5-SE神经网络架构图

1.1 检测网络架构

YOLO v5-SE检测网络架构基于YOLO v5算法,遵循端到端训练和实时监测的原则,检测网络分为主干网络特征提取网络和特征融合网络两个模块.

1.1.1 特征提取网络

本文目标检测网络模型中的特征提取网络采用CSPDarknet53结构和Foucs结构.如图2所示,CSPDarknet53采用残差网络ResNet,通过1×1和3×3的2次卷积以及BN层和Leaky ReLu层处理,将主干的输入与输出结合.ResNet模块可以解决网络中因为深度增加而带来的梯度消失问题.

图2 残差网络ResNet

如图3所示,Focus网络采取切片操作,对输入的图像每隔1个像素取1个值,得到4个独立的特征层,再将这4个特征层进行堆叠,将宽高信息转换到通道信息,输入通道扩充4倍.检测网络输入为640×640×3的图像,切片操作后变成320×320×12的特征图,再经过1次卷积操作,最终变成320×320×32的特征图.Focus网络可以减少下采样带来的信息损失.

图3 Focus结构

1.1.2 特征融合网络

系统从特征提取网络得到3个有效特征层feat1=(80,80,256)、feat2=(40,40,512)、feat3=(20,20,1 024),利用这3个有效特征层进行特征融合网络层的构建.图4为骨干网络中的特征融合网络.特征融合网络可以帮助系统处理大范围尺度变换的物体,合理利用网络中各个层次特征.YOLO v5在Concat操作后添加CSPLayer,加强了网络特征融合能力.

图4 特征融合网络

1.2 通道注意力机制

为加强对特征的融合,在网络模型中加入通道注意力机制(SEnet).SEnet的本质是为每1个通道学习1个权重,通过这个权重来判断通道的重要程度.通过SEnet可以使得目标的特征利用最大化,减少特征丢失,对于小目标本就稀少的特征而言提高了特征的利用率.采用Squeeze操作将特征图进行全局池化,将H×W的特征图计算每个像素值的和,除以H×W,获取全局平均池化:

(1)

式中:uc为输入特征图;Fsq为Squeeze操作.

Excitation操作通过两次全连接操作和sigmoid激活函数得到相应的通道权重:

s=Fex(z,w)=σ(w2δ(w1z))

(2)

式中:δ为ReLu函数;z为输入特征图;σ为sigmoid函数;w1为第一个全连接操作,目的是降维;w2为第二个全连接操作,将维度恢复为输入维度.

SEnet嵌入到网络中后,首先对输入信息进行全局平均池化,特征图尺寸变成1×1,通道数不变,随后紧跟全连接层、激活层、全连接层、sigmoid激活层,如图5所示.

图5 通道注意力机制

1.3 损失函数

损失函数是用来衡量预测值与真实值的差距的,一定程度上决定了模型的性能.YOLO v5-SE的损失函数由分类损失、边界框损失和置信度损失构成.其中边界框损失函数采用GIOUloss,弥补了IOU只能反映出目标框与检测框的交并比相同、不能准确反映实际检测效果的缺点.YOLO v5-SE损失函数计算式为:

(1-Ci)log(1-Ci)]-

(1-Ci)log(1-Ci)]+

(1-pi(c))log(1-pi(c))]

(3)

2 试验数据分析

2.1 数据集处理

本文数据集来源于第二届全国信息融合挑战赛,由中国航空学会和中国船舶集团第七一六研究所联合提供.数据集一共2 667幅图像(不包含增强后的图像),含目标数量如图6所示.图6中横轴为目标名称,纵轴为这2 667幅图像中所包含的目标数量,每一幅图像可能包含多个目标.

图6 数据集中目标数量

2.2 试验及结果

本试验通过远程链接GPU服务器RTX3090进行训练,epochs为300,batch_size为16;算法评价指标为准确率(Precision)、召回率(Recall)、mAP@0.5,计算公式为:

P=TP/(TP+FP),R=TP/P,

(4)

式中:P为准确率;R为召回率;TP为预测为正的样本;FP为预测为负的样本;AP为P-R曲线下的面积,用来衡量数据集中一类的平均分类准确率.

为验证本文改进算法的检测性能,在YOLO v5算法基础上分别进行表1、表2所示的试验.

表1 消融试验

表2 各标签AP对比 %

由表1可见,原始模型YOLO v5的mAP仅达到27.61%;在经过一系列数据增强后,算法模型精确度提高,mAP达到53.44%;模型的网络结构加入SEnet改进后,mAP可达到59.17%;随着模型的不断改进,算法权重也有所减小.

表2为YOLO v5与YOLO v5-SE的各标签AP对比.由表2可见,Floater类小目标的AP改进前为0.00%,改进后可达到22.04%;改进后的模型能够提高海面小目标的识别精度.

图7为训练过程中loss、mAP变化曲线图.由图7可见,随着训练的加深,算法的loss随之降低,mAP逐渐升高.

(a) loss变化曲线

图8为海面目标检测结果.由图8(a)及图8(b)可见,海岸边停靠目标密集、部分目标被遮挡、目标较小,检测置信度(置信度是表示框内确实有物体的自信程度和框能将整个物体的所有特征都包括进来的自信程度)约为0.6～0.9.图8(c)为海面小目标检测,目标较为分散且尺寸极小,尤其是floater类别,肉眼只可见为一白色不规则形状点,该类别检测难度较高,在YOLO v5算法中检测AP为0,添加SEnet改进后,检测AP可达22.04%,海面小目标置信度约为0.5～0.7.综上所述,YOLO v5-SE算法在海面小目标识别中是有效的.

(a) 近海岸目标

3 结语

针对海面目标识别算法准确率较低、不能适应目标尺度变化和视角变化的情况,本文开发了海面小目标识别算法YOLO v5-SE,该算法在YOLO v5算法基础上进行了改进.基于公开数据制作一个数量充足、内容复杂的数据集,并利用高斯变换、增强对比度以及随机变换等图像处理方法对数据集进行数据增强,将mAP从27.61%提高到了59.17%.对网络进行改进,在主干网络与特征融合网络之间、特征融合网络中加入SEnet,加强特征融合.经过试验表明,相较于YOLO v5算法,改进后算法YOLO v5-SE能够提高对小目标的检测精度.