学习资源站

YOLOv11改进-特殊场景检测篇-轻量级的低照度图像增强网络IAT改进yolov11暗光检测（全网独家首发改进）

一、本文介绍

本文给大家带来的改进机制是 轻量级 的变换器 模型 ： Illumination Adaptive Transformer (IAT) ，用于 图像增强 和曝光校正。其基本原理是通过 分解图像信号处理器（ISP）管道到局部和全局图像组件 ，从而 恢复在低光或过/欠曝光条件下的正常光照sRGB图像 。具体来说，IAT使用注意力查询来表示和调整ISP相关参数，例如颜色校正、伽马校正。模型具有约90k参数和约0.004s的处理速度，能够在低光增强和曝光校正的基准数据集上持续实现优于最新技术（State-of-The-Art, SOTA）的性能，我们将其用于YOLOv11上来改进我们模型的暗光检测能力，同时本文的内容不影响其它的模块改进。

欢迎大家订阅我的专栏一起学习YOLO！

二、基本原理

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

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

2.1 IAT原理

本文提出了一个轻量级的 变换器 模型： Illumination Adaptive Transformer (IAT) ，用于图像增强和曝光校正。其基本原理是通过 分解图像信号处理器（ISP）管道到局部和全局图像组件 ，从而 恢复在低光或过/欠曝光条件下的正常光照sRGB图像 。具体来说，IAT使用注意力查询来表示和调整ISP相关参数，例如颜色校正、伽马校正。模型具有约90k参数和约0.004s的处理速度，能够在低光增强和曝光校正的基准数据集上持续实现优于最新技术（State-of-The-Art, SOTA）的性能。

Illumination Adaptive Transformer (IAT)的 基本原理 如下：

1. 轻量级变换器架构： IAT设计为一个轻量级的模型，具有大约90,000个参数，专注于图像增强和曝光校正任务。这使得它在处理速度和资源消耗上非常高效，适用于实时或资源受限的应用场景。

2. 图像信号处理器（ISP）管道分解： IAT的核心原理是模拟并改进传统的ISP管道。通过分解ISP处理过程中的局部和全局图像成分，IAT能够针对特定的光照条件调整图像的视觉表现。

3. 适应性光照调整： IAT能够根据输入图像的光照条件动态调整处理策略，有效地处理低光、过曝光和欠曝光等情况，恢复正常光照下的sRGB图像。

下面为大家展示 Illumination Adaptive Transformer (IAT)的结构 分为两个主要部分：局部分支和全局分支。

1. 局部分支 (Local Branch)： 处理图像的局部特征。这一分支通过多次使用参数增强模块（PEM）来提取局部特征，并通过卷积层来进一步处理这些特征。

2. 全局分支 (Global Branch)： 处理图像的全局信息。它同样包含多个PEM和卷积层，不过处理的是全局图像内容。

3. 参数生成 (黑色线条)： 黑色线条表示参数生成路径，即如何通过网络生成ISP管道中需要的参数，如颜色矩阵和伽马值。

4. 图像处理 (黄色线条)： 黄色线条表示实际的图像处理路径。图像经过局部和全局分支的处理后，获得的特征会被用于调整图像的颜色和曝光。

5. 交叉注意力 (Cross Attention) ：这一 组件 在全局分支中，负责整合局部和全局分支的信息，以更准确地调整颜色矩阵和伽马值。

6. 最终输出 ：处理过的图像特征通过一个重塑操作和卷积层的处理，将局部和全局的调整应用到原始输入图像上，最终输出增强后的图像。

2.2 IAT的核心模块

下面这张图为大家直观地展示了Illumination Adaptive Transformer (IAT)中的两个核心模块： 像素级增强模块（Pixel-wise Enhancement Module, PEM） 和 全局预测模块（Global Prediction Module, GPM） 。

（a）像素级增强模块（PEM）:
输入: 大小为 B × C × H × W 的特征图，其中 B 表示批次大小， C 表示通道数， H × W 表示特征图的高和宽。
流程:
1. 通过一系列的1x1卷积层，对特征图进行逐点的线性变换，以增强或调整特定像素点的特性。
2. 每个1x1卷积层之后，进行元素级的相乘（表示为黄色的圆圈和相乘符号）。
3. 操作结束后，特征图被重塑成原始的 B × C × H × W 形状。

（b）全局预测模块（GPM）:
流程:
1. 特征图首先经过一个全连接层（FC），产生 V ，代表全局信息的值向量。
2. 另一个全连接层生成 K ，代表键向量。
3. K 和 V 通过交叉注意力机制与查询 Q 相结合，查询 Q 通常来自于局部特征。
4. 结果通过重塑操作，形成颜色校正矩阵和伽马校正值。

两个模块共同工作， PEM负责增强局部特征细节 ，而 GPM则负责生成全局调整参数 ，两者合作为图像增强提供更精细的控制。通过这种方法，IAT能够在处理不同光照条件下的图像时提供细腻的调整，实现出色的图像增强效果。

三、核心代码

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

  1. import math
  2. import torch
  3. import torch.nn as nn
  4. from timm.models.layers import trunc_normal_, DropPath, to_2tuple
  6. __all__ = ['IAT']
  8. class query_Attention(nn.Module):
  9.     def __init__(self, dim, num_heads=2, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
  10.         super().__init__()
  11.         self.num_heads = num_heads
  12.         head_dim = dim // num_heads
  13.         # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
  14.         self.scale = qk_scale or head_dim ** -0.5
  16.         self.q = nn.Parameter(torch.ones((1, 10, dim)), requires_grad=True)
  17.         self.k = nn.Linear(dim, dim, bias=qkv_bias)
  18.         self.v = nn.Linear(dim, dim, bias=qkv_bias)
  19.         self.attn_drop = nn.Dropout(attn_drop)
  20.         self.proj = nn.Linear(dim, dim)
  21.         self.proj_drop = nn.Dropout(proj_drop)
  23.     def forward(self, x):
  24.         B, N, C = x.shape
  25.         k = self.k(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
  26.         v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
  28.         q = self.q.expand(B, -1, -1).view(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
  29.         attn = (q @ k.transpose(-2, -1)) * self.scale
  30.         attn = attn.softmax(dim=-1)
  31.         attn = self.attn_drop(attn)
  33.         x = (attn @ v).transpose(1, 2).reshape(B, 10, C)
  34.         x = self.proj(x)
  35.         x = self.proj_drop(x)
  36.         return x
  39. class query_SABlock(nn.Module):
  40.     def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
  41.                  drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
  42.         super().__init__()
  43.         self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
  44.         self.norm1 = norm_layer(dim)
  45.         self.attn = query_Attention(
  46.             dim,
  47.             num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
  48.             attn_drop=attn_drop, proj_drop=drop)
  49.         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
  50.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  51.         self.norm2 = norm_layer(dim)
  52.         mlp_hidden_dim = int(dim * mlp_ratio)
  53.         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
  55.     def forward(self, x):
  56.         x = x + self.pos_embed(x)
  57.         x = x.flatten(2).transpose(1, 2)
  58.         x = self.drop_path(self.attn(self.norm1(x)))
  59.         x = x + self.drop_path(self.mlp(self.norm2(x)))
  60.         return x
  63. class conv_embedding(nn.Module):
  64.     def __init__(self, in_channels, out_channels):
  65.         super(conv_embedding, self).__init__()
  66.         self.proj = nn.Sequential(
  67.             nn.Conv2d(in_channels, out_channels // 2, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
  68.             nn.BatchNorm2d(out_channels // 2),
  69.             nn.GELU(),
  70.             # nn.Conv2d(out_channels // 2, out_channels // 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
  71.             # nn.BatchNorm2d(out_channels // 2),
  72.             # nn.GELU(),
  73.             nn.Conv2d(out_channels // 2, out_channels, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
  74.             nn.BatchNorm2d(out_channels),
  75.         )
  77.     def forward(self, x):
  78.         x = self.proj(x)
  79.         return x
  82. class Global_pred(nn.Module):
  83.     def __init__(self, in_channels=3, out_channels=64, num_heads=4, type='exp'):
  84.         super(Global_pred, self).__init__()
  85.         if type == 'exp':
  86.             self.gamma_base = nn.Parameter(torch.ones((1)), requires_grad=False) # False in exposure correction
  87.         else:
  88.             self.gamma_base = nn.Parameter(torch.ones((1)), requires_grad=True)
  89.         self.color_base = nn.Parameter(torch.eye((3)), requires_grad=True)  # basic color matrix
  90.         # main blocks
  91.         self.conv_large = conv_embedding(in_channels, out_channels)
  92.         self.generator = query_SABlock(dim=out_channels, num_heads=num_heads)
  93.         self.gamma_linear = nn.Linear(out_channels, 1)
  94.         self.color_linear = nn.Linear(out_channels, 1)
  96.         self.apply(self._init_weights)
  98.         for name, p in self.named_parameters():
  99.             if name == 'generator.attn.v.weight':
  100.                 nn.init.constant_(p, 0)
  102.     def _init_weights(self, m):
  103.         if isinstance(m, nn.Linear):
  104.             trunc_normal_(m.weight, std=.02)
  105.             if isinstance(m, nn.Linear) and m.bias is not None:
  106.                 nn.init.constant_(m.bias, 0)
  107.         elif isinstance(m, nn.LayerNorm):
  108.             nn.init.constant_(m.bias, 0)
  109.             nn.init.constant_(m.weight, 1.0)
  112.     def forward(self, x):
  113.         #print(self.gamma_base)
  114.         x = self.conv_large(x)
  115.         x = self.generator(x)
  116.         gamma, color = x[:, 0].unsqueeze(1), x[:, 1:]
  117.         gamma = self.gamma_linear(gamma).squeeze(-1) + self.gamma_base
  118.         #print(self.gamma_base, self.gamma_linear(gamma))
  119.         color = self.color_linear(color).squeeze(-1).view(-1, 3, 3) + self.color_base
  120.         return gamma, color
  123. # ResMLP's normalization
  124. class Aff(nn.Module):
  125.     def __init__(self, dim):
  126.         super().__init__()
  127.         # learnable
  128.         self.alpha = nn.Parameter(torch.ones([1, 1, dim]))
  129.         self.beta = nn.Parameter(torch.zeros([1, 1, dim]))
  131.     def forward(self, x):
  132.         x = x * self.alpha + self.beta
  133.         return x
  135. # Color Normalization
  136. class Aff_channel(nn.Module):
  137.     def __init__(self, dim, channel_first = True):
  138.         super().__init__()
  139.         # learnable
  140.         self.alpha = nn.Parameter(torch.ones([1, 1, dim]))
  141.         self.beta = nn.Parameter(torch.zeros([1, 1, dim]))
  142.         self.color = nn.Parameter(torch.eye(dim))
  143.         self.channel_first = channel_first
  145.     def forward(self, x):
  146.         if self.channel_first:
  147.             x1 = torch.tensordot(x, self.color, dims=[[-1], [-1]])
  148.             x2 = x1 * self.alpha + self.beta
  149.         else:
  150.             x1 = x * self.alpha + self.beta
  151.             x2 = torch.tensordot(x1, self.color, dims=[[-1], [-1]])
  152.         return x2
  154. class Mlp(nn.Module):
  155.     # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
  156.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
  157.         super().__init__()
  158.         out_features = out_features or in_features
  159.         hidden_features = hidden_features or in_features
  160.         self.fc1 = nn.Linear(in_features, hidden_features)
  161.         self.act = act_layer()
  162.         self.fc2 = nn.Linear(hidden_features, out_features)
  163.         self.drop = nn.Dropout(drop)
  165.     def forward(self, x):
  166.         x = self.fc1(x)
  167.         x = self.act(x)
  168.         x = self.drop(x)
  169.         x = self.fc2(x)
  170.         x = self.drop(x)
  171.         return x
  173. class CMlp(nn.Module):
  174.     # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
  175.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
  176.         super().__init__()
  177.         out_features = out_features or in_features
  178.         hidden_features = hidden_features or in_features
  179.         self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
  180.         self.act = act_layer()
  181.         self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
  182.         self.drop = nn.Dropout(drop)
  184.     def forward(self, x):
  185.         x = self.fc1(x)
  186.         x = self.act(x)
  187.         x = self.drop(x)
  188.         x = self.fc2(x)
  189.         x = self.drop(x)
  190.         return x
  192. class CBlock_ln(nn.Module):
  193.     def __init__(self, dim, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
  194.                  drop_path=0., act_layer=nn.GELU, norm_layer=Aff_channel, init_values=1e-4):
  195.         super().__init__()
  196.         self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
  197.         #self.norm1 = Aff_channel(dim)
  198.         self.norm1 = norm_layer(dim)
  199.         self.conv1 = nn.Conv2d(dim, dim, 1)
  200.         self.conv2 = nn.Conv2d(dim, dim, 1)
  201.         self.attn = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
  202.         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
  203.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  204.         #self.norm2 = Aff_channel(dim)
  205.         self.norm2 = norm_layer(dim)
  206.         mlp_hidden_dim = int(dim * mlp_ratio)
  207.         self.gamma_1 = nn.Parameter(init_values * torch.ones((1, dim, 1, 1)), requires_grad=True)
  208.         self.gamma_2 = nn.Parameter(init_values * torch.ones((1, dim, 1, 1)), requires_grad=True)
  209.         self.mlp = CMlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
  211.     def forward(self, x):
  212.         x = x + self.pos_embed(x)
  213.         B, C, H, W = x.shape
  214.         #print(x.shape)
  215.         norm_x = x.flatten(2).transpose(1, 2)
  216.         #print(norm_x.shape)
  217.         norm_x = self.norm1(norm_x)
  218.         norm_x = norm_x.view(B, H, W, C).permute(0, 3, 1, 2)
  221.         x = x + self.drop_path(self.gamma_1*self.conv2(self.attn(self.conv1(norm_x))))
  222.         norm_x = x.flatten(2).transpose(1, 2)
  223.         norm_x = self.norm2(norm_x)
  224.         norm_x = norm_x.view(B, H, W, C).permute(0, 3, 1, 2)
  225.         x = x + self.drop_path(self.gamma_2*self.mlp(norm_x))
  226.         return x
  229. def window_partition(x, window_size):
  230.     """
  231.     Args:
  232.         x: (B, H, W, C)
  233.         window_size (int): window size
  234.     Returns:
  235.         windows: (num_windows*B, window_size, window_size, C)
  236.     """
  237.     B, H, W, C = x.shape
  238.     #print(x.shape)
  239.     x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
  240.     windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
  241.     return windows
  244. def window_reverse(windows, window_size, H, W):
  245.     """
  246.     Args:
  247.         windows: (num_windows*B, window_size, window_size, C)
  248.         window_size (int): Window size
  249.         H (int): Height of image
  250.         W (int): Width of image
  251.     Returns:
  252.         x: (B, H, W, C)
  253.     """
  254.     B = int(windows.shape[0] / (H * W / window_size / window_size))
  255.     x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
  256.     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
  257.     return x
  260. class WindowAttention(nn.Module):
  261.     r""" Window based multi-head self attention (W-MSA) module with relative position bias.
  262.     It supports both of shifted and non-shifted window.
  263.     Args:
  264.         dim (int): Number of input channels.
  265.         window_size (tuple[int]): The height and width of the window.
  266.         num_heads (int): Number of attention heads.
  267.         qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
  268.         qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
  269.         attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
  270.         proj_drop (float, optional): Dropout ratio of output. Default: 0.0
  271.     """
  273.     def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
  274.         super().__init__()
  275.         self.dim = dim
  276.         self.window_size = window_size  # Wh, Ww
  277.         self.num_heads = num_heads
  278.         head_dim = dim // num_heads
  279.         self.scale = qk_scale or head_dim ** -0.5
  281.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
  282.         self.attn_drop = nn.Dropout(attn_drop)
  283.         self.proj = nn.Linear(dim, dim)
  284.         self.proj_drop = nn.Dropout(proj_drop)
  286.         self.softmax = nn.Softmax(dim=-1)
  288.     def forward(self, x):
  289.         B_, N, C = x.shape
  290.         qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
  291.         q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
  293.         q = q * self.scale
  294.         attn = (q @ k.transpose(-2, -1))
  296.         attn = self.softmax(attn)
  298.         attn = self.attn_drop(attn)
  300.         x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
  301.         x = self.proj(x)
  302.         x = self.proj_drop(x)
  303.         return x
  305. ## Layer_norm, Aff_norm, Aff_channel_norm
  306. class SwinTransformerBlock(nn.Module):
  307.     r""" Swin Transformer Block.
  308.     Args:
  309.         dim (int): Number of input channels.
  310.         input_resolution (tuple[int]): Input resulotion.
  311.         num_heads (int): Number of attention heads.
  312.         window_size (int): Window size.
  313.         shift_size (int): Shift size for SW-MSA.
  314.         mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
  315.         qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
  316.         qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
  317.         drop (float, optional): Dropout rate. Default: 0.0
  318.         attn_drop (float, optional): Attention dropout rate. Default: 0.0
  319.         drop_path (float, optional): Stochastic depth rate. Default: 0.0
  320.         act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
  321.         norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
  322.     """
  324.     def __init__(self, dim, num_heads=2, window_size=8, shift_size=0,
  325.                  mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
  326.                  act_layer=nn.GELU, norm_layer=Aff_channel):
  327.         super().__init__()
  328.         self.dim = dim
  329.         self.num_heads = num_heads
  330.         self.window_size = window_size
  331.         self.shift_size = shift_size
  332.         self.mlp_ratio = mlp_ratio
  334.         self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
  335.         #self.norm1 = norm_layer(dim)
  336.         self.norm1 = norm_layer(dim)
  337.         self.attn = WindowAttention(
  338.             dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
  339.             qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
  341.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  342.         #self.norm2 = norm_layer(dim)
  343.         self.norm2 = norm_layer(dim)
  344.         mlp_hidden_dim = int(dim * mlp_ratio)
  345.         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
  347.     def forward(self, x):
  348.         x = x + self.pos_embed(x)
  349.         B, C, H, W = x.shape
  350.         x = x.flatten(2).transpose(1, 2)
  352.         shortcut = x
  353.         x = self.norm1(x)
  354.         x = x.view(B, H, W, C)
  356.         # cyclic shift
  357.         if self.shift_size > 0:
  358.             shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
  359.         else:
  360.             shifted_x = x
  362.         # partition windows
  363.         x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
  364.         x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
  366.         # W-MSA/SW-MSA
  367.         attn_windows = self.attn(x_windows)  # nW*B, window_size*window_size, C
  369.         # merge windows
  370.         attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
  371.         shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
  373.         x = shifted_x
  374.         x = x.view(B, H * W, C)
  376.         # FFN
  377.         x = shortcut + self.drop_path(x)
  378.         x = x + self.drop_path(self.mlp(self.norm2(x)))
  379.         x = x.transpose(1, 2).reshape(B, C, H, W)
  381.         return x
  384. class Local_pred(nn.Module):
  385.     def __init__(self, dim=16, number=4, type='ccc'):
  386.         super(Local_pred, self).__init__()
  387.         # initial convolution
  388.         self.conv1 = nn.Conv2d(3, dim, 3, padding=1, groups=1)
  389.         self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
  390.         # main blocks
  391.         block = CBlock_ln(dim)
  392.         block_t = SwinTransformerBlock(dim)  # head number
  393.         if type == 'ccc':
  394.             # blocks1, blocks2 = [block for _ in range(number)], [block for _ in range(number)]
  395.             blocks1 = [CBlock_ln(16, drop_path=0.01), CBlock_ln(16, drop_path=0.05), CBlock_ln(16, drop_path=0.1)]
  396.             blocks2 = [CBlock_ln(16, drop_path=0.01), CBlock_ln(16, drop_path=0.05), CBlock_ln(16, drop_path=0.1)]
  397.         elif type == 'ttt':
  398.             blocks1, blocks2 = [block_t for _ in range(number)], [block_t for _ in range(number)]
  399.         elif type == 'cct':
  400.             blocks1, blocks2 = [block, block, block_t], [block, block, block_t]
  401.         #    block1 = [CBlock_ln(16), nn.Conv2d(16,24,3,1,1)]
  402.         self.mul_blocks = nn.Sequential(*blocks1, nn.Conv2d(dim, 3, 3, 1, 1), nn.ReLU())
  403.         self.add_blocks = nn.Sequential(*blocks2, nn.Conv2d(dim, 3, 3, 1, 1), nn.Tanh())
  405.     def forward(self, img):
  406.         img1 = self.relu(self.conv1(img))
  407.         mul = self.mul_blocks(img1)
  408.         add = self.add_blocks(img1)
  410.         return mul, add
  413. # Short Cut Connection on Final Layer
  414. class Local_pred_S(nn.Module):
  415.     def __init__(self, in_dim=3, dim=16, number=4, type='ccc'):
  416.         super(Local_pred_S, self).__init__()
  417.         # initial convolution
  418.         self.conv1 = nn.Conv2d(in_dim, dim, 3, padding=1, groups=1)
  419.         self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
  420.         # main blocks
  421.         block = CBlock_ln(dim)
  422.         block_t = SwinTransformerBlock(dim)  # head number
  423.         if type == 'ccc':
  424.             blocks1 = [CBlock_ln(16, drop_path=0.01), CBlock_ln(16, drop_path=0.05), CBlock_ln(16, drop_path=0.1)]
  425.             blocks2 = [CBlock_ln(16, drop_path=0.01), CBlock_ln(16, drop_path=0.05), CBlock_ln(16, drop_path=0.1)]
  426.         elif type == 'ttt':
  427.             blocks1, blocks2 = [block_t for _ in range(number)], [block_t for _ in range(number)]
  428.         elif type == 'cct':
  429.             blocks1, blocks2 = [block, block, block_t], [block, block, block_t]
  430.         #    block1 = [CBlock_ln(16), nn.Conv2d(16,24,3,1,1)]
  431.         self.mul_blocks = nn.Sequential(*blocks1)
  432.         self.add_blocks = nn.Sequential(*blocks2)
  434.         self.mul_end = nn.Sequential(nn.Conv2d(dim, 3, 3, 1, 1), nn.ReLU())
  435.         self.add_end = nn.Sequential(nn.Conv2d(dim, 3, 3, 1, 1), nn.Tanh())
  436.         self.apply(self._init_weights)
  438.     def _init_weights(self, m):
  439.         if isinstance(m, nn.Linear):
  440.             trunc_normal_(m.weight, std=.02)
  441.             if isinstance(m, nn.Linear) and m.bias is not None:
  442.                 nn.init.constant_(m.bias, 0)
  443.         elif isinstance(m, nn.LayerNorm):
  444.             nn.init.constant_(m.bias, 0)
  445.             nn.init.constant_(m.weight, 1.0)
  446.         elif isinstance(m, nn.Conv2d):
  447.             fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
  448.             fan_out //= m.groups
  449.             m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
  450.             if m.bias is not None:
  451.                 m.bias.data.zero_()
  453.     def forward(self, img):
  454.         img1 = self.relu(self.conv1(img))
  455.         # short cut connection
  456.         mul = self.mul_blocks(img1) + img1
  457.         add = self.add_blocks(img1) + img1
  458.         mul = self.mul_end(mul)
  459.         add = self.add_end(add)
  461.         return mul, add
  464. class IAT(nn.Module):
  465.     def __init__(self, in_dim=3, with_global=True, type='lol'):
  466.         super(IAT, self).__init__()
  467.         # self.local_net = Local_pred()
  469.         self.local_net = Local_pred_S(in_dim=in_dim)
  471.         self.with_global = with_global
  472.         if self.with_global:
  473.             self.global_net = Global_pred(in_channels=in_dim, type=type)
  475.     def apply_color(self, image, ccm):
  476.         shape = image.shape
  477.         image = image.view(-1, 3)
  478.         image = torch.tensordot(image, ccm, dims=[[-1], [-1]])
  479.         image = image.view(shape)
  480.         return torch.clamp(image, 1e-8, 1.0)
  482.     def forward(self, img_low):
  483.         # print(self.with_global)
  484.         mul, add = self.local_net(img_low)
  485.         img_high = (img_low.mul(mul)).add(add)
  487.         if not self.with_global:
  488.             return img_high
  490.         else:
  491.             gamma, color = self.global_net(img_low)
  492.             b = img_high.shape[0]
  493.             img_high = img_high.permute(0, 2, 3, 1)  # (B,C,H,W) -- (B,H,W,C)
  494.             img_high = torch.stack(
  495.                 [self.apply_color(img_high[i, :, :, :], color[i, :, :]) ** gamma[i, :] for i in range(b)], dim=0)
  496.             img_high = img_high.permute(0, 3, 1, 2)  # (B,H,W,C) -- (B,C,H,W)
  497.             return img_high
  500. if __name__ == "__main__":
  501.     img = torch.Tensor(1, 3, 640, 640)
  502.     net = IAT()
  503.     imghigh = net(img)
  504.     print(imghigh.size())
  505.     print('total parameters:', sum(param.numel() for param in net.parameters()))
  506.     _, _, high = net(img)

四、手把手教你添加IAT低照度图像增强网络

4.1 修改一

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

​

4.2 修改二

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

4.3 修改三

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

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

​

4.4 修改四

按照我的添加在parse_model里添加即可，红框内的添加即可，没有的都是其它文章里的改进机制。

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

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

5.1 IAT的yaml文件

此版本的训练信息：YOLO11-IAT summary: 435 layers, 2,681,969 parameters, 2,681,953 gradients, 24.4 GFLOPs 

  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  4. # Parameters
  5. nc: 80 # number of classes
  6. scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  7.   # [depth, width, max_channels]
  8.   n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  9.   s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  10.   m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  11.   l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  12.   x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
  14. # YOLO11n backbone
  15. backbone:
  16.   # [from, repeats, module, args]
  17.   - [-1, 1, IAT, []] # 0-P1/2
  18.   - [-1, 1, Conv, [64, 3, 2]] # 1-P1/2
  19.   - [-1, 1, Conv, [128, 3, 2]] # 2-P2/4
  20.   - [-1, 2, C3k2, [256, False, 0.25]]
  21.   - [-1, 1, Conv, [256, 3, 2]] # 4-P3/8
  22.   - [-1, 2, C3k2, [512, False, 0.25]]
  23.   - [-1, 1, Conv, [512, 3, 2]] # 6-P4/16
  24.   - [-1, 2, C3k2, [512, True]]
  25.   - [-1, 1, Conv, [1024, 3, 2]] # 8-P5/32
  26.   - [-1, 2, C3k2, [1024, True]]
  27.   - [-1, 1, SPPF, [1024, 5]] # 10
  28.   - [-1, 2, C2PSA, [1024]] # 11
  30. # YOLO11n head
  31. head:
  32.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  33.   - [[-1, 7], 1, Concat, [1]] # cat backbone P4
  34.   - [-1, 2, C3k2, [512, False]] # 14
  36.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  37.   - [[-1, 5], 1, Concat, [1]] # cat backbone P3
  38.   - [-1, 2, C3k2, [256, False]] # 17 (P3/8-small)
  40.   - [-1, 1, Conv, [256, 3, 2]]
  41.   - [[-1, 14], 1, Concat, [1]] # cat head P4
  42.   - [-1, 2, C3k2, [512, False]] # 20 (P4/16-medium)
  44.   - [-1, 1, Conv, [512, 3, 2]]
  45.   - [[-1, 11], 1, Concat, [1]] # cat head P5
  46.   - [-1, 2, C3k2, [1024, True]] # 23 (P5/32-large)
  48.   - [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.2 训练代码

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

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

5.3 IAT的训练过程截图

五、本文总结

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