3. 多尺度特征融合的RAW去马赛克

任务说明

本案例实现一个轻量级多尺度特征融合网络，将单通道RAW Bayer图像转换为三通道RGB图像，重点抑制去马赛克过程中产生的伪彩色伪影。网络采用双分支结构：分支1使用不同膨胀率的卷积捕获多尺度上下文，分支2保留高分辨率细节，最终通过特征融合与亚像素卷积完成上采样。

环境依赖

需要 Python 3.8+、PyTorch 1.10+、torchvision（仅用于数据集加载）、numpy、opencv-python（用于图像读写）。若使用GPU训练，需安装对应CUDA版本的PyTorch。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# ---------- Bayer通道分离（RGGB） ----------
def bayer_to_channels(raw):
    """输入: (B,1,H,W)  Bayer RAW (RGGB)
       输出: 4通道, 每个通道对应R/G/G/B位置下采样后的图像"""
    B, C, H, W = raw.shape
    assert C == 1
    # 确保尺寸为偶数
    H_even, W_even = H - H % 2, W - W % 2
    raw = raw[:, :, :H_even, :W_even]
    R = raw[:, :, 0::2, 0::2]   # (B,1,H/2,W/2)
    G1 = raw[:, :, 0::2, 1::2]
    G2 = raw[:, :, 1::2, 0::2]
    B = raw[:, :, 1::2, 1::2]
    return torch.cat([R, G1, G2, B], dim=1)  # (B,4,H/2,W/2)

# ---------- 多尺度特征融合模块 ----------
class MSFF(nn.Module):
    """多尺度特征融合: 3个膨胀率分支 + 1个细节分支"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        # 膨胀卷积分支 (感受野递增)
        self.branch1 = nn.Conv2d(in_ch, out_ch, 3, padding=1, dilation=1)
        self.branch2 = nn.Conv2d(in_ch, out_ch, 3, padding=2, dilation=2)
        self.branch3 = nn.Conv2d(in_ch, out_ch, 3, padding=3, dilation=3)
        # 细节分支 (1x1卷积保留局部信息)
        self.detail = nn.Conv2d(in_ch, out_ch, 1)
        # 融合权重 (可学习)
        self.fuse_weights = nn.Parameter(torch.ones(4) / 4.0)

    def forward(self, x):
        b1 = F.relu(self.branch1(x))
        b2 = F.relu(self.branch2(x))
        b3 = F.relu(self.branch3(x))
        det = F.relu(self.detail(x))
        # 加权融合
        w = F.softmax(self.fuse_weights, dim=0)
        return w[0]*b1 + w[1]*b2 + w[2]*b3 + w[3]*det

# ---------- 去马赛克网络 ----------
class DemosaicNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 输入: 4通道 (R,G1,G2,B 下采样图)
        self.head = nn.Conv2d(4, 32, 3, padding=1)
        # 多尺度特征提取 (3层)
        self.msff1 = MSFF(32, 64)
        self.msff2 = MSFF(64, 64)
        self.msff3 = MSFF(64, 32)
        # 亚像素上采样: 将 (B,32,H/2,W/2) -> (B,8,H,W) -> (B,3,H,W)
        self.up_conv = nn.Conv2d(32, 8*3, 3, padding=1)
        self.pixel_shuffle = nn.PixelShuffle(2)  # 输出 (B,3,H,W)

    def forward(self, raw):
        # raw: (B,1,H,W)
        x = bayer_to_channels(raw)          # (B,4,H/2,W/2)
        x = F.relu(self.head(x))
        x = self.msff1(x)
        x = self.msff2(x)
        x = self.msff3(x)
        x = self.up_conv(x)                 # (B,24,H/2,W/2)
        rgb = self.pixel_shuffle(x)         # (B,3,H,W)
        return rgb

# ---------- 训练/验证辅助 ----------
def rgb_to_bayer(rgb):
    """将RGB图像转换为RGGB Bayer RAW (用于生成训练数据)"""
    B, C, H, W = rgb.shape
    raw = torch.zeros(B, 1, H, W, device=rgb.device)
    raw[:, 0, 0::2, 0::2] = rgb[:, 0, 0::2, 0::2]  # R
    raw[:, 0, 0::2, 1::2] = rgb[:, 1, 0::2, 1::2]  # G1
    raw[:, 0, 1::2, 0::2] = rgb[:, 1, 1::2, 0::2]  # G2
    raw[:, 0, 1::2, 1::2] = rgb[:, 2, 1::2, 1::2]  # B
    return raw

def psnr(pred, target, max_val=1.0):
    """计算PSNR (峰值信噪比)"""
    mse = F.mse_loss(pred, target)
    return 10 * torch.log10(max_val**2 / mse)

# ---------- 最小训练闭环 ----------
def train_demo():
    # 生成模拟数据: 8张16x16随机彩色图像
    torch.manual_seed(42)
    B, H, W = 8, 16, 16
    rgb_gt = torch.rand(B, 3, H, W)  # 真实RGB
    raw = rgb_to_bayer(rgb_gt)        # 模拟RAW

    model = DemosaicNet()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.L1Loss()

    model.train()
    for epoch in range(5):
        optimizer.zero_grad()
        rgb_pred = model(raw)
        # 损失: L1 + 轻微颜色一致性正则 (抑制伪彩色)
        loss = loss_fn(rgb_pred, rgb_gt)
        # 颜色差异惩罚: 计算RGB通道间的梯度差异
        grad_r = torch.abs(rgb_pred[:, 0:1, :, :-1] - rgb_pred[:, 0:1, :, 1:])
        grad_g = torch.abs(rgb_pred[:, 1:2, :, :-1] - rgb_pred[:, 1:2, :, 1:])
        grad_b = torch.abs(rgb_pred[:, 2:3, :, :-1] - rgb_pred[:, 2:3, :, 1:])
        color_loss = (torch.abs(grad_r - grad_g).mean() +
                      torch.abs(grad_g - grad_b).mean())
        total_loss = loss + 0.1 * color_loss
        total_loss.backward()
        optimizer.step()
        if epoch % 2 == 0:
            print(f"Epoch {epoch}: Loss = {total_loss.item():.4f}")

    # 验证
    model.eval()
    with torch.no_grad():
        pred = model(raw)
        p = psnr(pred, rgb_gt)
        print(f"训练后PSNR: {p.item():.2f} dB")
        # 检查伪彩色: 计算色差 (应较小)
        color_diff = torch.abs(pred[:, 0] - pred[:, 1]).mean() + \
                     torch.abs(pred[:, 1] - pred[:, 2]).mean()
        print(f"平均通道差异 (越小伪彩色越少): {color_diff.item():.4f}")

if __name__ == "__main__":
    train_demo()

运行方式

将上述代码保存为 demosaic_msff.py，直接运行：

python demosaic_msff.py

代码会使用随机生成的8张16×16模拟RAW图像进行5轮训练，并输出训练损失与验证指标。

验证方法与通过标准

训练收敛：训练过程中Loss应逐轮下降，最终Loss小于0.1。
重建质量：训练后PSNR应高于20dB（随机初始化时通常低于15dB），表明网络学到了有效的去马赛克映射。
伪彩色抑制：平均通道差异应小于0.15。若该值过大，说明RGB通道间出现不协调的独立伪影，即伪彩色明显。

如需在真实数据上验证，可将 rgb_gt 替换为任意RGB图像（需归一化到[0,1]），通过 rgb_to_bayer 生成RAW后输入模型，观察输出图像的色彩还原度与边缘清晰度。

关键说明

网络输入为单通道RAW，内部先分离为RGGB四通道，在1/2分辨率下进行特征提取，最后通过PixelShuffle恢复到全分辨率，避免在原始分辨率上直接卷积的计算开销。
多尺度融合模块使用膨胀率1/2/3的并行卷积，分别捕获局部细节、中等上下文与大范围颜色信息，有助于在平坦区域抑制颜色混叠。
损失函数中的颜色一致性正则项通过约束RGB三通道的梯度差异，直接针对伪彩色伪影进行优化。
本实现为教学演示的最小闭环，未包含数据增强、批归一化或复杂训练策略。实际应用中建议使用更大分辨率图像（如512×512）并增加训练轮数。

📚 课程目录