YOLOv8改进 | 自定义数据集训练 | AirNet助力YOLOv8检测

一、本文介绍

二、AirNet原理介绍

2.1 对比基降解编码器（CBDE）

2.2 降解引导修复网络（DGRN）

三、yolov8与AirNet结合修改教程

3.1 核心代码文件的创建与添加

3.1.1 AirNet.py文件添加

3.1.2 __init__.py文件添加

3.2 yolov8 的tasks.py 文件修改

3.2.1 导入AirNet模块

3.2.2 添加AirNet模块

3.3 模型配置文件添加

3.3.1 构建模型配置文件yolov8-airnet.yaml

3.3.2 配置网络结构

四、自定义数据集

4.1数据标注

4.2 数据格式转换

4.4 训练数据集配置文件添加

五、模型训练

5.1 训练文件的创建与配置

5.2 训练过程展示

六、半自动标注

七、总结

一、本文介绍

本文将为读者带来一种适用多种复杂场景的全能图像修复网络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 训练过程展示

六、半自动标注

七、总结

本文介绍了一种适用于多种复杂场景的全能图像修复网络AIrNet，AIrNet通过对比基降解编码器和降解引导修复网络两个核心组件，能够在未知损坏类型和程度的情况下有效恢复受损图像的质量。同时，本文还提供了将AIrNet与YOLOv8结合进行目标检测的方法，并介绍了自定义数据集的训练和半自动标注的流程。这些方法和流程可以帮助读者在自己的数据集上进行实验与应用，提高图像修复和目标检测的准确性和效率。

注：上述构建自定义的数据的方法，针对所有yolo系列模型检测数据构建与方法均有效。全部代码私信博主获取。结合全部代码阅读本博客更为清晰明了