1. 基于CNN的RAW图像去噪

任务说明

本案例构建一个轻量级CNN模型，对单通道RAW图像（模拟拜耳阵列中单一通道数据）进行去噪处理。输入为添加高斯噪声的RAW图像，输出为去噪后的干净图像，并以PSNR（峰值信噪比）作为客观评估指标。通过本案例，掌握RAW域图像去噪的基本数据流程与CNN训练/评估闭环。

环境与依赖

需要Python 3.8+，并安装以下库：numpy、torch、torchvision、scipy（用于PSNR计算）。建议使用CUDA加速训练，若无GPU则自动回退CPU。

实现代码

以下代码完整实现了数据生成、模型定义、训练与PSNR评估。核心逻辑包括：

• 数据模拟：从公开数据集（如BSD400）读取图像并裁剪为RAW单通道数据，或直接使用随机噪声图模拟；本案例采用随机生成平滑块+噪声的方式构造训练对，保证可复现性。
• 网络结构：3层卷积的编码-解码结构，使用ReLU激活，残差学习预测噪声。
• 训练流程：Adam优化器，MSE损失，每轮后计算验证集PSNR。

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from scipy.signal import convolve2d

# ---------- 1. 数据生成 ----------
def generate_data(num_samples=1000, patch_size=64, noise_sigma=25):
    """生成训练/验证数据对 (clean, noisy)"""
    clean_batch = []
    noisy_batch = []
    for _ in range(num_samples):
        # 生成平滑RAW块（模拟真实RAW的平滑区域）
        base = np.random.rand(patch_size, patch_size).astype(np.float32) * 255
        # 添加低频纹理
        kernel = np.ones((5,5)) / 25
        base = convolve2d(base, kernel, mode='same', boundary='wrap')
        clean = np.clip(base, 0, 255)
        noise = np.random.normal(0, noise_sigma, clean.shape).astype(np.float32)
        noisy = np.clip(clean + noise, 0, 255)
        clean_batch.append(clean)
        noisy_batch.append(noisy)
    return np.stack(clean_batch, 0), np.stack(noisy_batch, 0)

# ---------- 2. CNN模型 ----------
class RawDenoiser(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 1, 3, padding=1)
        )
    def forward(self, x):
        return self.net(x)  # 直接预测干净图像

# ---------- 3. PSNR计算 ----------
def psnr(img1, img2, max_val=255.0):
    mse = np.mean((img1 - img2) ** 2)
    if mse == 0:
        return float('inf')
    return 10 * np.log10(max_val**2 / mse)

# ---------- 4. 训练与验证 ----------
def train():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 生成数据
    clean_train, noisy_train = generate_data(800)
    clean_val, noisy_val = generate_data(200)
    # 转为Tensor
    train_clean = torch.from_numpy(clean_train).unsqueeze(1).float().to(device)
    train_noisy = torch.from_numpy(noisy_train).unsqueeze(1).float().to(device)
    val_clean = torch.from_numpy(clean_val).unsqueeze(1).float().to(device)
    val_noisy = torch.from_numpy(noisy_val).unsqueeze(1).float().to(device)

    model = RawDenoiser().to(device)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    # 训练
    model.train()
    for epoch in range(20):
        optimizer.zero_grad()
        output = model(train_noisy)
        loss = criterion(output, train_clean)
        loss.backward()
        optimizer.step()
        if (epoch+1) % 5 == 0:
            print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')

    # 验证
    model.eval()
    with torch.no_grad():
        val_output = model(val_noisy).cpu().numpy()
    val_clean_np = val_clean.cpu().numpy()
    val_noisy_np = val_noisy.cpu().numpy()
    # 计算PSNR
    noisy_psnr = np.mean([psnr(val_clean_np[i,0], val_noisy_np[i,0]) for i in range(200)])
    denoised_psnr = np.mean([psnr(val_clean_np[i,0], val_output[i,0]) for i in range(200)])
    print(f'Noisy PSNR: {noisy_psnr:.2f} dB')
    print(f'Denoised PSNR: {denoised_psnr:.2f} dB')
    return model

if __name__ == '__main__':
    model = train()
    torch.save(model.state_dict(), 'raw_denoiser.pth')

运行与验证

直接执行脚本即可完成训练与验证：

python raw_denoiser.py

通过条件：训练完成后，控制台输出的Denoised PSNR应明显高于Noisy PSNR（通常提升5dB以上）。例如，当噪声标准差为25时，Noisy PSNR约为20dB，去噪后应达到28dB以上。若未达到，可适当增加训练轮数或调整网络层数。

常见问题

• 数据范围：RAW图像通常为12-14bit，本案例简化为8bit（0-255）便于演示，实际应用需按bit深度调整max_val参数。
• 噪声模型：真实RAW噪声为泊松-高斯混合分布，本案例仅使用高斯噪声近似，适用于教学演示。
• 训练速度：若CPU训练过慢，可减少num_samples至500或降低patch_size至32。