一种基于图像增强和改进YOLOv5的陨石坑检测方法

技术领域

本发明涉及深空探测领域，具体涉及了一种基于图像增强和改进YOLOv5的陨石坑检测方法。本发明的应用主要是在小行星环境暗弱，陨石坑尺度差异大、形状不规则等情况下对小行星上的难以检测的陨石坑进行高准确率的检测。

背景技术

深空探测任务在世界各航天国家的竞相发展，科学探测逐渐成为新热点。行星的地面由于内部熔岩喷发或陨石体撞击填充着各种陨石坑。石坑作为行星表面最为常见的地形特征之一，常被用作探测器着陆过程中的导航陆标及障碍规避的主要对象。由于空间探索任务发展尚未成熟，能够查阅到的各行星陨石坑图片资料相对较少，为了构造更为丰富的陨石坑样本数据库，我们采取不同类别的各星体表面的陨石坑样本，为陨石坑网络识别模型提供充足且品质优良的训练数据来源。

近年来，国内外研究学者们对小行星陨石坑检测进行了诸多研究。传统的基于形态特征的陨石坑检测算法已经发展的十分成熟。但由于在太阳系中太阳光照角度较小导致小行星表面处于暗弱环境，导致图像整体对比度偏低，且小行星尺度差异大，在复杂地形中传统的陨石坑检测能力有限，识别率不高。

随着科技发展，深度学习技术突飞猛进，特别是在目标检测领域，出现了YOLO、R-CNN、SSD等优秀的目标检测框架。将人工智能引入深空探测等领域的研究将具有更加广阔的应用实践和发展前景，基于深度学习的陨石坑识别方法研究将伴随着科技浪潮的席卷得到进一步的创新和完善。

小行星目标环境暗弱，光照度不强导致很多陨石坑特征不明显，很容易识别有误，影响后续网络模型的训练。针对小行星图像暗弱、陨石坑尺度差异大、高分辨率下小陨石坑漏检的问题，我们引入了一种基于图像增强和改进YOLOv5的陨石坑检测方法，增强了陨石坑特征。该算法相对于传统的基于直方图均衡化图像增强方法和传统CNN卷积网络算法来说，图像增强和深度卷积神经网络的方法相结合后输入到改进后的YOLOv5的网络结构进行训练，大大减少了陨石坑图像误检和漏检的情况，对于尺度差异大、形状不规则的陨石坑依旧能够进行有效的检测，可以更强的提取陨石坑的特征，提升检测的精确度。

发明内容

本发明提供了一种基于图像增强和改进YOLOv5的陨石坑检测方法，结合深度卷积神经网络的方法，以解决小行星图像暗弱、陨石坑尺度差异大、高分辨率下小陨石坑漏检等问题。

本发明采取的技术方案为：

一种基于图像增强和改进YOLOv5的陨石坑检测方法，具体步骤如下：

S1：小行星陨石坑样本数据库的构建；

S2：对S1中小行星数据库中的陨石坑图像进行图像预处理，采用基于小波变换和改进传统的单尺度Retinex图像增强的方法；

S3：将步骤S2预处理后的图像送入主干网络中进行特征提取，获得神经网络的输出；

S4：将S3输出的陨石坑特征图其输入到网络中的Neck部分，Neck部分对不同尺度的特征图进行下采样和特征融合后获得不同尺度的张量格式图像；

S5：将S4得到的不同尺度的张量格式图像送入检测网络中，分别得到每个检测框的类别、置信度和坐标位置，对S4得到的小行星陨石坑的检测框，利用平均检测准确率P、检测率R与泛化交并比3个参数进行评价。

步骤S1中小行星陨石坑样本数据库的构建：

世界各国积极开展多次小行星探测任务，本文采用爱达、爱神星及灶神星、谷神星等4个小行星带有的陨石坑图像作为原始数据集，来源于美国国家航空航天局(NASA)科学可视化工作室官网的小行星图片库。基于专家经验，利用开源图像标注工具对各图像样本中的陨石坑及其位置信息进行人工标注，得到JSON格式的标注文件，用python脚本将JSON格式的标注文件转换txt格式文件。由于采集到的陨石坑数量难以满足深度学习网络训练的需求，容易产生过拟合现象，本发明使用Mosaic、镜像、随即裁剪、随机旋转、均值滤波、加入高斯噪声和少量椒盐噪声等数据增强方法来扩充数据，将小行星图像按照一定比例扩充到8000张，增强了样本的多样性，构成了陨石坑样本数据集。将带有类别标签的图像数据集划分为训练集和测试集，其中80％作为训练集，20％作为测试集。

步骤S2中对陨石坑图像进行预处理：

先使用基于改进传统的Retinex算法对原始陨石坑图像进行处理，该算法采用引导滤波代替高斯滤波提取原始陨石坑图像的入射分量，与传统的Retinex算法相比，能更好地保留图像的边缘细节。将陨石坑图像从RGB颜色空间变换到HSV颜色空间，提取H和V分量，在V分量通道中，以原始图像为导像，通过导滤波器计算V分量的照度估计分量。得到的照度估计分量在频域内减去V分量数据，得到V分量的反射分量。S分量的自适应上拉增强校正。从频域转换到实数域，拼接H、S、V通道。颜色空间转换为RGB颜色空间，得到物体的反射分量S1。

然后再利用sym4小波的多分辨率分析能力，把图像分解得到高频分量，使用小波变换对于高频分量进行处理的具体步骤如下：

1)利用3×3大小的窗口内最大和最小灰度值I

2)根据每个像素点的局部对比度D(i,j)，利用sigm(x)＝1/[1+exp(-x)]计算得到低细节隶属度、中细节隶属度和高细节隶属度，分别用η

3)使用增强系数对高频分量的细节信息进行局部增强，解决其局部细节丢失的缺点。设低、中、高区域的增强因子分为λ

增强系数

4)增强系数λ[η(i,j)]与原始小波系数相乘得到对高频子带图增强后的小波系数，将处理后部分进行小波逆变换，重构得到增强分量S2。

最后再将两个分量S1和S2按照一定比例进行融合，得到图像增强后的陨石坑图像。这种基于小波变换和改进的Retinex图像增强方法让图像显得更加自然，图像边缘信息损失减小，图像细节显著增强，近处的小、扁状陨石坑轮廓变得立体清晰，视觉效果变得更好。

步骤S3中的主干网络：

先将经过S1、S2数据增强和图像增强等预处理后的640×640×3陨石坑图像输入到Focus结构，采用切片操作，变成320×320×12的特征图，再经过一次卷积操作，最终变成320×320×32的特征图，再送入特征提取网络中。使用比传统ResNet的网络具有更强的特征提取能力的Swin Transformer block作为主干网络提取特征，替换掉YOLOv5原始的CSPDarknet，采取滑窗操作和顶层设计，不仅可以探索图像各像素间的空间信息联系还能对特征的加权选择。一方面能引入CNN卷积操作的局部性，另一方面能解决计算繁琐问题。

所述Swin Transformer网络包括：首先使用Partition，用kernel size和stride均为4的卷积操作将图片的H和W降为原来的1/4。然后分成4个阶段，每个阶段中包括两个部分，分别是patch Merging(第一个块是线性层)和Swin Transformer Block。所述SwinTransformer块中由多个窗口多头自注意力W-MSA(window based self-attention)层和移动窗口多头自注意力SW-MSA层、多层感知机(Multi-layer perceptron，MLP)和层归一化(LayerNormaliazation，LN)等计算单元通过残差连接组成，分别够获取窗口内部的信息和窗口之间的信息。SW-MSA和W-MSA结构完全一样，区别在于SW-MSA在进行self-attention操作之前进行对向量进行了shift操作，可以理解为先剪切一张图片的上部和左部，拼接到其下部和右部，再进行窗口划分，这样每个窗口就包含划分之前相邻四个窗口的特征信息，提高了特征的利用效果，对各窗口的信息进行了融合。

步骤S4中的Neck模块：

将SwinTransformer网络中最的多尺度特征图输出作为YOLOv5改进网络的Neck模块的输入，采用了借鉴CSPnet设计的CSP2结构，经过一个Conv-BN-LeakyReLU的卷积模块，再将特征层串行通过多个池化核大小为5×5的最大池化层，然后在通道方向进行concat拼接再做进一步融合。在真实的小行星环境中，不同尺度的待检测的多类陨石坑会呈现在同一副图中，但采用的识别网络输入尺寸是固定的，很容易产生漏检。针对此问题，在Neck模块中添加了FPN网络结构，用自顶向下的上采样将深层次的特征图通过经过两倍上采样，与浅一层的特征图进行逐像素相加，向下传达强语义特征；同时将底层特征图又通过自底向上的下采样操作与高层特征图进行融合，自底向上传达强定位特征将底层特征与含有丰富语义信息的高层特征进行融合，使得每一层特征图都具有丰富的信息。Neck模块加强网络特征融合的能力，从不同的主干层对不同的检测层进行参数聚合，提升了多尺度目标的检测效果。

步骤S5中的头部检测网络：

把S4获得的特征图输入到YOLOv5的头部检测网络进行预测，头部结构的深度表示边界框偏移量、置信度、类别和先验框，且每个尺度输出都有三个不同尺寸的先验框。依据陨石坑的不同对每个区域都生成了3个高宽比不同的预测网络的anchor大小，尺寸为原输入尺寸1/8、1/16和1/32，对每个锚框进行置信度、类别和坐标位置(中心点x坐标、中心点y坐标、宽度w和高度h)预测，然后将预测网络通过NMS非极大值抑制输出的候选框映射成原图大小，框选出图像中的目标物体，最终得出检测结果。

利用泛化交并比构造回归目标框损失函数，其计算公式如下：

式中，GIOU表示泛化交并比，IOU表示交并比，A，B表示任意的两个陨石坑的检测框，C表示一个能够包住A，B的最小方框，|C\(A∪B)|表示为C减去A与B并集的面积，|C|表示为C的面积，|A∪B|表示A框与B框并集的面积，|A∩B|表示A框与B框交集的面积；

引入图像平均检测准确率P与检测率R作为评价性能指标，式中：T

本文算法的精确率及检测率最高，该算法优于现有主流目标检测算法，相比于传统的YOLOv5网络检测结果提高了5.97％，能够更好地完成小行星陨石坑检测任务。

本发明具有以下优点：

本发明针对现有的目标检测方法在暗弱环境下对于小行星陨石坑图像中陨石目标检测精度较差，中小尺寸陨石坑目标漏检多的问题，提出了一种基于图像增强和改进YOLOv5的陨石坑检测方法。

(1)与传统深度学习陨石坑目标检测算法相比，增加了在小行星暗弱环境下对陨石坑图像的预处理步骤，与传统视觉检测方法混合可以更好的增强与保留陨石坑特征，比现有技术检测速度快、识别率高的情况下还能降低网络功耗、内存和训练工作方面的要求，从而能够更好的完成小行星陨石坑检测任务。

(2)把YOLOv5本身的主干网络CSPDarknet替换为具有移动窗口的Swintransformer，改进后的YOLOv5算法与原算法相比，大大地缩小了运算量，进步了模型推理速度，检测出的陨石坑数量远远高于原算法的效果。

附图说明

图1是本发明的基于图像增强和改进YOLOv5的陨石坑检测方法的总体流程图；

图2是本发明的陨石坑原图像；

图3是本发明经过图像增强后的图像；

图4为本发明的YOLOv5网络结构框架图；

图5是本发明的Swin transfomer模块的层级结构图；

图6是本发明的Swin transfomer block结构示意图；

具体实施方式

本发明提供了一种基于图像增强和改进YOLOv5的陨石坑检测方法，利用小波变换和改进传统的单尺度Retinex算法，结合深度神经卷积网络，提高陨石坑图像的检测精度。

实施流程如图1所示，以下结合附图和具体实施案例对本发明作进一步说明。

步骤1：构建小行星陨石坑样本数据库

世界各国积极开展多次小行星探测任务，本文采用爱达、爱神星及灶神星、谷神星等4个小行星带有的陨石坑图像作为原始数据集，来源于美国国家航空航天局(NASA)科学可视化工作室官网的小行星图片库。基于专家经验，利用开源图像标注工具对各图像样本中的陨石坑及其位置信息进行人工标注，得到JSON格式的标注文件，用python脚本将JSON格式的标注文件转换txt格式文件。由于采集到的陨石坑数量难以满足深度学习网络训练的需求，容易产生过拟合现象，本发明使用Mosaic、镜像、随即裁剪、随机旋转、均值滤波、加入高斯噪声和少量椒盐噪声等数据增强方法来扩充数据，将小行星图像按照一定比例扩充到8000张，增强了样本的多样性，构成了陨石坑样本数据集。将带有类别标签的图像数据集划分为训练集和测试集，其中80％作为训练集，20％作为测试集。

步骤2：图像预处理

先使用基于改进传统的Retinex算法对原始陨石坑图像进行处理，该算法采用引导滤波代替高斯滤波提取原始陨石坑图像的入射分量，与传统的Retinex算法相比，能更好地保留图像的边缘细节。将陨石坑图像从RGB颜色空间变换到HSV颜色空间，提取H和V分量，在V分量通道中，以原始图像为导像，通过导滤波器计算V分量的照度估计分量。得到的照度估计分量在频域内减去V分量数据，得到V分量的反射分量。S分量的自适应上拉增强校正。从频域转换到实数域，拼接H、S、V通道。颜色空间转换为RGB颜色空间，得到物体的反射分量S1。

然后再利用sym4小波的多分辨率分析能力，把图像分解得到高频分量，使用小波变换对于高频分量进行处理的具体步骤如下：

1)利用3×3大小的窗口内最大和最小灰度值I

2)根据每个像素点的局部对比度D(i，j)，利用sigm(x)＝1/[1+exp(-x)]计算得到低细节隶属度、中细节隶属度和高细节隶属度，分别用η

3)使用增强系数对高频分量的细节信息进行局部增强，解决其局部细节丢失的缺点。设低、中、高区域的增强因子分为λ

增强系数

4)增强系数λ[η(i，j)]与原始小波系数相乘得到对高频子带图增强后的小波系数，将处理后部分进行小波逆变换，重构得到增强分量S2。

最后再将两个分量S1和S2按照一定比例进行融合，得到图像增强后的陨石坑图像。这种基于小波变换和改进的Retinex图像增强方法让图像显得更加自然，图像边缘信息损失减小，图像细节显著增强，近处的小、扁状陨石坑轮廓变得立体清晰，视觉效果变得更好。

步骤3：主干网络

先将经过S1、S2数据增强和图像增强等预处理后的640×640×3陨石坑图像输入到Focus结构，采用切片操作，变成320×320×12的特征图，再经过一次卷积操作，最终变成320×320×32的特征图，再送入特征提取网络中。使用比传统ResNet的网络具有更强的特征提取能力的Swin Transformer block作为主干网络提取特征，替换掉YOLOv5原始的CSPDarknet，采取滑窗操作和顶层设计，不仅可以探索图像各像素间的空间信息联系还能对特征的加权选择。一方面能引入CNN卷积操作的局部性，另一方面能解决计算繁琐问题。

所述Swin Transformer网络包括：首先使用Partition，用kernel size和stride均为4的卷积操作将图片的H和W降为原来的1/4。然后分成4个阶段，每个阶段中包括两个部分，分别是patch Merging(第一个块是线性层)和Swin Transformer Block。所述SwinTransformer块中由多个窗口多头自注意力W-MSA(window based self-attention)层和移动窗口多头自注意力SW-MSA层、多层感知机(Multi-layerperceptron，MLP)和层归一化(LayerNormaliazation，LN)等计算单元通过残差连接组成，分别够获取窗口内部的信息和窗口之间的信息。SW-MSA和W-MSA结构完全一样，区别在于SW-MSA在进行self-attention操作之前进行对向量进行了shift操作，可以理解为先剪切一张图片的上部和左部，拼接到其下部和右部，再进行窗口划分，这样每个窗口就包含划分之前相邻四个窗口的特征信息，提高了特征的利用效果，对各窗口的信息进行了融合。

步骤4：Neck模块

将SwinTransformer网络中最的多尺度特征图输出作为YOLOv5改进网络的Neck模块的输入，采用了借鉴CSPnet设计的CSP2结构，经过一个Conv-BN-LeakyReLU的卷积模块，再将特征层串行通过多个池化核大小为5×5的最大池化层，然后在通道方向进行concat拼接再做进一步融合。在真实的小行星环境中，不同尺度的待检测的多类陨石坑会呈现在同一副图中，但采用的识别网络输入尺寸是固定的，很容易产生漏检。针对此问题，在Neck模块中添加了FPN网络结构，用自顶向下的上采样将深层次的特征图通过经过两倍上采样，与浅一层的特征图进行逐像素相加，向下传达强语义特征；同时将底层特征图又通过自底向上的下采样操作与高层特征图进行融合，自底向上传达强定位特征将底层特征与含有丰富语义信息的高层特征进行融合，使得每一层特征图都具有丰富的信息。Neck模块加强网络特征融合的能力，从不同的主干层对不同的检测层进行参数聚合，提升了多尺度目标的检测效果。

步骤5：头部检测网络

把步骤4获得的特征图输入到YOLOv5的头部检测网络进行预测，头部结构的深度表示边界框偏移量、置信度、类别和先验框，且每个尺度输出都有三个不同尺寸的先验框。依据陨石坑的不同对每个区域都生成了3个高宽比不同的预测网络的anchor大小，尺寸为原输入尺寸1/8、1/16和1/32，对每个锚框进行类别、置信度和坐标位置(中心点x坐标、中心点y坐标、宽度w和高度h)预测，然后将预测网络通过非极大值抑制输出的候选框映射成原图大小，框选出图像中的目标物体，最终得出检测结果。

利用泛化交并比构造回归目标框损失函数，其计算公式如下：

式中，GIOU表示泛化交并比，IOU表示交并比，A、B表示任意的两个陨石坑的检测框，C表示一个能够包住A，B的最小方框，|C\(A∪B)|表示为C减去A与B并集的面积，|C|表示为C的面积，|A∪B|表示A框与B框并集的面积，|A∩B|表示A框与B框交集的面积。

引入图像平均检测准确率P与检测率R作为评价性能指标，式中：T

本文算法的精确率及检测率最高，该算法优于现有主流目标检测算法，相比于传统的YOLOv5网络检测结果提高了5.97％，能够更好地完成小行星陨石坑检测任务。