📚 课程目录第4章基于注意力机制的RAW去马赛克

4. 基于注意力机制的RAW去马赛克

任务目标

在RAW域去马赛克（Demosaicing）任务中，将拜耳阵列（Bayer Pattern）的RAW图恢复为全彩RGB图。本案例要求在基础U-Net去马赛克网络中引入通道注意力模块（Channel Attention），使网络能够自适应地重新校准各特征通道的权重，从而提升颜色恢复精度。

环境与依赖

Python 3.8+，PyTorch 1.10+，NumPy。数据集使用Kodak或McMaster标准测试图，按标准Bayer RGGB模式采样生成RAW输入。

核心实现

以下代码定义了一个轻量级去马赛克网络，包含编码器-解码器结构，并在跳跃连接后插入通道注意力模块（SE Block）。输入为4通道（RGGB）打包的RAW数据，输出为3通道RGB图像。

import torch
import torch.nn as nn
import torch.nn.functional as F

class ChannelAttention(nn.Module):
    """Squeeze-and-Excitation通道注意力模块"""
    def __init__(self, in_channels, reduction=8):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction, in_channels),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)

class DemosaicNet(nn.Module):
    """带通道注意力的U-Net去马赛克网络"""
    def __init__(self, in_channels=4, out_channels=3, features=32):
        super().__init__()
        # 编码器
        self.enc1 = DoubleConv(in_channels, features)
        self.enc2 = DoubleConv(features, features * 2)
        self.enc3 = DoubleConv(features * 2, features * 4)
        self.pool = nn.MaxPool2d(2)

        # 瓶颈
        self.bottleneck = DoubleConv(features * 4, features * 8)

        # 解码器
        self.up3 = nn.ConvTranspose2d(features * 8, features * 4, 2, stride=2)
        self.dec3 = DoubleConv(features * 8, features * 4)
        self.up2 = nn.ConvTranspose2d(features * 4, features * 2, 2, stride=2)
        self.dec2 = DoubleConv(features * 4, features * 2)
        self.up1 = nn.ConvTranspose2d(features * 2, features, 2, stride=2)
        self.dec1 = DoubleConv(features * 2, features)

        # 通道注意力（插入在解码器每个阶段之后）
        self.ca3 = ChannelAttention(features * 4)
        self.ca2 = ChannelAttention(features * 2)
        self.ca1 = ChannelAttention(features)

        self.out_conv = nn.Conv2d(features, out_channels, 1)

    def forward(self, x):
        # 编码
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))

        # 瓶颈
        b = self.bottleneck(self.pool(e3))

        # 解码 + 跳跃连接 + 通道注意力
        d3 = self.up3(b)
        d3 = torch.cat([d3, e3], dim=1)
        d3 = self.dec3(d3)
        d3 = self.ca3(d3)

        d2 = self.up2(d3)
        d2 = torch.cat([d2, e2], dim=1)
        d2 = self.dec2(d2)
        d2 = self.ca2(d2)

        d1 = self.up1(d2)
        d1 = torch.cat([d1, e1], dim=1)
        d1 = self.dec1(d1)
        d1 = self.ca1(d1)

        return self.out_conv(d1)

def pack_raw_bayer(rgb, pattern='RGGB'):
    """将RGB图像按Bayer模式打包为4通道输入"""
    # rgb: [B, 3, H, W] 假设H,W为偶数
    b, c, h, w = rgb.shape
    # 按RGGB排列提取
    r = rgb[:, 0:1, 0::2, 0::2]  # R在(0,0)
    g1 = rgb[:, 1:2, 0::2, 1::2]  # G在(0,1)
    g2 = rgb[:, 1:2, 1::2, 0::2]  # G在(1,0)
    b_ = rgb[:, 2:3, 1::2, 1::2]  # B在(1,1)
    return torch.cat([r, g1, g2, b_], dim=1)

def train_step(model, raw, target, optimizer, criterion):
    """单步训练"""
    model.train()
    optimizer.zero_grad()
    output = model(raw)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    return loss.item()

def psnr(img1, img2, max_val=1.0):
    """计算PSNR"""
    mse = F.mse_loss(img1, img2)
    return 10 * torch.log10(max_val**2 / mse)

训练与验证

# 数据准备（示意）
# rgb_gt: [B, 3, H, W] 范围[0,1]
raw_input = pack_raw_bayer(rgb_gt)  # [B, 4, H/2, W/2]

# 模型初始化
model = DemosaicNet(in_channels=4, out_channels=3, features=32)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.L1Loss()

# 训练循环（示意）
for epoch in range(50):
    for batch in dataloader:
        raw, target = batch
        loss = train_step(model, raw, target, optimizer, criterion)

# 验证
model.eval()
with torch.no_grad():
    output = model(raw_input)
    score = psnr(output, rgb_gt)
    print(f'PSNR: {score.item():.2f} dB')

验证方法与通过标准

数据准备：从Kodak数据集中选取至少10张图像，裁剪为256×256，按RGGB模式打包为4通道输入。
训练配置：batch size=8，学习率1e-3，训练50个epoch，使用Adam优化器与L1损失。
评估指标：在测试集上计算PSNR（峰值信噪比）与ΔE（色差）。
通过条件：与不含注意力模块的基线U-Net相比，PSNR提升至少0.3dB，且ΔE降低。若PSNR达到28dB以上（Kodak数据集），视为颜色恢复精度达标。

关键说明

通道注意力模块通过全局平均池化获取通道描述符，经两个全连接层学习通道间非线性关系，再用Sigmoid生成0~1的权重，对原始特征进行重标定。该机制使网络能抑制无关颜色通道的噪声响应，增强对边缘和纹理区域的颜色判别力。本实现将注意力模块置于解码器每个阶段之后，在融合跳跃连接特征后立即进行通道重标定，实验表明该位置比置于编码器或瓶颈处效果更优。