目录
一、本文介绍
本文将为读者带来一种适用多种复杂场景的全能图像修复网络AIrNet网络,其由对比基降解编码器(CBDE)和降解引导修复网络(DGRN)两个网络组成,能够在未知损坏类型和程度的情况下恢复受损图像。此外、还提供将其与YOLOv8等目标检测模型有效结合改进的方法,使其能显著提升模型在处理受损图像时的性能,并对自定义数据集的训练以及半自动化标注的方法进行简单介绍,帮助力读者快速在自己的数据集上进行实验与应用。
二、AirNet原理介绍
官方论文地址:All-in-One Image Restoration for Unknown Corruption (thecvf.com)
官方代码地址:https://github. com/XLearning-SCU/2022-CVPR-AirNet
摘要:在本文中,我们研究了一个图像恢复领域的挑战性问题,即如何开发一种全能方法,能够从各种未知类型和程度的损坏中恢复图像。为此,我们提出了一种全能图像恢复网络(AirNet),该网络由两个神经模块组成,分别是对比基降解编码器(CBDE)和降解引导修复网络(DGRN)。AirNet的主要优势体现在两个方面。首先,它是一个全能解决方案,能够在一个网络中恢复各种受损图像。其次,AirNet不依赖于损坏类型和程度的先验知识,仅使用观察到的受损图像进行推理。这两个优势使得AirNet在现实世界场景中更加灵活和经济,因为在这些场景中,损坏的先验知识很难得知,且损坏情况会随空间和时间而变化。广泛的实验结果表明,所提出的方法在四个具有挑战性的数据集上优于17种图像恢复基线方法。
AIrNet是一种创新的图像修复网络,专为处理多种未知和复杂的图像损坏场景设计。它结合了对比基降解编码器(CBDE)和降解引导修复网络(DGRN)两个核心组件,能够在不依赖具体损坏类型或程度先验知识的情况下,有效恢复受损图像的质量。这种网络结构不仅提升了图像修复的泛化能力,还显著增强了修复效果的自然度和准确性。
2.1 对比基降解编码器(CBDE)
CBDE是AIrNet的第一阶段,负责从受损图像中提取关键特征,并识别出图像中的降解模式(如模糊、噪声、压缩失真等)。该编码器通过对比学习机制,将受损图像与一系列预设的“干净-降解”图像对进行对比,从而学习如何准确表征图像中的降解特性。这一过程为后续的修复工作提供了重要的先验信息。
2.2 降解引导修复网络(DGRN)
DGRN利用CBDE提取的特征和降解信息,指导图像修复过程。该网络设计有强大的生成能力,能够基于降解特征生成对应的修复策略,并应用于受损图像,以恢复其原始质量。DGRN通过多层卷积、残差连接和注意力机制等先进技术,确保修复结果既符合图像的自然规律,又能有效去除各种形式的损坏。
三、yolov8与AirNet结合修改教程
3.1 核心代码文件的创建与添加
3.1.1 AirNet.py文件添加
找到ultralytics/nn文件夹,建立Addmoudules文件夹,并在该文件夹下建立AirNet.py文件,如下图。
AirNet.py添加核心代码,如下
import math
import torch
import torch.nn as nn
from torch.nn.modules.utils import _pair
from mmcv.ops import modulated_deform_conv2d
__all__ = ['AirNet']
class DCN_layer(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1,
groups=1, deformable_groups=1, bias=True, extra_offset_mask=True):
super(DCN_layer, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = _pair(kernel_size)
self.stride = stride
self.padding = padding
self.dilation = dilation
self.groups = groups
self.deformable_groups = deformable_groups
self.with_bias = bias
.........
注:全部代码私信博主获取
3.1.2 __init__.py文件添加
在Addmoudules文件夹建立__init__.py文件,如下图
__init__.py文件 添加代码,如下
from .AirNet import *
3.2 yolov8 的tasks.py 文件修改
找到ultralytics/nn文件夹下的task.py 文件并打开
3.2.1 导入AirNet模块
3.2.2 添加AirNet模块
注:其中红色框为添加部分
3.3 模型配置文件添加
3.3.1 构建模型配置文件yolov8-airnet.yaml
找到ultralytics/cfg/models/v8文件夹,并在该文件夹下建立yolov8-airnet.yaml文件,如下图
3.3.2 配置网络结构
打开 yolov8-airnet.yaml 文件,在其中配置网络结构,如下
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs
s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs
m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs
l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, AirNet, []] # 0-P1/2
- [-1, 1, Conv, [64, 3, 2]] # 1-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 2-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 4-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 6-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 8-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 10
# YOLOv8.0n head
...
注:全部代码私信博主获取
四、自定义数据集
4.1数据标注
新建两个文件夹如下,其中images用于保存待标注的图片数据,Annotations用于保存标注过后xml格式的标签文件
打开 labelimg 标注工具(标注工具自行检索下载),如下图,配置images,Annotations文件夹所在路径,然后进行标注
标注后将在Annotations文件夹下生成如下文件
4.2 数据格式转换
如下图新建label文件夹,用于保存由Annotations文件夹下的xml标签文件转换后,适合yolov8训练的txt格式的标签文件
定位至Yolov8_Dataset_Processing/xmls2yolo.py文件并打开,再定位如下图到红色框处,需配置检测的标签列表,以及xml标签路径(Annotations),转换后txt格式的标签路径(label)。
配置完成后运行xmls2yolo.py,即可在label文件夹下生成所需txt格式的标签文件,如下图
注:全部代码联系博主获得
4.3 数据集划分
定位至Yolov8_Dataset_Processing/dataset_split.py文件,打开并定位至红色框处,需配置
classes、source_file_path 、images_file_path ,如下图所示
配置完成后运行dataset_split.py,即可在lmages同级文件夹下生成划分后的训练集、验证集、测试集,如下图
4.4 训练数据集配置文件添加
定位至ultralytics/cfg/datasets文件夹下,新建训练数据集的配置文件a_sleep.yaml如下图
在a_sleep.yaml 文件中进行训练、验证、测试数据的配置,如下图
注:若有多个训练、验证、测试数据,只需在对应位置继续添加即可,参考注释掉的部分
五、模型训练
5.1 训练文件的创建与配置
在ultralytics的同级目录处创建train.py 文件
点击打开train.py文件,如下图需配置 1,与 2处,其中1配置的是本文3.3构建模型配置文件yolov8-airnet.yaml。2处配置的是本文4.4 新建训练数据集的配置文件a_sleep.yaml。
5.2 训练过程展示
六、半自动标注
半自动化标注,解决的是标注数据量庞大,难以人工从头进行一一标注的问题。主要步骤是,先标注两千张数据,进行模型训练,然后用训练好的模型,对剩下未标注的数据集进行推理,得到模型推理后的结果,将这些结果转换为xml格式的标签文件。最好利用标注软件labelimg打开,就得到大量已经有预标注的图片,后续工作仅需微调这些标注即可,大大减少了标注的工作量。以下是半自动化标注流程
6.1.加载训练权重进行待标注数据集的推理
在ultralytics的同级目录处创建predict.py 文件,如下图中1处所示,并在其中配置训练好的权重模型、待推理标注的图片文件夹路径,以及是否将推理结果保存为txt文件,如下图中2、3、4处所示.配置后点击运行
运行结束后将在ultralytics-main/runs/detect/predict/labels/目录下生成txt标签文件
6.2.将推理结果txt转换为xml标签文件
在Yolov8_Dataset_Processing/的目录下创建yolo2xml.py文件,如下图中1处所示,并在其中配置训练好的图片路径、推理后的txt路径,转换为xml待微调的标签路径,以及类别对照字典,如下图中2、3、4、5处所示.配置后点击运行
6.3.打开标注软件labelimg,微调标注
打开labelimg后配置1与2 两个地方,1是图片文件夹路径,2是转换为xml待微调的标签路径。完成后图片将生成对应的标注框,若标注准确就跳转到下一张图,不准确则调整
七、训练结果可视化
7.1多训练结果对比
在ultralytics的同级目录处创建Multi_Training_Results_Comparison.py文件,如下图中1处所示,并在其中配置训练好训练结果保持时指定的路径以及需要对比的哪几次训练,如下图中2、3处所示,对比的是exp2、与exp4的训练结果,配置后点击运行
运行后将在运行文件的同级目录下生成对比后的结果图metrics_curves.png、loss_curves.png
7.2可视化热力图
在ultralytics的同级目录处创建Multi_Training_Results_Comparison.py文件,如下图中1处所示,并在其中2处配置相关参数,参数对照信息如下
配置后点击运行
运行后在同级目录下找到result文件夹将获得对应热力图如下
7.3计算推理效率 FPS
在ultralytics的同级目录处创建yolov8_Inference_Speed_Test.py文件,如下图中1处所示,并在其中2处配置相关参数,参数对照信息如下
配置后 命令行运行 python yolov8_Inference_Speed_Test.py 命令,得下图
八、模型部署
8.1 web网页部署
在ultralytics的同级目录处创建YOLOv8-web-app-master文件夹,如下图中1处所示,并在2处的detection目录下放置用于推理的模型。
完成以上配置即可在命令行运行
streamlit run app.py
得到
将上述http地址输入浏览器搜索栏,得到
如上图,web app 应用的操作流程是,通过1处下拉栏,选择上面配置的detection目录下放置用于推理的模型;然后通过2处选择推理的类型,支持视频流和图片;接着通过3选择本地图片或者视频进行推理,最后点击4执行推理。
九、总结
本文介绍了一种适用于多种复杂场景的全能图像修复网络AIrNet,AIrNet通过对比基降解编码器和降解引导修复网络两个核心组件,能够在未知损坏类型和程度的情况下有效恢复受损图像的质量。同时,本文还提供了将AIrNet与YOLOv8结合进行目标检测的方法,并介绍了自定义数据集的训练和半自动标注的流程。这些方法和流程可以帮助读者在自己的数据集上进行实验与应用,提高图像修复和目标检测的准确性和效率。
注:上述构建自定义的数据的方法,针对所有yolo系列模型检测数据构建与方法均有效。全部代码私信博主获取。结合全部代码阅读本博客更为清晰明了