2. 深度学习去马赛克

任务目标

本案例训练一个轻量级U-Net模型，将单通道RAW Bayer图像（RGGB模式）转换为三通道RGB图像，并与双线性插值去马赛克方法进行PSNR和SSIM对比。核心目的是验证深度学习去马赛克相比传统插值方法的优势。

环境与依赖

需要Python 3.8+，安装以下依赖：

pip install torch torchvision numpy opencv-python scikit-image

数据准备与预处理

使用公开数据集（如MIT-Adobe FiveK或自建数据集）中的RGB图像，按以下流程生成训练对：

将RGB图像裁剪为128×128的patch；
模拟CFA采样：从RGB中提取RGGB Bayer模式，生成单通道RAW图像；
输入为RAW单通道图（归一化到[0,1]），标签为原始RGB三通道图。

Bayer采样核心逻辑：

def rgb_to_bayer(rgb):
    h, w, _ = rgb.shape
    bayer = np.zeros((h, w), dtype=np.float32)
    bayer[0::2, 0::2] = rgb[0::2, 0::2, 0]  # R
    bayer[0::2, 1::2] = rgb[0::2, 1::2, 1]  # G
    bayer[1::2, 0::2] = rgb[1::2, 0::2, 1]  # G
    bayer[1::2, 1::2] = rgb[1::2, 1::2, 2]  # B
    return bayer

U-Net模型实现

采用轻量级U-Net结构，输入1通道，输出3通道，编码器3层下采样，解码器对应上采样，跳跃连接保留细节：

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

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.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.conv(x)

class UNet(nn.Module):
    def __init__(self, in_channels=1, 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.out = 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.dec3(torch.cat([self.up3(b), e3], dim=1))
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
        return torch.sigmoid(self.out(d1))

训练与评估脚本

import numpy as np
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim

class BayerDataset(Dataset):
    def __init__(self, rgb_images):
        self.rgb_images = rgb_images  # list of HxWx3 float32 arrays in [0,1]
    def __len__(self):
        return len(self.rgb_images)
    def __getitem__(self, idx):
        rgb = self.rgb_images[idx]
        bayer = rgb_to_bayer(rgb)
        bayer_t = torch.from_numpy(bayer).unsqueeze(0)
        rgb_t = torch.from_numpy(rgb.transpose(2,0,1))
        return bayer_t, rgb_t

def train(model, loader, epochs=30, lr=1e-3):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    opt = optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.L1Loss()
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            opt.zero_grad()
            pred = model(x)
            loss = loss_fn(pred, y)
            loss.backward()
            opt.step()
            total_loss += loss.item()
        print(f'Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(loader):.4f}')

def evaluate(model, test_rgb):
    """test_rgb: list of HxWx3 float32 arrays"""
    device = next(model.parameters()).device
    model.eval()
    psnr_unet, ssim_unet = [], []
    psnr_bilinear, ssim_bilinear = [], []
    with torch.no_grad():
        for rgb in test_rgb:
            bayer = rgb_to_bayer(rgb)
            # U-Net prediction
            x = torch.from_numpy(bayer).unsqueeze(0).unsqueeze(0).to(device)
            pred = model(x).squeeze(0).cpu().numpy().transpose(1,2,0)
            pred = np.clip(pred, 0, 1)
            # Bilinear interpolation (OpenCV demosaicing)
            bayer_8u = (bayer*255).astype(np.uint8)
            bilinear = cv2.cvtColor(bayer_8u, cv2.COLOR_BayerBG2RGB).astype(np.float32)/255.0
            # Compute metrics
            psnr_unet.append(psnr(rgb, pred, data_range=1.0))
            ssim_unet.append(ssim(rgb, pred, channel_axis=2, data_range=1.0))
            psnr_bilinear.append(psnr(rgb, bilinear, data_range=1.0))
            ssim_bilinear.append(ssim(rgb, bilinear, channel_axis=2, data_range=1.0))
    print(f'U-Net    - PSNR: {np.mean(psnr_unet):.2f} dB, SSIM: {np.mean(ssim_unet):.4f}')
    print(f'Bilinear - PSNR: {np.mean(psnr_bilinear):.2f} dB, SSIM: {np.mean(ssim_bilinear):.4f}')

运行方式

# 1. 准备数据（示例：加载RGB图像列表）
# rgb_list = [cv2.cvtColor(cv2.imread(f), cv2.COLOR_BGR2RGB).astype(np.float32)/255.0 for f in image_files]

# 2. 划分训练/测试集
# train_rgb, test_rgb = rgb_list[:800], rgb_list[800:]

# 3. 创建数据加载器
# dataset = BayerDataset(train_rgb)
# loader = DataLoader(dataset, batch_size=16, shuffle=True)

# 4. 训练
# model = UNet()
# train(model, loader, epochs=30)

# 5. 评估
# evaluate(model, test_rgb)

验证方法与通过标准

运行上述训练脚本后，观察测试集上的评估输出。判定训练成功的标准：

训练损失随epoch稳定下降，最终L1损失低于0.05；
U-Net的PSNR比双线性插值高至少2dB，SSIM高至少0.02；
U-Net输出图像无明显伪彩色（边缘无红蓝交错），细节纹理清晰。

若未达到上述指标，可增加训练epoch数至50或调整学习率为5e-4。

常见问题

Bayer模式匹配：OpenCV的COLOR_BayerBG2RGB对应RGGB排列，需与rgb_to_bayer中的采样位置一致；
显存不足：降低batch_size至8或裁剪patch为64×64；
输出偏灰：检查是否对输出做了sigmoid激活，且标签归一化在[0,1]区间。

📚 课程目录