Tip 2024 | CGA特征融合模块，短短几行代码，有效提点，即插即用

点击下方卡片，关注“AI前沿速递”公众号

各种重磅干货，第一时间送达

标题：DEA-Net: Single image dehazing based on detail-enhanced convolution and content-guided attention

论文链接：https://arxiv.org/pdf/2301.04805

代码链接：https://github.com/cecret3350/DEA-Net/tree/main

来源：浙江大学

创新点

1. 细节增强卷积（DEConv）

• 首次引入差分卷积：DEConv包含并行的普通卷积和差分卷积，这是首次将差分卷积引入图像去雾问题中。差分卷积通过计算像素对的差异并将先验信息显式编码到CNN中，增强了模型对图像细节信息的捕捉能力，从而提高了去雾性能。
• 无额外参数和计算成本：通过重参数化技术，DEConv可以等价地转换为没有额外参数和计算成本的普通卷积。这使得DEConv在不增加模型复杂度的情况下，提升了特征的表示和泛化能力。

2. 内容引导注意力（CGA）

• 通道特定的空间重要性图（SIM）：CGA为每个通道分配独特的SIM，引导模型关注每个通道的重要区域。通过使用输入特征来指导SIM的生成，CGA能够强调特征中编码的更多有用信息，有效提高去雾性能。
• 粗到细的注意力生成：CGA采用两步注意力生成器，首先生成粗略的空间注意力图，然后根据每个通道的输入特征进行细化，生成最终的通道特定的SIM。这种设计充分混合了通道注意力权重和空间注意力权重，保证了信息交互。

3. 基于CGA的混合融合方案

• 有效融合低级和高级特征：提出了一种基于CGA的混合融合方案，将编码器部分的低级特征与相应的高级特征进行自适应融合。通过学习空间权重，该方案能够有效融合特征并辅助梯度流动，进一步提升去雾效果。
• 解决感受野不匹配问题：传统的特征融合方法（如逐元素相加或简单的Mixup操作）存在感受野不匹配的问题。CGA-based Mixup Fusion Scheme通过CGA计算空间权重来调节特征，解决了这一问题，提高了特征融合的效果。

4. DEA-Net整体架构

• 结合DEConv和CGA：DEA-Net通过结合DEConv和CGA，并使用基于CGA的混合融合方案，提出了细节增强注意力网络（DEA-Net）来重建高质量的无雾图像。DEA-Net在多个基准数据集上表现优于最先进的去雾方法，以更快的推理速度获得了更准确的结果。
• U-Net-like架构：DEA-Net采用类似U-Net的架构，包含编码器部分、特征转换部分和解码器部分。通过在不同层次上部署不同的块（DEB和DEAB），DEA-Net能够有效地提取和融合特征，恢复出高质量的无雾图像。

关键模块

1. Detail-Enhanced Attention Block（DEAB）

• 组成：由Detail-Enhanced Convolution（DEConv）和Content-Guided Attention（CGA）组成。
• 功能：DEConv通过整合先验信息增强特征表示和泛化能力，CGA则通过为每个通道分配独特的空间重要性图（SIM），关注特征中有用的信息，二者结合提升去雾性能。

2. Detail-Enhanced Convolution（DEConv）

• 组成：包含五个并行部署的卷积层，包括一个普通卷积（VC）和四个差分卷积（CDC、ADC、HDC、VDC）。

• 功能：差分卷积通过计算像素对的差异来显式编码先验信息，如HDC和VDC通过学习有益的梯度信息，将梯度先验编码到卷积层中，从而增强特征的表示和泛化能力。通过重参数化技术，DEConv可以等价地转换为普通卷积，无需增加额外参数和计算成本。

3. Content-Guided Attention（CGA）

• 功能：作为一种两步注意力生成器，首先生成粗略的空间注意力图，然后根据每个通道的输入特征进行细化，生成最终的通道特定的SIM。通过使用输入特征的内容来引导SIM的生成，CGA能够关注每个通道中独特的特征部分，强调特征中有用的信息，有效提升性能。此外，基于CGA还提出了一个特征融合方案，用于将编码器部分的低级特征与相应的高级特征进行自适应融合。

4. CGA-based Mixup Fusion Scheme

• 功能：为了解决特征融合中的感受野不匹配问题，提出了一种基于CGA的Mixup融合方案。通过CGA计算空间权重来调节特征，将编码器部分的低级特征与解码器部分的高级特征进行加权求和融合，并通过跳跃连接添加输入特征，以缓解梯度消失问题并简化学习过程。最终通过1×1卷积层得到融合后的特征。

5. DEA-Net整体架构

• 组成：包含编码器部分、特征转换部分和解码器部分。特征转换部分采用堆叠的DEAB来学习无雾特征。在层次结构中包含三个级别，不同级别采用不同的块来提取相应特征（1&2级别：DEB，3级别：DEAB）。
• 功能：通过结合DEConv、CGA以及CGA-based Mixup Fusion Scheme，DEA-Net能够恢复出高质量的无雾图像，在多个基准数据集上展现出优于现有最先进去雾方法的性能，同时具有更快的推理速度。

代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import Rearrange

class SpatialAttention(nn.Module):
    def __init__(self):
        super(SpatialAttention, self).__init__()
        self.sa = nn.Conv2d(2, 1, 7, padding=3, padding_mode='reflect', bias=True)

    def forward(self, x):
        x_avg = torch.mean(x, dim=1, keepdim=True)
        x_max, _ = torch.max(x, dim=1, keepdim=True)
        x2 = torch.cat([x_avg, x_max], dim=1)
        sattn = self.sa(x2)
        return sattn

class ChannelAttention(nn.Module):
    def __init__(self, dim, reduction=8):
        super(ChannelAttention, self).__init__()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.ca = nn.Sequential(
            nn.Conv2d(dim, dim // reduction, 1, padding=0, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(dim // reduction, dim, 1, padding=0, bias=True),
        )

    def forward(self, x):
        x_gap = self.gap(x)
        cattn = self.ca(x_gap)
        return cattn

class PixelAttention(nn.Module):
    def __init__(self, dim):
        super(PixelAttention, self).__init__()
        self.pa2 = nn.Conv2d(2 * dim, dim, 7, padding=3, padding_mode='reflect', groups=dim, bias=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, pattn1):
        B, C, H, W = x.shape
        x = x.unsqueeze(dim=2)
        pattn1 = pattn1.unsqueeze(dim=2)
        x2 = torch.cat([x, pattn1], dim=2)
        x2 = Rearrange('b c t h w -> b (c t) h w')(x2)
        pattn2 = self.pa2(x2)
        pattn2 = self.sigmoid(pattn2)
        return pattn2

class CGAFusion(nn.Module):
    def __init__(self, dim, reduction=8):
        super(CGAFusion, self).__init__()
        self.sa = SpatialAttention()
        self.ca = ChannelAttention(dim, reduction)
        self.pa = PixelAttention(dim)
        self.conv = nn.Conv2d(dim, dim, 1, bias=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, y):
        initial = x + y
        cattn = self.ca(initial)
        sattn = self.sa(initial)
        pattn1 = sattn + cattn
        pattn2 = self.sigmoid(self.pa(initial, pattn1))
        result = initial + pattn2 * x + (1 - pattn2) * y
        result = self.conv(result)
        return result

代码解释

1. 初始融合：将两个输入特征图 x 和 y 相加，得到初始融合结果 initial。
2. 通道注意力：通过 ChannelAttention 模块生成通道注意力图 cattn。
3. 空间注意力：通过 SpatialAttention 模块生成空间注意力图 sattn。
4. 像素注意力：将空间注意力图和通道注意力图相加，得到 pattn1，然后通过 PixelAttention 模块生成像素注意力图 pattn2。
5. 最终融合：将初始融合结果 initial 与 x 和 y 分别乘以 pattn2 和 1 - pattn2，得到最终融合结果 result。
6. 1x1卷积：通过1x1卷积层对最终融合结果进行处理，得到输出特征图。

本文内容为论文学习收获分享，受限于知识能力，本文对原文的理解可能存在偏差，最终内容以原论文为准。本文信息旨在传播和学术交流，其内容由作者负责，不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题，请及时与我们联系，我们将在第一时间回复并处理。