YOLOv11改进-特殊场景检测篇-适用多种复杂场景的全能图像修复网络AirNet助力yolov11检测（全网独家首发）

一、本文介绍

本文给大家带来的改进机制是一种适用多种复杂场景的全能图像修复网络AirNet，其由 对比基降解编码器（CBDE） 和 降解引导修复网络（DGRN） 两个神经模块组成，能够在未知损坏类型和程度的情况下恢复受损图像。这两部分共同工作，能够处理多种类型的图像退化，而无需预先知道损坏的具体信息。本文的内容为专栏读者指定发布。

二、原理介绍

官方论文地址： 官方代码地址点击此处即可跳转 

官方代码地址：

文章提出的一种名为 AirNet （All-in-one Image Restoration Network ）的全能图像修复网络。AirNet由 对比基降解编码器（CBDE） 和 降解引导修复网络（DGRN） 两个神经模块组成，能够在未知损坏类型和程度的情况下恢复受损图像。这两部分共同工作，能够处理多种类型的图像退化，而无需预先知道损坏的具体信息。

一. 对比基降解编码器（CBDE）:

对比基降解编码器（Contrastive-Based Degraded Encoder, CBDE）是全能型图像恢复网络（AirNet）的核心组成部分之一，它主要负责从受损图像中学习损坏的表示。这种编码器采用了对比学习方法，通过对比不同的图像损坏样式来提取关键的损坏特征，从而支持后续的图像恢复过程。

CBDE的主要任务是从输入的受损图像中提取一个隐含的表示（标记为z），这个表示能够捕捉到图像损坏的本质特征。这个表示是在没有先验知识（即不知道具体的损坏类型和程度）的情况下提取的，这是其设计中的一大创新。CBDE使用对比学习的框架来训练。在这个框架中，网络通过最大化相同损坏类型之间的一致性（正样本对）并最小化不同损坏类型之间的一致性（负样本对）来学习损坏的表示。同时，CBDE不仅仅学习一个简单的损坏标签，而是学习一个能够反映损坏复杂性和多样性的高维特征表示。CBDE的一个显著特点是它不依赖于任何关于损坏类型或程度的先验知识。这使得AirNet可以灵活应对多种未知或未预见的图像损坏情况，这在实际应用中是非常宝贵的。

二. 降解引导修复网络（DGRN）：

降解引导修复网络（Degradation-Guided Restoration Network, DGRN）使用从对比基降解编码器（CBDE）学习到的降解特征来引导图像的恢复过程。其目标是利用这些降解特征恢复出清晰的图像，即使在损坏未知的情况下也能有效工作。DGRN的设计旨在充分利用CBDE提取的损坏表示（z），根据这些信息来恢复图像。以下是DGRN的一些关键组件和功能：

1.降解引导模块（Degradation-Guided Modules, DGM）：每个DGM都包含可变形卷积（Deformable Convolution, DCN）层和空间特征变换（Spatial Feature Transform, SFT）层。这些模块根据输入的降解特征z动态调整其处理方式，例如，通过改变卷积核的形状（DCN）或调整特征的尺度和偏移（SFT）。

2.降解引导组（Degradation-Guided Groups, DGG）：DGRN由多个DGG组成，每个DGG包含若干个DGM。这种分组结构有助于分阶段地细化恢复过程，每一组都在前一组的基础上进一步提高图像质量。

3.特征融合与递归处理：在DGRN中，不同层级的输出可以被递归地反馈到网络中，用于进一步的处理。

三、核心代码

核心代码的使用方式看章节四！

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch
__all__ = ['AOD_pono_net']
class AODnet(nn.Module):
    def __init__(self):
        super(AODnet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=1, stride=1, padding=0)
        self.conv2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=5, stride=1, padding=2)
        self.conv4 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=7, stride=1, padding=3)
        self.conv5 = nn.Conv2d(in_channels=12, out_channels=3, kernel_size=3, stride=1, padding=1)
        self.b = 1
    def forward(self, x):
        x1 = F.relu(self.conv1(x))
        x2 = F.relu(self.conv2(x1))
        cat1 = torch.cat((x1, x2), 1)
        x3 = F.relu(self.conv3(cat1))
        cat2 = torch.cat((x2, x3), 1)
        x4 = F.relu(self.conv4(cat2))
        cat3 = torch.cat((x1, x2, x3, x4), 1)
        k = F.relu(self.conv5(cat3))
        if k.size() != x.size():
            raise Exception("k, haze image are different size!")
        output = k * x - k + self.b
        return F.relu(output)
class AOD_pono_net(nn.Module):
    def __init__(self,S):
        super(AOD_pono_net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=1, stride=1, padding=0)
        self.conv2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=5, stride=1, padding=2)
        self.conv4 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=7, stride=1, padding=3)
        self.conv5 = nn.Conv2d(in_channels=12, out_channels=3, kernel_size=3, stride=1, padding=1)
        self.b = 1
        self.pono = PONO(affine=False)
        self.ms = MS()
    def forward(self, x):
        x1 = F.relu(self.conv1(x))
        x2 = F.relu(self.conv2(x1))
        cat1 = torch.cat((x1, x2), 1)
        x1, mean1, std1 = self.pono(x1)
        x2, mean2, std2 = self.pono(x2)
        x3 = F.relu(self.conv3(cat1))
        cat2 = torch.cat((x2, x3), 1)
        x3 = self.ms(x3, mean1, std1)
        x4 = F.relu(self.conv4(cat2))
        x4 = self.ms(x4, mean2, std2)
        cat3 = torch.cat((x1, x2, x3, x4), 1)
        k = F.relu(self.conv5(cat3))
        if k.size() != x.size():
            raise Exception("k, haze image are different size!")
        output = k * x - k + self.b
        output = F.relu(output)
        return output
class PONO(nn.Module):
    def __init__(self, input_size=None, return_stats=False, affine=True, eps=1e-5):
        super(PONO, self).__init__()
        self.return_stats = return_stats
        self.input_size = input_size
        self.eps = eps
        self.affine = affine
        if affine:
            self.beta = nn.Parameter(torch.zeros(1, 1, *input_size))
            self.gamma = nn.Parameter(torch.ones(1, 1, *input_size))
        else:
            self.beta, self.gamma = None, None
    def forward(self, x):
        mean = x.mean(dim=1, keepdim=True)
        std = (x.var(dim=1, keepdim=True) + self.eps).sqrt()
        x = (x - mean) / std
        if self.affine:
            x = x * self.gamma + self.beta
        return x, mean, std
class MS(nn.Module):
    def __init__(self, beta=None, gamma=None):
        super(MS, self).__init__()
        self.gamma, self.beta = gamma, beta
    def forward(self, x, beta=None, gamma=None):
        beta = self.beta if beta is None else beta
        gamma = self.gamma if gamma is None else gamma
        if gamma is not None:
            y = x.mul(gamma)  # 使用非原地操作mul
        else:
            y = x  # 如果不乘gamma，保持y不变
        if beta is not None:
            y = y.add(beta)  # 使用非原地操作add
        return y
if __name__ == "__main__":
    # Generating Sample image
    image_size = (1, 3, 640, 640)
    image = torch.rand(*image_size)
    out = AOD_pono_net()
    out = out(image)
    print(out.size())

四、添加教程

4.1 修改一

第一还是建立文件，我们找到如下 ultralytics /nn文件夹下建立一个目录名字呢就是'Addmodules'文件夹( 用群内的文件的话已经有了无需新建) ！然后在其内部建立一个新的py文件将核心代码复制粘贴进去即可。

4.2 修改二

第二步我们在该目录下创建一个新的py文件名字为'__init__.py'( 用群内的文件的话已经有了无需新建) ，然后在其内部导入我们的检测头如下图所示。

4.3 修改三

第三步我门中到如下文件'ultralytics/nn/tasks.py'进行导入和注册我们的模块( 用群内的文件的话已经有了无需重新导入直接开始第四步即可) ！

从今天开始以后的教程就都统一成这个样子了，因为我默认大家用了我群内的文件来进行修改！！

4.4 修改四

按照我的添加在parse_model里添加即可。

到此就修改完成了，大家可以复制下面的yaml文件运行。

五、AirNet的yaml文件和运行记录

5.1 AirNet的yaml文件

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLO11 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=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
# YOLO11n 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, 2, C3k2, [256, False, 0.25]]
  - [-1, 1, Conv, [256, 3, 2]] # 4-P3/8
  - [-1, 2, C3k2, [512, False, 0.25]]
  - [-1, 1, Conv, [512, 3, 2]] # 6-P4/16
  - [-1, 2, C3k2, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]] # 8-P5/32
  - [-1, 2, C3k2, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]] # 10
  - [-1, 2, C2PSA, [1024]] # 11
# YOLO11n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 7], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 14
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 5], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 17 (P3/8-small)
  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 14], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 20 (P4/16-medium)
  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 11], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 23 (P5/32-large)
  - [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.2 训练代码

大家可以创建一个py文件将我给的代码复制粘贴进去，配置好自己的文件路径即可运行。

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/v8/yolov8-C2f-FasterBlock.yaml')
    # model.load('yolov8n.pt') # loading pretrain weights
    model.train(data=r'替换数据集yaml文件地址',
                # 如果大家任务是其它的'ultralytics/cfg/default.yaml'找到这里修改task可以改成detect, segment, classify, pose
                cache=False,
                imgsz=640,
                epochs=150,
                single_cls=False,  # 是否是单类别检测
                batch=4,
                close_mosaic=10,
                workers=0,
                device='0',
                optimizer='SGD', # using SGD
                # resume='', # 如过想续训就设置last.pt的地址
                amp=False,  # 如果出现训练损失为Nan可以关闭amp
                project='runs/train',
                name='exp',
                )

5.3 AirNet的训练过程截图

五、本文总结

到此本文的正式分享内容就结束了，在这里给大家推荐我的YOLOv11改进有效涨点专栏，本专栏目前为新开的平均质量分98分，后期我会根据各种最新的前沿顶会进行论文复现，也会对一些老的改进机制进行补充，如果大家觉得本文帮助到你了，订阅本专栏，关注后续更多的更新~