1. 基于神经网络的自动白平衡

任务说明

本案例实现一个轻量级CNN，输入RAW格式的Bayer图像（RGGB），输出每个通道的增益系数（R/G/B），用于自动白平衡校正。同时实现传统灰度世界算法作为对比基线，并给出可量化的评估指标。

环境依赖

Python 3.8+，NumPy，PyTorch 2.0+，OpenCV（仅用于图像读写）
数据集：使用公开的MIT-Adobe FiveK数据集中的RAW图像（或任意包含RAW格式的相机图像），本案例提供数据加载与预处理接口

核心实现

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import cv2
import os

# ---------- 1. 数据预处理 ----------
def load_raw_image(path):
    """读取RAW Bayer图像（假设为16bit RGGB，尺寸HxW）"""
    raw = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if raw.dtype != np.uint16:
        raw = (raw / 256).astype(np.uint16)
    return raw.astype(np.float32) / 65535.0  # 归一化到[0,1]

def bayer_to_rgb_channels(bayer):
    """从RGGB Bayer中提取R/G/G/B四个通道（每个通道尺寸减半）"""
    H, W = bayer.shape
    R = bayer[0:H:2, 0:W:2]
    G1 = bayer[0:H:2, 1:W:2]
    G2 = bayer[1:H:2, 0:W:2]
    B = bayer[1:H:2, 1:W:2]
    return R, G1, G2, B

def pack_bayer(bayer):
    """将Bayer图像打包为4通道输入（R,G1,G2,B）"""
    R, G1, G2, B = bayer_to_rgb_channels(bayer)
    return np.stack([R, G1, G2, B], axis=-1)  # (H/2, W/2, 4)

# ---------- 2. 数据集 ----------
class RAWDataset(Dataset):
    def __init__(self, raw_paths, gt_gains, patch_size=64):
        self.raw_paths = raw_paths
        self.gt_gains = gt_gains  # (N, 3) 每行[R_gain, G_gain, B_gain]
        self.patch_size = patch_size

    def __len__(self):
        return len(self.raw_paths)

    def __getitem__(self, idx):
        raw = load_raw_image(self.raw_paths[idx])
        packed = pack_bayer(raw)  # (H/2, W/2, 4)
        # 随机裁剪patch
        H, W, _ = packed.shape
        ph, pw = self.patch_size // 2, self.patch_size // 2
        if H > ph and W > pw:
            y = np.random.randint(0, H - ph)
            x = np.random.randint(0, W - pw)
            patch = packed[y:y+ph, x:x+pw]
        else:
            patch = cv2.resize(packed, (pw, ph))
        # 转Tensor: (C, H, W)
        patch_t = torch.from_numpy(patch.transpose(2, 0, 1)).float()
        gain_t = torch.from_numpy(self.gt_gains[idx]).float()
        return patch_t, gain_t

# ---------- 3. CNN模型 ----------
class AWB_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(4, 16, 3, padding=1), nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        self.regressor = nn.Sequential(
            nn.Linear(32, 16), nn.ReLU(),
            nn.Linear(16, 3)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        gains = self.regressor(x)
        # 约束增益为正且接近1（通过sigmoid+offset）
        gains = torch.sigmoid(gains) * 2.0 + 0.5  # 范围[0.5, 2.5]
        return gains

# ---------- 4. 传统灰度世界算法 ----------
def gray_world_gains(raw_bayer):
    """灰度世界假设：R/G/B通道均值相等"""
    R, G1, G2, B = bayer_to_rgb_channels(raw_bayer)
    G = (G1 + G2) / 2.0
    mean_R, mean_G, mean_B = R.mean(), G.mean(), B.mean()
    # 以G通道均值为基准
    gain_R = mean_G / (mean_R + 1e-6)
    gain_B = mean_G / (mean_B + 1e-6)
    return np.array([gain_R, 1.0, gain_B])

# ---------- 5. 校正与评估 ----------
def apply_gains(raw_bayer, gains):
    """将增益应用到Bayer各通道"""
    R, G1, G2, B = bayer_to_rgb_channels(raw_bayer)
    gain_R, gain_G, gain_B = gains
    R = np.clip(R * gain_R, 0, 1)
    G1 = np.clip(G1 * gain_G, 0, 1)
    G2 = np.clip(G2 * gain_G, 0, 1)
    B = np.clip(B * gain_B, 0, 1)
    # 重组为Bayer
    H, W = raw_bayer.shape
    out = np.zeros_like(raw_bayer)
    out[0:H:2, 0:W:2] = R
    out[0:H:2, 1:W:2] = G1
    out[1:H:2, 0:W:2] = G2
    out[1:H:2, 1:W:2] = B
    return out

def compute_angular_error(gain_pred, gain_gt):
    """计算预测增益与真值增益的夹角误差（度）"""
    pred = gain_pred / (np.linalg.norm(gain_pred) + 1e-6)
    gt = gain_gt / (np.linalg.norm(gain_gt) + 1e-6)
    cos_theta = np.clip(np.dot(pred, gt), -1, 1)
    return np.degrees(np.arccos(cos_theta))

# ---------- 6. 训练与验证 ----------
def train_model(model, dataloader, epochs=20, lr=1e-3):
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()
    model.train()
    for epoch in range(epochs):
        total_loss = 0.0
        for patches, gains_gt in dataloader:
            optimizer.zero_grad()
            gains_pred = model(patches)
            loss = criterion(gains_pred, gains_gt)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(dataloader):.4f}")

def validate(model, raw_paths, gt_gains):
    """在验证集上比较CNN与灰度世界"""
    model.eval()
    cnn_errors, gw_errors = [], []
    with torch.no_grad():
        for path, gt in zip(raw_paths, gt_gains):
            raw = load_raw_image(path)
            # CNN预测
            packed = pack_bayer(raw)
            patch_t = torch.from_numpy(packed.transpose(2, 0, 1)).float().unsqueeze(0)
            gains_cnn = model(patch_t).squeeze().numpy()
            # 灰度世界
            gains_gw = gray_world_gains(raw)
            # 计算误差
            cnn_errors.append(compute_angular_error(gains_cnn, gt))
            gw_errors.append(compute_angular_error(gains_gw, gt))
    print(f"CNN平均角度误差: {np.mean(cnn_errors):.2f}°")
    print(f"灰度世界平均角度误差: {np.mean(gw_errors):.2f}°")
    return np.mean(cnn_errors), np.mean(gw_errors)

# ---------- 7. 主流程 ----------
if __name__ == "__main__":
    # 假设已有数据列表（替换为实际路径）
    train_raw = ["train1.dng", "train2.dng"]  # 实际需转换为可读格式
    train_gt = np.array([[1.2, 1.0, 1.8], [1.1, 1.0, 1.6]])
    val_raw = ["val1.dng"]
    val_gt = np.array([[1.3, 1.0, 1.7]])

    # 创建数据加载器
    dataset = RAWDataset(train_raw, train_gt)
    loader = DataLoader(dataset, batch_size=4, shuffle=True)

    # 初始化并训练
    model = AWB_CNN()
    train_model(model, loader, epochs=10)

    # 验证对比
    validate(model, val_raw, val_gt)

运行方式

将RAW图像（.dng或.tiff格式）转换为16bit单通道Bayer数组，或直接使用模拟Bayer数据测试
准备训练数据：每张RAW图像对应一个3维增益真值（R_gain, G_gain, B_gain），可通过相机标定或手动白平衡获得
执行脚本：python awb_cnn.py

验证方法

训练过程中观察MSE损失下降，最终收敛至较低水平（如<0.01）
验证阶段输出CNN与灰度世界的平均角度误差，CNN误差应显著低于灰度世界（通常低3°以上）
目视检查：将校正后的Bayer图像转换为RGB显示，确认色偏被有效消除

常见问题

若训练不收敛，可降低学习率至1e-4，或增加训练轮次
输入图像尺寸过小时，AdaptiveAvgPool仍可工作，但建议patch≥32×32
真值增益获取困难时，可使用灰度世界结果作为伪标签进行自监督预训练

📚 课程目录

1. 基于神经网络的自动白平衡

任务说明

环境依赖

核心实现

运行方式

验证方法

常见问题