devbox
一键难忘520
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

YOLOv9改进策略 | 低照度图像篇 | 轻量级的低照度图像增强网络IAT改进YOLOv9暗光检测（全网独家首发）_yolov9弱光照

作者：一键难忘520 | 2024-07-12 11:43:06

赞

踩

yolov9弱光照

  一、本文介绍

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

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

 专栏地址：YOLOv9有效涨点专栏-持续复现各种顶会内容-有效涨点-全网改进最全的专栏 

目录

 一、本文介绍

二、基本原理

2.1 IAT原理

2.2 IAT的核心模块

三、核心代码 

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

  4.1 修改一

4.2 修改二 

4.3 修改三 

4.4 修改四 

五、yaml文件和运行记录

5.1 yaml文件

5.2 训练过程截图 

五、本文总结

二、基本原理

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

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

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
  5.  
  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 修改一

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

 

4.2 修改二 

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

4.3 修改三 

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

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

​

4.4 修改四 

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

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

五、yaml文件和运行记录

5.1 yaml文件

主干和Neck全部用上该卷积轻量化到机制的yaml文件。

  1. # YOLOv9
  2.  
  3. # parameters
  4. nc: 80  # number of classes
  5. depth_multiple: 1  # model depth multiple
  6. width_multiple: 1  # layer channel multiple
  7. #activation: nn.LeakyReLU(0.1)
  8. #activation: nn.ReLU()
  9.  
  10. # anchors
  11. anchors: 3
  12.  
  13. # YOLOv9 backbone
  14. backbone:
  15.   [
  16.    [-1, 1, IAT, []],
  17.    [-1, 1, Silence, []],
  18.    # conv down
  19.    [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
  20.    # conv down
  21.    [-1, 1, Conv, [128, 3, 2]],  # 2-P2/4
  22.    # elan-1 block
  23.    [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]],  # 3
  24.    # conv down
  25.    [-1, 1, Conv, [256, 3, 2]],  # 4-P3/8
  26.    # elan-2 block
  27.    [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]],  # 5
  28.    # conv down
  29.    [-1, 1, Conv, [512, 3, 2]],  # 6-P4/16
  30.    # elan-2 block
  31.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 7
  32.    # conv down
  33.    [-1, 1, Conv, [512, 3, 2]],  # 8-P5/32
  34.    # elan-2 block
  35.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 9
  36.   ]
  37.  
  38. # YOLOv9 head
  39. head:
  40.   [
  41.    # elan-spp block
  42.    [-1, 1, SPPELAN, [512, 256]],  # 10
  43.  
  44.    # up-concat merge
  45.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  46.    [[-1, 8], 1, Concat, [1]],  # cat backbone P4
  47.  
  48.    # elan-2 block
  49.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 13
  50.  
  51.    # up-concat merge
  52.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  53.    [[-1, 6], 1, Concat, [1]],  # cat backbone P3
  54.  
  55.    # elan-2 block
  56.    [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]],  # 16 (P3/8-small)
  57.  
  58.    # conv-down merge
  59.    [-1, 1, Conv, [256, 3, 2]],
  60.    [[-1, 14], 1, Concat, [1]],  # cat head P4
  61.  
  62.    # elan-2 block
  63.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 19 (P4/16-medium)
  64.  
  65.    # conv-down merge
  66.    [-1, 1, Conv, [512, 3, 2]],
  67.    [[-1, 11], 1, Concat, [1]],  # cat head P5
  68.  
  69.    # elan-2 block
  70.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 22 (P5/32-large)
  71.    
  72.    # routing
  73.    [6, 1, CBLinear, [[256]]], # 23
  74.    [8, 1, CBLinear, [[256, 512]]], # 24
  75.    [10, 1, CBLinear, [[256, 512, 512]]], # 25
  76.    
  77.    # conv down
  78.    [0, 1, Conv, [64, 3, 2]],  # 26-P1/2
  79.  
  80.    # conv down
  81.    [-1, 1, Conv, [128, 3, 2]],  # 27-P2/4
  82.  
  83.    # elan-1 block
  84.    [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]],  # 28
  85.  
  86.    # conv down fuse
  87.    [-1, 1, Conv, [256, 3, 2]],  # 29-P3/8
  88.    [[24, 25, 26, -1], 1, CBFuse, [[0, 0, 0]]], # 30
  89.  
  90.    # elan-2 block
  91.    [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]],  # 31
  92.  
  93.    # conv down fuse
  94.    [-1, 1, Conv, [512, 3, 2]],  # 32-P4/16
  95.    [[25, 26, -1], 1, CBFuse, [[1, 1]]], # 33
  96.  
  97.    # elan-2 block
  98.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 34
  99.  
  100.    # conv down fuse
  101.    [-1, 1, Conv, [512, 3, 2]],  # 35-P5/32
  102.    [[26, -1], 1, CBFuse, [[2]]], # 36
  103.  
  104.    # elan-2 block
  105.    [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 37
  106.  
  107.    # detect
  108.    [[32, 35, 38, 17, 20, 23], 1, DualDDetect, [nc]],  # DualDDetect(A3, A4, A5, P3, P4, P5)
  109.   ]

5.2 训练过程截图 

五、本文总结

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

 专栏地址：YOLOv9有效涨点专栏-持续复现各种顶会内容-有效涨点-全网改进最全的专栏 

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/一键难忘520/article/detail/813975
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号