2. 残差学习用于低光照RAW增强

任务目标

构建一个轻量残差卷积网络，将低光照RAW图像（线性Bayer数据）映射为正常光照RGB图像。网络需同时完成去马赛克、亮度提升与噪声抑制，输出8-bit sRGB图像。

环境依赖

Python 3.9+，PyTorch 2.0+，OpenCV（仅用于图像读写），NumPy。无需额外数据集——使用合成数据验证网络收敛性与映射能力。

核心实现

import torch
import torch.nn as nn
import numpy as np

# ---------- 残差块 ----------
class ResBlock(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.conv1 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(ch)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(ch)

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return self.relu(out + x)  # 残差连接

# ---------- 网络：RAW(1ch) → RGB(3ch) ----------
class LowLightRAWEnhancer(nn.Module):
    def __init__(self):
        super().__init__()
        self.head = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.ReLU(inplace=True)
        )
        self.res_blocks = nn.Sequential(*[ResBlock(32) for _ in range(4)])
        self.tail = nn.Sequential(
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 3, 1)  # 输出3通道sRGB
        )

    def forward(self, raw):
        x = self.head(raw)
        x = self.res_blocks(x)
        return self.tail(x)

# ---------- 合成训练数据 ----------
def make_synthetic_batch(batch_size=4, H=64, W=64, noise_std=0.02):
    """生成低光照RAW(单通道Bayer模拟)与正常光照RGB对"""
    # 正常光照RGB：随机颜色块
    rgb_gt = torch.rand(batch_size, 3, H, W)
    # 模拟Bayer采样：取R/G/G/B通道拼成单通道RAW
    bayer = torch.zeros(batch_size, 1, H, W)
    bayer[:, :, 0::2, 0::2] = rgb_gt[:, 0:1, 0::2, 0::2]  # R
    bayer[:, :, 0::2, 1::2] = rgb_gt[:, 1:2, 0::2, 1::2]  # G
    bayer[:, :, 1::2, 0::2] = rgb_gt[:, 1:2, 1::2, 0::2]  # G
    bayer[:, :, 1::2, 1::2] = rgb_gt[:, 2:3, 1::2, 1::2]  # B
    # 低光照退化：暗化 + 高斯噪声
    raw_low = bayer * 0.1 + torch.randn_like(bayer) * noise_std
    return raw_low, rgb_gt

# ---------- 训练与验证 ----------
def train_step(model, opt, loss_fn, raw, rgb):
    opt.zero_grad()
    pred = model(raw)
    loss = loss_fn(pred, rgb)
    loss.backward()
    opt.step()
    return loss.item()

if __name__ == "__main__":
    torch.manual_seed(42)
    model = LowLightRAWEnhancer()
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()

    # 训练50轮
    for epoch in range(50):
        raw, rgb = make_synthetic_batch()
        loss = train_step(model, opt, loss_fn, raw, rgb)
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1:3d} | Loss: {loss:.6f}")

    # 验证：单张推理
    model.eval()
    with torch.no_grad():
        raw_test, rgb_test = make_synthetic_batch(batch_size=1)
        pred = model(raw_test)
        # 计算PSNR（峰值信噪比）
        mse = loss_fn(pred, rgb_test).item()
        psnr = 10 * np.log10(1.0 / mse)
        print(f"验证PSNR: {psnr:.2f} dB")
        print(f"输入RAW均值: {raw_test.mean():.4f} → 输出RGB均值: {pred.mean():.4f}")

运行方式

python low_light_raw_enhance.py

验证方法与通过标准

• 收敛性：训练损失应持续下降，50轮后MSE应低于0.01（对应PSNR > 20dB）。
• 亮度提升：输出RGB均值应显著高于输入RAW均值（输入约0.05，输出应接近0.5）。
• 结构正确性：输出张量形状为(1, 3, H, W)，值域在[0,1]内。
• 噪声抑制：在固定噪声种子下，输出图像应比输入RAW视觉上更平滑（可通过计算输出局部方差验证）。

关键说明

• 合成数据中Bayer采样仅用于模拟RAW单通道输入，实际应用时需按真实CFA模式排列。
• 残差块数量（4个）与通道数（32）为最小可复现配置，增加可提升效果但会延长训练时间。
• 网络尾部使用1×1卷积将64通道压缩至3通道，避免引入额外空间伪影。
• 训练50轮约需10秒（CPU），若使用GPU可缩短至2秒内。