CN116152591A - 模型训练方法、红外小目标检测方法、装置及电子设备 - Google Patents
模型训练方法、红外小目标检测方法、装置及电子设备 Download PDFInfo
- Publication number
- CN116152591A CN116152591A CN202211490559.8A CN202211490559A CN116152591A CN 116152591 A CN116152591 A CN 116152591A CN 202211490559 A CN202211490559 A CN 202211490559A CN 116152591 A CN116152591 A CN 116152591A
- Authority
- CN
- China
- Prior art keywords
- network
- enhancement
- data set
- model
- model training
- Prior art date
- Legal status (The legal status is an assumption and is not a legal conclusion. Google has not performed a legal analysis and makes no representation as to the accuracy of the status listed.)
- Granted
Links
- 238000001514 detection method Methods 0.000 title claims abstract description 84
- 238000012549 training Methods 0.000 title claims abstract description 78
- 238000000034 method Methods 0.000 title claims abstract description 59
- 230000007246 mechanism Effects 0.000 claims abstract description 21
- 238000005070 sampling Methods 0.000 claims abstract description 8
- 210000000988 bone and bone Anatomy 0.000 claims abstract description 5
- 238000012545 processing Methods 0.000 claims description 19
- 238000000605 extraction Methods 0.000 claims description 8
- 230000006870 function Effects 0.000 claims description 8
- 238000007781 pre-processing Methods 0.000 claims description 7
- 238000004590 computer program Methods 0.000 claims description 4
- ORILYTVJVMAKLC-UHFFFAOYSA-N Adamantane Natural products C1C(C2)CC3CC1CC2C3 ORILYTVJVMAKLC-UHFFFAOYSA-N 0.000 claims description 3
- 238000009499 grossing Methods 0.000 claims description 3
- 230000004048 modification Effects 0.000 abstract description 4
- 238000012986 modification Methods 0.000 abstract description 4
- 238000010586 diagram Methods 0.000 description 15
- 230000008569 process Effects 0.000 description 11
- 238000004891 communication Methods 0.000 description 7
- 230000006872 improvement Effects 0.000 description 5
- 208000037170 Delayed Emergence from Anesthesia Diseases 0.000 description 3
- 230000009471 action Effects 0.000 description 3
- 230000004913 activation Effects 0.000 description 3
- 238000004364 calculation method Methods 0.000 description 3
- 230000000694 effects Effects 0.000 description 3
- 230000004927 fusion Effects 0.000 description 3
- 238000011176 pooling Methods 0.000 description 3
- 230000009286 beneficial effect Effects 0.000 description 2
- 238000003331 infrared imaging Methods 0.000 description 2
- 238000010606 normalization Methods 0.000 description 2
- 230000005540 biological transmission Effects 0.000 description 1
- 238000010276 construction Methods 0.000 description 1
- 230000001351 cycling effect Effects 0.000 description 1
- 238000013135 deep learning Methods 0.000 description 1
- 230000007547 defect Effects 0.000 description 1
- 230000001419 dependent effect Effects 0.000 description 1
- 238000013461 design Methods 0.000 description 1
- 238000010801 machine learning Methods 0.000 description 1
- 238000012544 monitoring process Methods 0.000 description 1
- 230000003287 optical effect Effects 0.000 description 1
- 230000011218 segmentation Effects 0.000 description 1
- 238000006467 substitution reaction Methods 0.000 description 1
- 230000009466 transformation Effects 0.000 description 1
- 239000013598 vector Substances 0.000 description 1
Images
Classifications
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V10/00—Arrangements for image or video recognition or understanding
- G06V10/70—Arrangements for image or video recognition or understanding using pattern recognition or machine learning
- G06V10/77—Processing image or video features in feature spaces; using data integration or data reduction, e.g. principal component analysis [PCA] or independent component analysis [ICA] or self-organising maps [SOM]; Blind source separation
- G06V10/774—Generating sets of training patterns; Bootstrap methods, e.g. bagging or boosting
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06N—COMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
- G06N3/00—Computing arrangements based on biological models
- G06N3/02—Neural networks
- G06N3/08—Learning methods
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06T—IMAGE DATA PROCESSING OR GENERATION, IN GENERAL
- G06T3/00—Geometric image transformations in the plane of the image
- G06T3/40—Scaling of whole images or parts thereof, e.g. expanding or contracting
- G06T3/4053—Scaling of whole images or parts thereof, e.g. expanding or contracting based on super-resolution, i.e. the output image resolution being higher than the sensor resolution
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V10/00—Arrangements for image or video recognition or understanding
- G06V10/70—Arrangements for image or video recognition or understanding using pattern recognition or machine learning
- G06V10/82—Arrangements for image or video recognition or understanding using pattern recognition or machine learning using neural networks
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V2201/00—Indexing scheme relating to image or video recognition or understanding
- G06V2201/07—Target detection
-
- Y—GENERAL TAGGING OF NEW TECHNOLOGICAL DEVELOPMENTS; GENERAL TAGGING OF CROSS-SECTIONAL TECHNOLOGIES SPANNING OVER SEVERAL SECTIONS OF THE IPC; TECHNICAL SUBJECTS COVERED BY FORMER USPC CROSS-REFERENCE ART COLLECTIONS [XRACs] AND DIGESTS
- Y02—TECHNOLOGIES OR APPLICATIONS FOR MITIGATION OR ADAPTATION AGAINST CLIMATE CHANGE
- Y02T—CLIMATE CHANGE MITIGATION TECHNOLOGIES RELATED TO TRANSPORTATION
- Y02T10/00—Road transport of goods or passengers
- Y02T10/10—Internal combustion engine [ICE] based vehicles
- Y02T10/40—Engine management systems
Landscapes
- Engineering & Computer Science (AREA)
- Theoretical Computer Science (AREA)
- Physics & Mathematics (AREA)
- General Physics & Mathematics (AREA)
- Evolutionary Computation (AREA)
- Artificial Intelligence (AREA)
- Computing Systems (AREA)
- Health & Medical Sciences (AREA)
- General Health & Medical Sciences (AREA)
- Software Systems (AREA)
- Databases & Information Systems (AREA)
- Medical Informatics (AREA)
- Computer Vision & Pattern Recognition (AREA)
- Multimedia (AREA)
- Life Sciences & Earth Sciences (AREA)
- Biomedical Technology (AREA)
- Biophysics (AREA)
- Computational Linguistics (AREA)
- Data Mining & Analysis (AREA)
- Molecular Biology (AREA)
- General Engineering & Computer Science (AREA)
- Mathematical Physics (AREA)
- Image Analysis (AREA)
Abstract
本申请公开了一种模型训练方法、红外小目标检测方法、装置及电子设备，其中的模型训练方法包括：利用特征增强网络对数据集进行图像超分辨率重构，利用改进YOLOv5模型进行模型训练时，结合Mosaic策略和Mixup策略对数据集进行数据增强，把增强数据集输入改进的YOLOv5模型进行训练；对YOLOv5模型的改进包括：在BackBone网络的特征层之间引入坐标注意力机制，在Neck网络增加浅层特征层P2，把Head网络中C3模块的BottleNeck块替换为Swin Transformer Block结构，训练得到的检测模型能够更好地把握图像的全局上下文信息，聚焦感兴趣的目标区域，解决过小目标过采样的特征丢失问题，有更好的鲁棒性。
Description
技术领域
本发明属于图像处理领域，具体涉及一种模型训练方法、红外小目标检测方法、装置及电子设备。
背景技术
红外成像监测在军事和民用领域被广泛应用。由于红外小目标少像素，低信噪比且特征匮乏，它们的检测一直是一个难点。基于模型驱动的传统检测方法包括基于滤波器的方法、基于局部对比度机制的方法、基于数据结构的方法这几大类。这些方法存在诸多局限性，比如对目标尺寸大小的要求较高、在对比度小的图像上表现不佳等，无法适应复杂的动态检测环境。区别于需要专家知识经验来人工设计特征提取器的传统机器学习方法，数据驱动的检测方法可以通过模型训练过程自适应地从数据集中学习特征。近年来，许多基于目标检测和目标分割的深度学习方法使检测性能得到了一定的提高。然而，主流目标检测网络重复使用卷积层做下采样容易造成小目标特征信息在深层网络中的丢失，而且这些方法在图像的长距离依赖关系上建模能力差，检测结果往往存在较多的漏检和虚警。
发明内容
鉴于此，本发明旨在提出一种基于改进YOLOv5模型的红外小目标检测模型，解决过采样导致的特征丢失问题，并克服现有方法的缺陷。
第一方面，本发明提供一种模型训练方法，该方法训练得到的模型用于红外小目标检测，包括：
获取用于模型训练的数据集；
利用特征增强网络对数据集进行超分辨率特征增强，得到第一增强数据集；
对第一增强数据集进行二次数据增强得到第二增强数据集；
把第二增强数据集输入改进的YOLOv5模型进行模型训练，得到用于红外小目标检测的检测模型；
改进的YOLOv5模型的BackBone网络引入坐标注意力机制，和/或，
Neck网络增加浅层特征层P2，配合Head网络构成四层检测层，和/或，
把Head网络中C3模块中的BottleNeck块替换为Swin Transformer Block结构。
进一步地，对第一增强数据集进行二次数据增强得到第二增强数据集，包括：
利用Mosaic策略对第一增强数据集进行处理得到拼接图像；
随机抽取设定数量的拼接图像，利用Mixup策略对拼接图像处理生成新的数据样本组成第二增强数据集；
达到Mosaic策略的抽样次数要求时输出第二增强数据集。
进一步地，在YOLOv5模型的BackBone网络引入坐标注意力机制，包括：
在BackBone网络的当前特征提取层与下一特征提取层之间插入3层的坐标注意力机制，使得BackBone网络在每次下采样得到的特征图经过坐标注意力机制编码后再进入Neck网络。
进一步地，特征增强网络采用SRResNet网络。
进一步地，上述方法还包括：
对SRResNet网络进行预训练，采用平滑常数为0.9的Adam优化器优化损失函数，训练学习率为10-4，学习率迭代更新大小为106，训练迭代200次，工作线程数为4。
第二方面，本发明提供一种红外小目标检测方法，包括：
获取待检测图像；
利用特征增强网络对待检测图像进行超分辨率预处理得到重构图像；
利用上述的模型训练方法得到的目标检测模型处理重构图像，得到目标检测结果。
进一步地，特征增强网络采用SRResNet网络。
第三方面，本发明提供一种模型训练装置，包括：
数据获取单元，用于获取用于模型训练的数据集；
数据增强单元，内置有特征增强网络，用于：
利用特征增强网络对数据集进行超分辨率特征增强，得到第一增强数据集，及，对第一增强数据集进行二次数据增强得到第二增强数据集；
数据处理单元，内置有改进的YOLOv5模型，用于利用输入的第二增强数据集进行模型训练，得到用于红外小目标检测的检测模型；
改进的YOLOv5模型的BackBone网络引入坐标注意力机制，和/或，
Neck网络增加浅层特征层P2，配合Head网络构成四层检测层，和/或，
把Head网络中C3模块中的BottleNeck块替换为Swin Transformer Block结构。
第四方面，本发明提供一种红外小目标检测装置，包括：
图像获取单元，用于获取待检测图像；
图像重构单元，用于对待检测图像进行超分辨率预处理，得到重构图像；
数据处理单元，用于利用第三方面的模型训练装置训练得到的检测模型处理待检测图像，得到目标检测结果。
第五方面，本发明提供一种电子设备，包括：存储器和处理器；
所述存储器，用于存储程序；
所述处理器，用于调用存储于所述存储器中的程序，以执行上述第一方面实施例和/或结合第一方面实施例的任一种可能的实施方式提供的模型训练方法，或，执行上述第二方面实施例提供的红外小目标检测方法。
第六方面，本发明还提供一种可读存储介质，其上存储有计算机程序，所述计算机程序被处理器执行时，实现上述第一方面实施例和/或结合第一方面实施例的任一种可能的实施方式提供的模型训练方法，或，执行上述第二方面实施例提供的红外小目标检测方法。
与现有的目标检测网络相比，本发明有如下有益效果：
本发明提出的用于红外小目标检测的模型训练方法，数据进入YOLOv5模型之前利用特征增强网络对数据集进行图像超分辨率重构，显著提高检测模型的平均精度均值；利用YOLOv5模型进行模型训练时，结合Mosaic策略和Mixup策略对数据集进行数据增强，丰富小目标图像的训练样本数量；对YOLOv5模型的改进包括：在BackBone网络的特征层之间引入坐标注意力机制，利用目标的坐标位置信息降低其在深层网络的丢失率；在Neck网络增加浅层特征层P2，P2层感受野最小且分辨率最高，能够减少特征在下采样过程中位置信息的丢失；把Head网络中C3模块的BottleNeck块替换为Swin Transformer Block结构，增强模型对整幅图像的上下文理解能力，有效减少检测时发生的虚假警报。
附图说明
为了更清楚地说明本发明实施例或现有技术中的技术方案，下面将对实施例或现有技术描述中所需要使用的附图作简单地介绍，显而易见地，下面描述中的附图仅仅是本发明的实施例，对于本领域普通技术人员来讲，在不付出创造性劳动的前提下，还可以根据提供的附图获得其他的附图。
图1是本发明实施例提供的模型训练方法流程图；
图2是本发明实施例用到的SRResNet网络结构示意图；
图3是本发明实施例利用Mosaic和Mixup策略进行数据增强的流程图；
图4是本发明实施例改进的YOLOv5模型结构图；
图5是本发明实施例在YOLOv5模型中引入的坐标注意力机制示意图；
图6是本发明实施例在YOLOv5模型中添加了浅层特征层P2的FPN+PAN结构示意图；
图7是本发明实施例引入的Swin Transformer Block结构的示意图；
图8是本发明实施例提出的C3STR模块结构示意图；
图9是本发明实施例提出的模型训练工作流程图；
图10是本发明实施例提供的红外小目标检测方法流程图；
图11是本发明实施例提供的模型训练装置结构图；
图12是本发明实施例提供的红外小目标检测装置结构图；
图13是本发明实施例提供的电子设备架构图。
具体实施方式
下面将结合本发明实施例中的附图，对本发明实施例中的技术方案进行清楚、完整地描述，显然，所描述的实施例仅仅是本发明一部分实施例，而不是全部的实施例。基于本发明中的实施例，本领域普通技术人员在没有做出创造性劳动前提下所获得的所有其他实施例，都属于本发明保护的范围。
为了改进现有的目标检测网络用于检测红外小目标时表现出的不足，本发明以下提供的实施例将说明对主流目标检测网络YOLOv5模型的改进，除此之外还对模型训练过程的其他环节进行了改进，使得训练好的模型在红外小目标检测方面能够有更好的表现。
结合图1，对本发明实施例提供的模型训练方法进行说明。
S101.获取用于模型训练的数据集。
S102.利用特征增强网络对数据集进行超分辨率特征增强，得到第一增强数据集。
考虑到在实际的红外成像中理想高分辨率图像往往很难得到，此外，检测网络中的卷积层对低分辨率图像多次下采样容易造成小目标的特征丢失，本发明提出对原始的训练数据集进行超分辨率特征增强，经过特征增强网络恢复的图像可以获得较好的信噪比。
在可选的实施方式中，选择SRResNet网络作为超分辨率重构的特征增强网络。SRResNet网络不仅可以提高图像分辨率大小，还能在超分辨率过程中完整保留小目标的高频细节信息。
具体地在示例方面，用到的SRREsNet网络的结构图如图2所示，由16个用于特征提取的基于ResNet的残差模块和2个用于尺度缩放的子像素卷积模块构成。在每个残差块中，把两个64通道、步长为1的3×3的卷积核作为残差权重模块，每个卷积核后面连接批规范化层(BN)，激活函数选择的是PReLu函数。此外，将这16个残差块整体作为一个残差卷积层，与这些残差块之外添加的一个残差卷积层通过跳链构成了一个更大的残差模块。每个子像素卷积模块包含一个256通道、步长为1的3×3的卷积核，一个子像素卷积层(PixelShuffler)以及PReLu激活函数。缩放因子设置为4，因此训练好的SRResNet可以将图像分辨率提高4倍。
在COCO数据集上进行SRResNet网络的预训练。训练采用平滑常数为0.9的Adam优化器优化损失函数，训练学习率为10-4，学习率迭代更新大小为106，训练迭代200轮，工作线程数为4。对于每个小批量，批处理大小为400。使用训练好的SRResNet网络对训练数据集的所有红外图像做超分辨率重构，参考原数据集图像的xml标注文件，通过Labellmg对超分辨率后的图像重新进行手工标注。训练好的SRResNet网络用于红外小目标检测时，对待测图像做超分辨率预处理。
S103.对第一增强数据集进行二次数据增强得到第二增强数据集。
在目标检测领域常用的数据增强方法有Mosaic、Mixup、Cutout、CutMix等，为了进一步丰富小样本数据来增加训练的样本数量，同时提高模型的泛化能力，本发明实施例在示例方面提供一种结合Mosaic和Mixup策略的数据增强方法。
在示例方面，本发明实施例提供的数据增强过程如图3所示，对于输入的训练集图像，在任意一轮迭代(epoch)中，定义Batch size为一次训练所抓取的数据样本数量。针对少样本数据集的模型训练，实施例进行的数据增强的主要流程是：
(1)对每一次训练抓取的Batch size张图像，通过调整像素值、色阶饱和度增强图像的全局光学特征，通过图像旋转缩放增强全局几何学特征；(2)采用Mosaic策略，即随机裁剪、随机缩放输入的任意四张图像，然后通过任意排列、拼接它们组合成一张图像；(3)采用Mixup策略，即训练过程中，随机抽取两幅经过前述步骤(2)拼接后的训练样本，将这两个训练样本与它们各自标签构成的两个训练对按比例随机相加生成一个新的图像-标签样本；(4)根据Mosaic策略抽样次数循环上述步骤(2)和(3)，直到抽样次数n大于(Batchsize/4),将数据增强后的训练样本集合输出到检测网络进行训练。
S104.把第二增强数据集输入改进的YOLOv5模型进行模型训练，得到用于红外小目标检测的检测模型。
本发明实施例改进的YOLOv5模型的网络结构如图4所示，接下来结合图4对本发明实施例中对YOLOv5模型的改进进一步描述。
在示例方面，针对改进后的YOLOv5模型，本发明实施例将其命名为YOLOSR-IST模型。原始的YOLOv5模型分为主干网络(BackBone网络)、特征融合网络(Neck网络)和检测网络(Head网络)三大部分。BackBone网络负责目标的特征提取，由Focus模块、4个C3层、SPP模块组成；Neck网络对BackBone网络提取的特征进行增强，采用FPN+PAN结构；Head网络有3个检测头，分别对原图像进行8倍、16倍、32倍下采样，分别生成三个不同大小的特征向量，用于检测不同大小的目标。
在原始的YOLOv5结构的基础上，本发明实施例对YOLOv5模型的改进包括以下方面：
(1)在BackBone网络中引入坐标注意力机制。
具体在每一特征提取层的C3层与下一特征提取层的Conv层之间插入3层的坐标注意力机制(CA模块)，使得每一次下采样得到的不同特征图经过CA模块编码后才能进入特征金字塔网络。这样做有利于改善特征融合前不同分辨率特征图的注意力权重分配，将图注意力聚焦到小目标边缘，通过充分利用目标的坐标位置信息来降低其在深层网络中的丢失率。CA模块将通道注意力分解成两个不同坐标方向的一维特征编码，使模型可以沿着一个坐标维度保留浅层特征图上的目标位置信息，在另外一个坐标维度上捕获更高层次特征图长范围依赖的语义信息。坐标注意力机制结构简单，几乎不会增加额外的计算开销，很适合YOLOv5s这种轻量级网络。
本发明实施例引入的坐标注意力机制如图5所示，维度为C×H×W的输入是一个C通道，H高，W宽的特征图，对该输入在水平坐标方向做尺寸为(H,1)的平均池化，在竖直坐标方向做(1,W)的平均池化。水平方向编码后h高度的第c个通道的输出表达式为：
竖直方向编码后宽度w的第c个通道的输出表达式为：
接下来对池化编码后两个通道输出的特征图进行级联，用1×1尺寸的卷积核对级联的特征图完成卷积操作Conv1，经过批量归一化和非线性变换δ函数后输出具有水平和竖直方向空间信息的f特征图，其表达式为：
f＝δ(Conv1([zh,zw]))
按照不同的空间维度切分上一步输出的特征图，得到两个独立的特征图fh、fw，然后通过两个1×1尺寸的卷积计算Convh和Convw分别调整特征张量fh、fw的通道数，之后再让它们通过Sigmoid激活函数变换到0-1的范围大小，作为代表重要性级别的参数。计算公式为：
gh＝σ(Convh(fh))
gw＝σ(Convw(fw))
最后，将包含不同维度空间信息的gh和gw作为注意力权重对输入进行加权编码，输出Y的表达式是：
(2)Neck网络增加浅层特征层P2。
由于红外小目标特征信息很少，经过多次下采样后很难保留低级特征图的位置信息。原始YOLOv5将BackBone网络中获取的特征图P3、P4、P5送入Neck进行特征融合。P2特征图在原图像的感受野最小、分辨率最高，非常适合用来对缺乏特征信息、难以精确定位的红外弱小目标做分类。为了有效减少特征下采样过程中位置信息的丢失，在特征金字塔网络(FPN)中引入了P2。
添加P2的特征金字塔网络(FPN)和路径聚合网络(PAN)如图6所示，引入P2后，FPN从F5层向F2层传递上层信息，PAN从T2层向T5层传递下层信息，把两种特征信息通过两次方向相反的传递过程更好地融合在不同分辨率的特征图中，帮助检测网络学习到更丰富的特征信息，提升红外小目标的检测效果。FPN中与P2特征图具有相同分辨率大小的特征图称为F2，F2的计算公式为：
F2＝Concat(P2,Upsample(Conv(C3(F3),1,1)))
引入浅层特征图P2后Neck网络配合Head网络构成四层检测层。
(3)把Head网络中C3模块中的BottleNeck块替换为Swin Transformer Block结构。
基于先验框回归的目标检测往往忽略了小目标与其邻域外的上下文之间的特征关系，并且缺乏对一幅图像中场景的整体把握，这个缺点让YOLOv5网络在高度可变的红外场景中不够鲁棒，比如在检测极端暗淡的红外图像时会因为无法消除一些像素级噪声的干扰而误报率升高。因此，本发明实施例将Swin Transformer网络中的Swin TransformerBlock结构引入YOLOv5来增强目标检测模型的对整幅图像的上下文理解能力。
Swin Transformer Block的结构图如图7所示，Swin Transformer Block是由两层基于窗口的连续子模块组成的一个基本计算单元，两个子层之间使用残差链接。在编码过程中，先将待提取的特征图输入到基于窗口的多头自注意机制(Window based Multi-head Self Attention，W-MSA)，再输入基于移动窗口的多头自注意机制(Shifted Windowbased Multi-head Self Attention，SW-MSA)，这种连续的窗口自注意力运算有利于窗口和窗口之间的信息交流。SW-MSA和W-MSA非常相似，主要区别是输入进入到SW-MSA后，会通过移动窗口划分来增加不同的patch之间的信息交流。
如果在YOLOv5的主干网络(BackBone网络)或特征金字塔网络(FPN结构)中添加Swin Transformer Block，提取和融合特征时高分辨率特征图的自注意力计算会产生高额的计算开销，而且对小样本数据集的检测精度的提高帮助不大。因此，仅在YOLOv5的Head网络用Swin Transformer Block替换了C3模块的瓶颈块(BottleNeck块)，提出了一种C3STR模块，如图8所示，四个C3STR模块分别和四种不同分辨率的检测头级联，形成多尺度的SwinTransformer检测头。多尺度的Swin Transformer检测头可以帮助模型更好地理解不同本地信息的前后语义联系，提高了模型在高度可变红外场景中的鲁棒性，有效减少检测时产生的虚假警报。
集合上述各种改进进行的模型训练工作流程如图9所示，把训练用的数据集先送入SRResNet网络做超分辨率上采样，并完成数据集增强后再输入到YOLOSR-IST网络进行模型训练，实际进行红外小目标检测时把待检测图像先做超分辨率预处理，再输入到上述训练好的检测模型中得到检测结果。
由于检测目标有大有小、尺度不一，使用不同的网格尺寸同时对1024*1024像素图像进行区域划分，分别划分为256*256，128*128，64*64，32*32个网格的特征图，对应的网格尺寸分别为4*4，8*8，16*16，32*32。不同尺寸特征图负责预测目标的尺寸也不同。由于每个检测头有三种不同的先验锚框大小，如下表1所示，每一个检测头会预测3个不同尺寸和位置的检测目标(这3个目标的尺寸属同一级别，位置都在同一网格内，但具体尺寸和坐标位置则不相同)。即256*256特征图输出3*256*256个极小尺寸(xs)检测目标的信息，128*128特征图输出3*128*128个小尺寸(s)目标的信息，64*64特征图输出3*64*64个中尺寸(m)目标的信息，32*32特征图输出3*32*32个大尺寸(l)目标的信息。本发明训练得到的检测模型主要用于检测红外小目标，所以最终输出的目标信息主要来自256*256和128*128的检测头。
表1
通过上述的模型训练方法可以得到鲁棒性更好、适用于红外小目标检测的检测模型，下面结合图10对本发明实施例一种利用该模型进行红外小目标检测的方法进行说明。
S201.获取待检测图像。
需要检测某张红外图像中的小目标时，通过图像采样设备获取原始的红外图像。
S202.利用特征增强网络对待检测图像进行超分辨率预处理得到重构图像。
具体地，可以采用SRResNet网络对图像进行超分辨率重构。
S203.利用上述的模型训练方法得到的目标检测模型处理重构图像，得到目标检测结果。
本申请实施例还提供一种模型训练装置100，如图11所示，包括：
数据获取单元110，用于获取用于模型训练的数据集；
数据增强单元120，内置有特征增强网络，用于：
利用特征增强网络对数据集进行超分辨率特征增强，得到第一增强数据集，及，对第一增强数据集进行二次数据增强得到第二增强数据集；
数据处理单元130，内置有改进的YOLOv5模型，用于利用输入的第二增强数据集进行模型训练，得到用于红外小目标检测的检测模型；
改进的YOLOv5模型的BackBone网络引入坐标注意力机制，和/或，
Neck网络增加浅层特征层P2，配合Head网络构成四层检测层，和/或，
把Head网络中C3模块中的BottleNeck块替换为Swin Transformer Block结构。
本申请实施例所提供的模型训练装置100，其实现原理及产生的技术效果和前述方法实施例相同，为简要描述，装置实施例部分未提及之处，可参考前述方法实施例中相应内容。
本申请实施例还提供一种红外小目标检测装置200，如图12所示，包括：
图像获取单元210，用于获取待检测图像；
图像重构单元220，用于对待检测图像进行超分辨率预处理，得到重构图像；
数据处理单元230，用于利用前述模型训练装置100训练得到的检测模型处理待检测图像，得到目标检测结果。
本申请实施例所提供的红外小目标检测装置200，其实现原理及产生的技术效果和前述方法实施例相同，为简要描述，装置实施例部分未提及之处，可参考前述方法实施例中相应内容。
参阅图13，其示出了一种电子设备的硬件结构框图，包括：至少一个处理器1，至少一个通信接口2，至少一个存储器3和至少一个通信总线4；
在本申请实施例中，处理器1、通信接口2、存储器3、通信总线4的数量为至少一个，且处理器1、通信接口2、存储器3通过通信总线4完成相互间的通信；
处理器1可能是一个中央处理器CPU，或者是特定集成电路ASIC(ApplicationSpecific Integrated Circuit)，或者是被配置成实施本发明实施例的一个或多个集成电路等；
存储器3可能包含高速RAM存储器，也可能还包括非易失性存储器(non-volatilememory)等，例如至少一个磁盘存储器；
其中，存储器存储有程序，处理器可调用存储器存储的程序，所述程序用于：实现前述模型训练方法的各个处理流程，或，实现前述红外小目标检测方法的各个处理流程。
本发明实施例还提供一种可读存储介质，其上存储有计算机程序，所述计算机程序被处理器执行时，实现上述实施例提供的模型训练方法，或，执行上述实施例提供的红外小目标检测方法。
最后，还需要说明的是，在本文中，诸如第一和第二等之类的关系术语仅仅用来将一个实体或者操作与另一个实体或操作区分开来，而不一定要求或者暗示这些实体或操作之间存在任何这种实际的关系或者顺序。而且，术语“包括”、“包含”或者其任何其他变体意在涵盖非排他性的包含，从而使得包括一系列要素的过程、方法、物品或者设备不仅包括那些要素，而且还包括没有明确列出的其他要素，或者是还包括为这种过程、方法、物品或者设备所固有的要素。在没有更多限制的情况下，由语句“包括一个……”限定的要素，并不排除在包括所述要素的过程、方法、物品或者设备中还存在另外的相同要素。
以上实施例仅用以说明本发明的技术方案，而非对其限制；尽管参照前述实施例对本发明进行了详细的说明，本领域的技术人员应当理解：其依然可以对前述实施例所记载的技术方案进行修改，或者对其中部分技术特征进行等同替换；而这些修改或者替换，并不使相应技术方案的本质脱离本发明实施例技术方案的精神和范围。
Claims (11)
1.一种模型训练方法，所述方法训练得到的模型用于红外小目标检测，包括：
获取用于模型训练的数据集；
利用特征增强网络对所述数据集进行超分辨率特征增强，得到第一增强数据集；
对所述第一增强数据集进行二次数据增强得到第二增强数据集；
把所述第二增强数据集输入改进的YOLOv5模型进行模型训练，得到用于红外小目标检测的检测模型；
所述改进的YOLOv5模型在BackBone网络引入坐标注意力机制，和/或，
Neck网络增加浅层特征层P2，配合Head网络构成四层检测层，和/或，
把Head网络中C3模块中的BottleNeck块替换为Swin Transformer Block结构。
2.根据权利要求1所述的模型训练方法，其特征在于，对所述第一增强数据集进行二次数据增强得到第二增强数据集，包括：
利用Mosaic策略对所述第一增强数据集进行处理得到拼接图像；
随机抽取设定数量的拼接图像，利用Mixup策略对拼接图像处理生成新的数据样本组成第二增强数据集；
达到Mosaic策略的抽样次数要求时输出第二增强数据集。
3.根据权利要求1所述的模型训练方法，其特征在于，所述在BackBone网络引入坐标注意力机制，包括：
在BackBone网络的当前特征提取层与下一特征提取层之间插入3层的坐标注意力机制，使得BackBone网络在每次下采样得到的特征图经过坐标注意力机制编码后再进入Neck网络。
4.根据权利要求1所述的模型训练方法，其特征在于，所述特征增强网络采用SRResNet网络。
5.根据权利要求1所述的模型训练方法，其特征在于，所述方法还包括：
对所述SRResNet网络进行预训练，采用平滑常数为0.9的Adam优化器优化损失函数，训练学习率为10-4，学习率迭代更新大小为106，训练迭代200次，工作线程数为4。
6.一种红外小目标检测方法，其特征在于，包括：
获取待检测图像；
利用特征增强网络对所述待检测图像进行超分辨率预处理得到重构图像；
利用权利要求1～5任一项所述的模型训练方法得到的检测模型处理所述重构图像，得到目标检测结果。
7.根据权利要求6所述的红外小目标检测方法，其特征在于，所述特征增强网络采用SRResNet网络。
8.一种模型训练装置，其特征在于，包括：
数据获取单元，用于获取用于模型训练的数据集；
数据增强单元，内置有特征增强网络，用于：
利用特征增强网络对数据集进行超分辨率特征增强，得到第一增强数据集，及，对第一增强数据集进行二次数据增强得到第二增强数据集；
数据处理单元，内置有改进的YOLOv5模型，用于利用输入的第二增强数据集进行模型训练，得到用于红外小目标检测的检测模型；
所述改进的YOLOv5模型在BackBone网络引入坐标注意力机制，和/或，
Neck网络增加浅层特征层P2，配合Head网络构成四层检测层，和/或，
把Head网络中C3模块中的BottleNeck块替换为Swin Transformer Block结构。
9.一种红外小目标检测装置，其特征在于，包括：
图像获取单元，用于获取待检测图像；
图像重构单元，用于对所述待检测图像进行超分辨率预处理，得到重构图像；
数据处理单元，用于利用如权利要求8所述的模型训练装置训练得到的检测模型处理待检测图像，得到目标检测结果。
10.一种电子设备，其特征在于，包括：存储器和处理器；
所述存储器，用于存储程序；
所述处理器，用于调用存储于所述存储器中的程序，以执行如权利要求1～5任一项所述的模型训练方法，或，执行如权利要求6～7任一项所述的红外小目标检测方法。
11.一种可读存储介质，其特征在于，所述可读存储介质上存储有计算机程序，所述计算机程序被处理器执行时，实现如权利要求1～5任一项所述的模型训练方法，或，执行如权利要求6～7任一项所述的红外小目标检测方法。
Priority Applications (1)
| Application Number | Priority Date | Filing Date | Title |
|---|---|---|---|
| CN202211490559.8A CN116152591B (zh) | 2022-11-25 | 2022-11-25 | 模型训练方法、红外小目标检测方法、装置及电子设备 |
Applications Claiming Priority (1)
| Application Number | Priority Date | Filing Date | Title |
|---|---|---|---|
| CN202211490559.8A CN116152591B (zh) | 2022-11-25 | 2022-11-25 | 模型训练方法、红外小目标检测方法、装置及电子设备 |
Publications (2)
| Publication Number | Publication Date |
|---|---|
| CN116152591A true CN116152591A (zh) | 2023-05-23 |
| CN116152591B CN116152591B (zh) | 2023-11-07 |
Family
ID=86351502
Family Applications (1)
| Application Number | Title | Priority Date | Filing Date |
|---|---|---|---|
| CN202211490559.8A Active CN116152591B (zh) | 2022-11-25 | 2022-11-25 | 模型训练方法、红外小目标检测方法、装置及电子设备 |
Country Status (1)
| Country | Link |
|---|---|
| CN (1) | CN116152591B (zh) |
Cited By (3)
| Publication number | Priority date | Publication date | Assignee | Title |
|---|---|---|---|---|
| CN116342596A (zh) * | 2023-05-29 | 2023-06-27 | 云南电网有限责任公司 | 一种基于YOLOv5改进的变电站设备螺母缺陷识别检测方法 |
| CN116912604A (zh) * | 2023-09-12 | 2023-10-20 | 浙江大华技术股份有限公司 | 模型训练方法、图像识别方法、装置以及计算机存储介质 |
| CN116996397A (zh) * | 2023-09-27 | 2023-11-03 | 之江实验室 | 一种网络丢包优化的方法、装置、存储介质及电子设备 |
Citations (5)
| Publication number | Priority date | Publication date | Assignee | Title |
|---|---|---|---|---|
| US20220066456A1 (en) * | 2016-02-29 | 2022-03-03 | AI Incorporated | Obstacle recognition method for autonomous robots |
| CN114220015A (zh) * | 2021-12-21 | 2022-03-22 | 一拓通信集团股份有限公司 | 一种基于改进YOLOv5的卫星图像小目标检测方法 |
| CN114550148A (zh) * | 2022-01-14 | 2022-05-27 | 山东师范大学 | 基于深度学习的严重遮挡商品的识别检测计数方法与系统 |
| CN114677362A (zh) * | 2022-04-08 | 2022-06-28 | 四川大学 | 基于改进YOLOv5的表面缺陷检测方法 |
| CN114758255A (zh) * | 2022-04-02 | 2022-07-15 | 桂林电子科技大学 | 一种基于yolov5算法的无人机检测方法 |
-
2022
- 2022-11-25 CN CN202211490559.8A patent/CN116152591B/zh active Active
Patent Citations (5)
| Publication number | Priority date | Publication date | Assignee | Title |
|---|---|---|---|---|
| US20220066456A1 (en) * | 2016-02-29 | 2022-03-03 | AI Incorporated | Obstacle recognition method for autonomous robots |
| CN114220015A (zh) * | 2021-12-21 | 2022-03-22 | 一拓通信集团股份有限公司 | 一种基于改进YOLOv5的卫星图像小目标检测方法 |
| CN114550148A (zh) * | 2022-01-14 | 2022-05-27 | 山东师范大学 | 基于深度学习的严重遮挡商品的识别检测计数方法与系统 |
| CN114758255A (zh) * | 2022-04-02 | 2022-07-15 | 桂林电子科技大学 | 一种基于yolov5算法的无人机检测方法 |
| CN114677362A (zh) * | 2022-04-08 | 2022-06-28 | 四川大学 | 基于改进YOLOv5的表面缺陷检测方法 |
Non-Patent Citations (2)
| Title |
|---|
| SUN, T 等: "Small object detection method based on YOLOv5 improved model", 《2022 IEEE 5TH INTERNATIONAL CONFERENCE ON INFORMATION SYSTEMS AND COMPUTER AIDED EDUCATION (ICISCAE》, pages 934 - 940 * |
| 张天宇: "基于深度学习的唐卡图像目标检测算法研究", 《中国优秀硕士学位论文全文数据库哲学与人文科学辑》, pages 30 - 35 * |
Cited By (6)
| Publication number | Priority date | Publication date | Assignee | Title |
|---|---|---|---|---|
| CN116342596A (zh) * | 2023-05-29 | 2023-06-27 | 云南电网有限责任公司 | 一种基于YOLOv5改进的变电站设备螺母缺陷识别检测方法 |
| CN116342596B (zh) * | 2023-05-29 | 2023-11-28 | 云南电网有限责任公司 | 一种基于YOLOv5改进的变电站设备螺母缺陷识别检测方法 |
| CN116912604A (zh) * | 2023-09-12 | 2023-10-20 | 浙江大华技术股份有限公司 | 模型训练方法、图像识别方法、装置以及计算机存储介质 |
| CN116912604B (zh) * | 2023-09-12 | 2024-01-16 | 浙江大华技术股份有限公司 | 模型训练方法、图像识别方法、装置以及计算机存储介质 |
| CN116996397A (zh) * | 2023-09-27 | 2023-11-03 | 之江实验室 | 一种网络丢包优化的方法、装置、存储介质及电子设备 |
| CN116996397B (zh) * | 2023-09-27 | 2024-01-09 | 之江实验室 | 一种网络丢包优化的方法、装置、存储介质及电子设备 |
Also Published As
| Publication number | Publication date |
|---|---|
| CN116152591B (zh) | 2023-11-07 |
Similar Documents
| Publication | Publication Date | Title |
|---|---|---|
| US11551333B2 (en) | Image reconstruction method and device | |
| CN116152591B (zh) | 模型训练方法、红外小目标检测方法、装置及电子设备 | |
| CN111369440B (zh) | 模型训练、图像超分辨处理方法、装置、终端及存储介质 | |
| Chen et al. | MICU: Image super-resolution via multi-level information compensation and U-net | |
| CN109064405A (zh) | 一种基于双路径网络的多尺度图像超分辨率方法 | |
| An et al. | TR-MISR: Multiimage super-resolution based on feature fusion with transformers | |
| CN113011562B (zh) | 一种模型训练方法及装置 | |
| CN112070664A (zh) | 一种图像处理方法以及装置 | |
| CN113865859A (zh) | 多尺度多源异构信息融合的齿轮箱状态故障诊断方法 | |
| CN114898284B (zh) | 一种基于特征金字塔局部差异注意力机制的人群计数方法 | |
| US20240062347A1 (en) | Multi-scale fusion defogging method based on stacked hourglass network | |
| CN116205962B (zh) | 基于完整上下文信息的单目深度估计方法及系统 | |
| CN107729885B (zh) | 一种基于多重残差学习的人脸增强方法 | |
| CN117408924A (zh) | 一种基于多重语义特征融合网络的低光照图像增强方法 | |
| CN117576483B (zh) | 基于多尺度卷积自编码器的多源数据融合地物分类方法 | |
| Ren et al. | A lightweight object detection network in low-light conditions based on depthwise separable pyramid network and attention mechanism on embedded platforms | |
| CN117392499A (zh) | 一种显著性目标的检测方法及装置 | |
| CN117058392A (zh) | 一种基于卷积局部增强的多尺度Transformer图像语义分割方法 | |
| CN115205148A (zh) | 基于双路径残差网络的图像去模糊方法 | |
| CN114565764A (zh) | 基于舰船实例分割的港口全景感知系统 | |
| CN113012072A (zh) | 一种基于注意力网络的图像运动去模糊方法 | |
| Zhang et al. | Dynamic Long-Short Range Structure Learning for Low-Illumination Remote Sensing Imagery HDR Reconstruction | |
| Li et al. | Super-resolution of fisheye rectified image based on deep multi-path cascaded network | |
| Feng et al. | An underwater image enhancement strategy based on pyramid attention mechanism | |
| CN116797485B (zh) | 一种基于数据合成的低照度图像增强方法及装置 |
Legal Events
| Date | Code | Title | Description |
|---|---|---|---|
| PB01 | Publication | ||
| PB01 | Publication | ||
| SE01 | Entry into force of request for substantive examination | ||
| SE01 | Entry into force of request for substantive examination | ||
| GR01 | Patent grant | ||
| GR01 | Patent grant |