用PyTorch实现深度学习白平衡编辑：从数据集准备到模型训练全流程
·
用PyTorch实现深度学习白平衡编辑：从数据集准备到模型训练全流程
你是否曾为一张照片的色温不对而烦恼？比如在室内暖黄灯光下拍出的照片，人脸显得蜡黄；或者在阴天户外拍出的照片，整体泛着冷蓝调。传统照片编辑软件中的白平衡工具往往需要手动调整，效果依赖经验，且难以处理已经经过相机ISP（图像信号处理器）非线性渲染的sRGB图像。这正是CVPR 2020论文《Deep White-Balance Editing》要解决的核心问题——利用深度学习，让计算机学会像专业摄影师一样，智能、精准地编辑照片的白平衡。
白平衡不仅仅是“去黄”或“去蓝”那么简单。它关乎图像色彩的真实还原与美学表达。相机在拍摄时，会根据场景光源的色温自动或手动设置白平衡，但一旦图像被处理成sRGB格式，再想修改白平衡就变得异常困难，因为ISP的非线性处理过程像一层“滤镜”，难以逆向剥离。这篇论文提出的方法，绕开了复杂的逆向工程，直接学习从“错误”白平衡的sRGB图像到“正确”或“目标”白平衡图像之间的端到端映射。
本文将带你深入复现这篇论文的核心思想与工程实践。我们将使用PyTorch框架，从零开始，一步步构建一个能够同时完成白平衡校正（将错误白平衡的照片修正为自然状态）和白平衡编辑（将照片调整到任意色温，如暖调的“钨丝灯”或冷调的“阴影”效果）的深度学习模型。整个过程将涵盖数据处理、模型架构设计、训练技巧以及实际应用部署的完整链路，并提供大量可运行的代码示例和调优经验。
1. 理解问题：为什么sRGB图像的白平衡编辑如此困难？
在深入代码之前，我们必须先理解问题的本质。当你用手机或相机拍照时，图像信号处理器（ISP）会执行一系列复杂的操作，将传感器捕捉的原始RAW数据转换成我们常见的JPEG或PNG（sRGB格式）图片。这个过程的第一步就是白平衡。
白平衡的本质是补偿场景光源的颜色偏差，让白色物体在任何光线下看起来都是白色。相机通常通过一个简单的3x3对角矩阵对RAW图像的RGB通道进行缩放来实现。问题在于，这个线性操作之后，ISP还会进行一系列非线性处理，比如伽马校正、色彩空间转换、锐化、降噪等。这些操作使得sRGB图像的颜色信息与原始场景光照之间的关系变得高度非线性且不可逆。
提示：你可以把ISP的渲染管道想象成一个复杂的烹饪过程。白平衡调整是第一步“调味”，但后续还有“煎炒烹炸”（非线性处理）。一旦菜品出锅（生成sRGB图像），你很难仅凭成品精确还原出最初“调味”的配方。
因此，传统的基于照明估计和线性对角校正的方法，在sRGB图像上往往效果不佳。论文作者构建的 Rendered WB数据集 巧妙地解决了训练数据问题。他们使用软件模拟了不同相机、不同白平衡设置下的ISP渲染流程，为同一场景生成了多张不同白平衡的sRGB图像对。这为我们提供了“错误输入-正确输出”的完美监督信号。
我们的目标，就是训练一个模型，让它学会这个复杂的、非线性的映射关系：f: 错误白平衡的sRGB图像 -> 目标白平衡的sRGB图像。
2. 环境搭建与数据准备
工欲善其事，必先利其器。我们先来搭建一个稳定、高效的PyTorch开发环境，并准备好Rendered WB数据集。
2.1 创建虚拟环境与安装依赖
为了避免包版本冲突，强烈建议使用Conda或venv创建独立的Python环境。以下以Conda为例：
# 创建并激活名为‘deep_wb’的Python 3.8环境
conda create -n deep_wb python=3.8 -y
conda activate deep_wb
# 安装PyTorch（请根据你的CUDA版本访问官网获取对应命令）
# 例如，对于CUDA 11.3：
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
# 安装其他必要的库
pip install numpy pillow scipy scikit-learn matplotlib tqdm tensorboard
2.2 获取与理解Rendered WB数据集
论文作者在项目页面提供了数据集下载链接。数据集结构通常如下：
Rendered_WB_Dataset/
├── Set1/
│ ├── input/ # 输入图像，命名如 `0001_WB_T_P.png` (T代表钨丝灯白平衡)
│ └── ground_truth/ # 对应的真实图像，命名如 `0001_G_AS.png` (A代表AWB，S代表阴影白平衡)
└── Set2/
├── input/
└── ground_truth/
关键命名规则解析：
0001_WB_T_P.png:T代表白炽灯（Tungsten）白平衡，P代表某种图片风格（Picture Style）。0001_G_AS.png:A代表自动白平衡（AWB，即“正确”白平衡），S代表阴影（Shade）白平衡。实际上，对于每个输入场景，我们有三个目标输出：AWB、Tungsten WB 和 Shade WB。
我们需要编写一个数据加载器，能够正确配对输入和三个目标图像。下面是一个简化的数据集类实现：
import os
from PIL import Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms as transforms
class RenderedWBDataset(Dataset):
def __init__(self, root_dir, set_name='Set1', transform=None, patch_size=128):
"""
Args:
root_dir (string): 数据集根目录路径。
set_name (string): 'Set1' 或 'Set2'。
transform (callable, optional): 可选的变换函数。
patch_size (int): 训练时随机裁剪的patch大小。
"""
self.root_dir = root_dir
self.set_name = set_name
self.patch_size = patch_size
self.transform = transform if transform else self.get_default_transform()
# 收集所有输入图像路径
input_dir = os.path.join(root_dir, set_name, 'input')
self.input_paths = sorted([os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith('.png')])
# 根据输入图像名，构建对应的三个真实图像路径
self.gt_paths = []
for inp_path in self.input_paths:
base_name = os.path.basename(inp_path).split('_')[0] # 如 '0001'
gt_dir = os.path.join(root_dir, set_name, 'ground_truth')
# 假设命名规则为 {base_name}_G_AS.png
# 该文件实际上包含了AWB和Shade两个目标，我们需要从论文中理解如何提取三个目标。
# 根据论文和代码，通常一个输入对应三个独立的GT文件，或一个文件包含多个通道。
# 这里是一个简化示例，实际需要根据数据集具体结构调整。
awb_path = os.path.join(gt_dir, f"{base_name}_G_A.png")
tungsten_path = os.path.join(gt_dir, f"{base_name}_G_T.png")
shade_path = os.path.join(gt_dir, f"{base_name}_G_S.png")
self.gt_paths.append([awb_path, tungsten_path, shade_path])
# 数据增强：随机裁剪和翻转
self.augmentation = transforms.Compose([
transforms.RandomCrop(patch_size),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
])
def get_default_transform(self):
"""将图像转换为Tensor并归一化到[-1, 1]"""
return transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
def __len__(self):
return len(self.input_paths)
def __getitem__(self, idx):
# 加载输入图像
input_img = Image.open(self.input_paths[idx]).convert('RGB')
# 加载三个目标图像
gt_paths = self.gt_paths[idx]
gt_awb = Image.open(gt_paths[0]).convert('RGB')
gt_tungsten = Image.open(gt_paths[1]).convert('RGB')
gt_shade = Image.open(gt_paths[2]).convert('RGB')
# 应用相同的数据增强（确保输入和GT的裁剪位置一致）
seed = torch.randint(0, 2**32, (1,)).item()
torch.manual_seed(seed)
input_img = self.augmentation(input_img)
torch.manual_seed(seed)
gt_awb = self.augmentation(gt_awb)
torch.manual_seed(seed)
gt_tungsten = self.augmentation(gt_tungsten)
torch.manual_seed(seed)
gt_shade = self.augmentation(gt_shade)
# 应用转换（ToTensor + Normalize）
input_tensor = self.transform(input_img)
gt_awb_tensor = self.transform(gt_awb)
gt_tungsten_tensor = self.transform(gt_tungsten)
gt_shade_tensor = self.transform(gt_shade)
# 将三个GT堆叠成一个张量 [3, H, W] -> [9, H, W]
gt_tensor = torch.cat([gt_awb_tensor, gt_tungsten_tensor, gt_shade_tensor], dim=0)
return input_tensor, gt_tensor
注意：上述代码中的文件路径匹配逻辑是示意性的。实际使用时，你必须仔细检查Rendered WB数据集的真实文件命名规则和结构。原始论文的Matlab/PyTorch代码仓库中提供了数据加载的具体实现，是调整上述代码的最佳参考。
2.3 数据加载与可视化检查
在开始训练前，务必检查数据加载是否正确。我们可以写一个简单的脚本，显示一些输入-目标对：
import matplotlib.pyplot as plt
def show_batch(sample_batch, num_samples=4):
"""显示一个批次的数据。"""
inputs, gts = sample_batch
batch_size = inputs.size(0)
num_samples = min(num_samples, batch_size)
fig, axes = plt.subplots(num_samples, 4, figsize=(12, 3*num_samples))
for i in range(num_samples):
# 反归一化并转换为HWC格式供matplotlib显示
inp = (inputs[i].permute(1,2,0).numpy() * 0.5 + 0.5).clip(0, 1)
gt_awb = (gts[i, 0:3].permute(1,2,0).numpy() * 0.5 + 0.5).clip(0, 1)
gt_tung = (gts[i, 3:6].permute(1,2,0).numpy() * 0.5 + 0.5).clip(0, 1)
gt_shade = (gts[i, 6:9].permute(1,2,0).numpy() * 0.5 + 0.5).clip(0, 1)
axes[i, 0].imshow(inp)
axes[i, 0].set_title('Input (e.g., Incorrect WB)')
axes[i, 0].axis('off')
axes[i, 1].imshow(gt_awb)
axes[i, 1].set_title('GT: AWB')
axes[i, 1].axis('off')
axes[i, 2].imshow(gt_tung)
axes[i, 2].set_title('GT: Tungsten')
axes[i, 2].axis('off')
axes[i, 3].imshow(gt_shade)
axes[i, 3].set_title('GT: Shade')
axes[i, 3].axis('off')
plt.tight_layout()
plt.show()
# 测试数据加载
if __name__ == '__main__':
dataset = RenderedWBDataset(root_dir='path/to/Rendered_WB_Dataset', set_name='Set1', patch_size=256)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True)
sample_batch = next(iter(dataloader))
show_batch(sample_batch)
通过可视化，你应该能看到输入图像有明显的色偏（如偏黄或偏蓝），而三个目标图像分别对应了自然（AWB）、暖调（Tungsten）和冷调（Shade）的白平衡效果。这是模型需要学习的三重映射。
3. 构建多解码器U-Net架构
论文的核心创新点在于其网络架构：一个共享的编码器配合三个独立的解码器。这种设计模仿了相机的工作流程——同一张原始图像，经过不同的白平衡处理管线，输出不同的结果。共享编码器提取图像的通用特征，而每个解码器专门学习生成特定白平衡下的图像。
3.1 编码器-解码器模块设计
我们基于U-Net结构，构建带有跳跃连接的多尺度特征提取网络。首先定义基础卷积块和下采样/上采样模块。
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvBlock(nn.Module):
"""一个简单的卷积块：Conv2d -> BatchNorm -> ReLU"""
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
super().__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
self.bn = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
return self.relu(self.bn(self.conv(x)))
class DownsampleBlock(nn.Module):
"""下采样块：两个ConvBlock后接一个MaxPool2d"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.conv1 = ConvBlock(in_channels, out_channels)
self.conv2 = ConvBlock(out_channels, out_channels)
self.pool = nn.MaxPool2d(2)
def forward(self, x):
# 保存跳跃连接的特征
skip = self.conv2(self.conv1(x))
x = self.pool(skip)
return x, skip
class UpsampleBlock(nn.Module):
"""上采样块：转置卷积上采样 -> 与跳跃连接特征拼接 -> 两个ConvBlock"""
def __init__(self, in_channels, skip_channels, out_channels):
super().__init__()
# 转置卷积用于上采样
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
# 拼接后通道数为 in_channels//2 + skip_channels
self.conv1 = ConvBlock(in_channels // 2 + skip_channels, out_channels)
self.conv2 = ConvBlock(out_channels, out_channels)
def forward(self, x, skip):
x = self.up(x)
# 处理可能的尺寸不匹配（由于池化舍入）
diffY = skip.size()[2] - x.size()[2]
diffX = skip.size()[3] - x.size()[3]
x = F.pad(x, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# 拼接跳跃连接
x = torch.cat([x, skip], dim=1)
x = self.conv2(self.conv1(x))
return x
3.2 构建完整的Deep WB网络
现在，我们将这些模块组装成完整的网络。编码器部分有4级下采样，每个解码器有对应的4级上采样。
class DeepWBModel(nn.Module):
def __init__(self, in_channels=3, base_channels=24, num_decoders=3):
"""
Args:
in_channels: 输入图像通道数，RGB为3。
base_channels: 第一层卷积的通道数，论文中为24。
num_decoders: 解码器数量，对应不同的白平衡目标（AWB, Tungsten, Shade）。
"""
super().__init__()
self.num_decoders = num_decoders
# ---------- 编码器部分 ----------
# 级别 1
self.enc1 = DownsampleBlock(in_channels, base_channels) # 输出通道: base_channels
# 级别 2
self.enc2 = DownsampleBlock(base_channels, base_channels * 2) # 48
# 级别 3
self.enc3 = DownsampleBlock(base_channels * 2, base_channels * 4) # 96
# 级别 4 (瓶颈层之前)
self.enc4_conv1 = ConvBlock(base_channels * 4, base_channels * 8) # 192
self.enc4_conv2 = ConvBlock(base_channels * 8, base_channels * 8) # 192
self.enc4_pool = nn.MaxPool2d(2)
# ---------- 瓶颈层 ----------
self.bottleneck = nn.Sequential(
ConvBlock(base_channels * 8, base_channels * 16), # 384
ConvBlock(base_channels * 16, base_channels * 16), # 384
nn.ConvTranspose2d(base_channels * 16, base_channels * 8, kernel_size=2, stride=2) # 上采样回192
)
# ---------- 解码器部分 (每个解码器独立) ----------
# 为每个白平衡目标创建解码器
self.decoders = nn.ModuleList()
for _ in range(num_decoders):
decoder = nn.ModuleDict({
'up4': UpsampleBlock(base_channels * 8, base_channels * 8, base_channels * 4), # 输入192+192 -> 输出96
'up3': UpsampleBlock(base_channels * 4, base_channels * 4, base_channels * 2), # 输入96+96 -> 输出48
'up2': UpsampleBlock(base_channels * 2, base_channels * 2, base_channels), # 输入48+48 -> 输出24
'up1': UpsampleBlock(base_channels, base_channels, base_channels // 2), # 输入24+24 -> 输出12
'final_conv': nn.Conv2d(base_channels // 2, in_channels, kernel_size=1) # 输出3通道RGB
})
self.decoders.append(decoder)
def forward(self, x):
# 编码器前向传播，保存各层特征用于跳跃连接
x, skip1 = self.enc1(x) # skip1: [B, 24, H/2, W/2]
x, skip2 = self.enc2(x) # skip2: [B, 48, H/4, W/4]
x, skip3 = self.enc3(x) # skip3: [B, 96, H/8, W/8]
x = self.enc4_conv1(x)
x = self.enc4_conv2(x)
skip4 = x # [B, 192, H/16, W/16]
x = self.enc4_pool(x) # [B, 192, H/32, W/32]
# 瓶颈层
x = self.bottleneck(x) # [B, 192, H/16, W/16]
# 多个解码器并行前向传播
outputs = []
for decoder in self.decoders:
d = decoder['up4'](x, skip4) # [B, 96, H/8, W/8]
d = decoder['up3'](d, skip3) # [B, 48, H/4, W/4]
d = decoder['up2'](d, skip2) # [B, 24, H/2, W/2]
d = decoder['up1'](d, skip1) # [B, 12, H, W]
d = decoder['final_conv'](d) # [B, 3, H, W]
# 使用Tanh激活将输出限制在[-1, 1]，对应我们之前归一化的范围
d = torch.tanh(d)
outputs.append(d)
# 将三个解码器的输出堆叠在通道维度 [B, 3, H, W] * 3 -> [B, 9, H, W]
# 与数据加载器中GT的格式保持一致
return torch.cat(outputs, dim=1)
这个架构有几个关键设计点：
- 共享编码器：无论输入图像原本是什么白平衡，编码器都学习提取其通用的内容特征（边缘、纹理、结构）。
- 独立解码器：每个解码器专门学习将共享特征“渲染”成特定白平衡下的图像。这比训练三个独立的U-Net更高效，且论文证明效果更好。
- 跳跃连接：U-Net的经典设计，将编码器各层的细节特征直接传递到解码器对应层，帮助网络更好地重建高频细节。
- 输出激活：最终使用
tanh激活函数，将输出值约束在[-1, 1]区间，与我们数据归一化的范围匹配。
我们可以快速验证一下网络结构：
if __name__ == '__main__':
# 创建一个随机输入张量 [batch_size, channels, height, width]
dummy_input = torch.randn(2, 3, 256, 256)
model = DeepWBModel()
output = model(dummy_input)
print(f"输入尺寸: {dummy_input.shape}")
print(f"输出尺寸: {output.shape}") # 应为 [2, 9, 256, 256]
print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f} M")
4. 训练策略与损失函数
模型架构搭建好后，下一步是定义如何训练它。我们需要一个合适的损失函数来指导网络学习，并配置优化策略。
4.1 损失函数选择
论文中对比了L1损失和L2（MSE）损失，发现两者效果相当。L1损失（平均绝对误差） 在图像生成任务中通常比L2损失更能产生清晰的边缘，因为它对异常值不那么敏感。我们将对三个解码器的输出分别计算L1损失，然后求和。
def multi_task_l1_loss(pred, target):
"""
计算多任务L1损失。
Args:
pred: 模型输出，形状为 [B, 9, H, W] (3个目标 * 3个通道)。
target: 真实标签，形状同pred。
Returns:
标量损失值。
"""
# 确保pred和target在[-1, 1]范围内
loss = F.l1_loss(pred, target, reduction='mean')
return loss
除了像素级的L1损失，在实践中，为了获得更逼真的视觉效果，我们还可以考虑加入感知损失（Perceptual Loss） 或对抗损失（Adversarial Loss）。但为了忠实复现论文，我们首先使用基础的L1损失。
4.2 训练循环与优化器配置
我们将使用Adam优化器，并采用学习率衰减策略。以下是训练循环的核心代码框架：
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
def train_model(model, train_loader, val_loader, device, num_epochs=100, lr=1e-4):
model.to(device)
# 使用Adam优化器，参数与论文保持一致
optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999))
# 学习率调度器：每25个epoch衰减为原来的0.5倍
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=25, gamma=0.5)
# 使用TensorBoard记录训练过程
writer = SummaryWriter(log_dir='./runs/deep_wb')
best_val_loss = float('inf')
for epoch in range(num_epochs):
# 训练阶段
model.train()
train_loss = 0.0
progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs} [Train]')
for batch_idx, (inputs, targets) in enumerate(progress_bar):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = multi_task_l1_loss(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.item()
progress_bar.set_postfix({'loss': loss.item()})
avg_train_loss = train_loss / len(train_loader)
writer.add_scalar('Loss/Train', avg_train_loss, epoch)
# 验证阶段
model.eval()
val_loss = 0.0
with torch.no_grad():
for inputs, targets in val_loader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
loss = multi_task_l1_loss(outputs, targets)
val_loss += loss.item()
avg_val_loss = val_loss / len(val_loader)
writer.add_scalar('Loss/Validation', avg_val_loss, epoch)
# 学习率调整
scheduler.step()
current_lr = scheduler.get_last_lr()[0]
writer.add_scalar('Learning Rate', current_lr, epoch)
print(f'Epoch {epoch+1}: Train Loss = {avg_train_loss:.4f}, Val Loss = {avg_val_loss:.4f}, LR = {current_lr:.6f}')
# 保存最佳模型
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'val_loss': best_val_loss,
}, 'best_deep_wb_model.pth')
print(f' -> 保存最佳模型 (Val Loss: {best_val_loss:.4f})')
writer.close()
print('训练完成！')
return model
4.3 关键训练技巧与参数设置
根据论文描述，有几个训练细节对最终效果至关重要：
- 数据增强：除了代码中实现的随机裁剪和翻转，论文还提到对原始RAW图像应用了随机的色温渲染来扩充数据。我们在使用Rendered WB数据集时，主要依赖其已有的丰富变换。
- Patch训练：由于全分辨率图像太大，通常随机裁剪出128x128或256x256的图块进行训练。这增加了数据的多样性，也减少了内存消耗。
- 批量大小与迭代次数：论文使用批量大小32，训练了165,000次迭代。我们可以根据GPU内存调整批量大小（如16或8），并相应增加epoch数量。
- 权重初始化：使用He初始化（
kaiming_normal_）来初始化卷积层权重，这对深层网络的稳定训练很有帮助。PyTorch的默认初始化已经是类似的方法。
一个完整的训练启动脚本可能如下所示：
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, required=True, help='Rendered WB数据集根目录')
parser.add_argument('--batch_size', type=int, default=16)
parser.add_argument('--patch_size', type=int, default=128)
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--lr', type=float, default=1e-4)
parser.add_argument('--num_workers', type=int, default=4)
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'使用设备: {device}')
# 准备数据集
full_dataset = RenderedWBDataset(root_dir=args.data_root, set_name='Set1', patch_size=args.patch_size)
# 按论文建议，使用前2/3的数据训练，后1/3测试。这里简单按比例分割。
train_size = int(0.67 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size])
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True
)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True
)
# 初始化模型
model = DeepWBModel()
# 开始训练
train_model(model, train_loader, val_loader, device, num_epochs=args.epochs, lr=args.lr)
5. 推理、编辑与部署
模型训练好后，我们如何用它来处理新图像，并实现白平衡的灵活编辑？
5.1 全卷积推理与颜色映射后处理
我们的网络是全卷积的，理论上可以处理任意尺寸的输入。但为了保持效率，论文建议将输入图像的最长边缩放到656像素（保持宽高比），且尺寸需为24的倍数（因为4级下采样因子为2^4=16，加上一些卷积的padding，24是安全的倍数）。推理后，为了进一步提升效果并与输入分辨率匹配，论文采用了一个全局多项式颜色映射的后处理步骤。
这个步骤的核心思想是：学习一个从降采样输入图像到网络降采样输出图像的全局颜色变换矩阵，然后将这个变换应用到原始分辨率的输入图像上。这能纠正网络可能存在的全局色偏，并适应任意尺寸。
import numpy as np
from scipy import linalg
def apply_global_color_mapping(input_fullres, output_downscaled, input_downscaled, degree=3):
"""
应用全局多项式颜色映射。
Args:
input_fullres: 原始高分辨率输入图像，形状 [H, W, 3]，值范围[0, 1]。
output_downscaled: 网络输出的降采样图像，形状 [h, w, 3]，值范围[-1, 1]。
input_downscaled: 降采样后的输入图像，形状 [h, w, 3]，值范围[-1, 1]。
degree: 多项式阶数，论文中使用3阶（对应11维特征）。
Returns:
final_output: 经过颜色映射后的高分辨率输出图像，形状 [H, W, 3]，值范围[0, 1]。
"""
# 1. 将输入输出图像从 [H,W,3] 重塑为 [N, 3] 的像素矩阵，N = H*W
h, w, c = input_downscaled.shape
input_flat = input_downscaled.reshape(-1, c) # [N, 3]
output_flat = output_downscaled.reshape(-1, c) # [N, 3]
# 2. 构建多项式特征。对于3阶多项式，每个RGB像素[x_r, x_g, x_b]映射到11维特征：
# [1, r, g, b, r*g, r*b, g*b, r^2, g^2, b^2, r*g*b]
# 这里我们实现一个通用的多项式展开函数
def polynomial_expansion(x, degree):
# x: [N, 3]
features = [np.ones((x.shape[0], 1))] # 偏置项
for d in range(1, degree+1):
# 生成所有次数为d的单项式组合（r^a * g^b * b^c, 其中a+b+c=d）
# 这里简化处理，仅包含论文提到的特定项。更通用的实现可使用 itertools。
if d == 1:
features.append(x) # r, g, b
elif d == 2:
# r^2, g^2, b^2, r*g, r*b, g*b
features.append(x**2)
features.append(x[:, 0:1] * x[:, 1:2]) # r*g
features.append(x[:, 0:1] * x[:, 2:3]) # r*b
features.append(x[:, 1:2] * x[:, 2:3]) # g*b
elif d == 3:
# r*g*b
features.append(x[:, 0:1] * x[:, 1:2] * x[:, 2:3])
return np.hstack(features) # [N, D]
Phi_input = polynomial_expansion(input_flat, degree) # [N, D]
# 3. 求解线性最小二乘问题： Phi_input * M ≈ output_flat
# M 是 [D, 3] 的映射矩阵
M, _, _, _ = linalg.lstsq(Phi_input, output_flat)
# 4. 将映射应用到原始高分辨率输入图像
input_fullres_flat = input_fullres.reshape(-1, 3) # [N_full, 3]
Phi_fullres = polynomial_expansion(input_fullres_flat, degree) # [N_full, D]
output_fullres_flat = np.dot(Phi_fullres, M) # [N_full, 3]
# 5. 重塑并裁剪到有效范围
H, W, _ = input_fullres.shape
final_output = output_fullres_flat.reshape(H, W, 3)
# 由于网络输出是tanh激活，范围在[-1,1]，我们映射回[0,1]
final_output = np.clip(final_output, -1, 1)
final_output = (final_output + 1) / 2.0 # 转换到 [0, 1]
return final_output
5.2 完整的推理流程
结合网络推理和后处理，我们可以构建一个端到端的处理管道：
from PIL import Image
import torchvision.transforms as transforms
def deep_wb_edit(input_image_path, model, device, task='all', target_size=656):
"""
对单张图像进行白平衡编辑。
Args:
input_image_path: 输入图像路径。
model: 加载好权重的DeepWB模型。
device: 计算设备。
task: 'AWB'仅输出自动白平衡结果，'editing'输出三个结果，'all'输出所有。
target_size: 网络推理的目标尺寸（最长边）。
Returns:
处理后的图像字典。键为 'awb', 'tungsten', 'shade'。
"""
# 1. 加载并预处理图像
input_img = Image.open(input_image_path).convert('RGB')
original_w, original_h = input_img.size
# 计算缩放尺寸，保持长宽比，且为24的倍数
scale = target_size / max(original_w, original_h)
new_w = int(round(original_w * scale / 24) * 24)
new_h = int(round(original_h * scale / 24) * 24)
resize_transform = transforms.Compose([
transforms.Resize((new_h, new_w)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
input_tensor = resize_transform(input_img).unsqueeze(0).to(device) # [1, 3, H, W]
# 2. 网络推理
model.eval()
with torch.no_grad():
output_tensor = model(input_tensor) # [1, 9, H, W]
# 3. 分离三个输出并反归一化
output_tensor = output_tensor.squeeze(0) # [9, H, W]
# 分割成三个图像
awb_out = output_tensor[0:3].cpu().numpy().transpose(1, 2, 0) # [H, W, 3], 范围[-1,1]
tungsten_out = output_tensor[3:6].cpu().numpy().transpose(1, 2, 0)
shade_out = output_tensor[6:9].cpu().numpy().transpose(1, 2, 0)
# 获取降采样后的输入图像（用于颜色映射）
input_downscaled = input_tensor.squeeze(0).cpu().numpy().transpose(1, 2, 0) # [H, W, 3], 范围[-1,1]
# 4. 颜色映射后处理
# 将原始输入图像转换为numpy数组，范围[0,1]
input_fullres = np.array(input_img).astype(np.float32) / 255.0
# 对每个输出应用颜色映射
awb_final = apply_global_color_mapping(input_fullres, awb_out, input_downscaled)
tungsten_final = apply_global_color_mapping(input_fullres, tungsten_out, input_downscaled)
shade_final = apply_global_color_mapping(input_fullres, shade_out, input_downscaled)
# 转换为PIL图像
def to_pil(img_array):
img_array = (np.clip(img_array, 0, 1) * 255).astype(np.uint8)
return Image.fromarray(img_array)
results = {
'awb': to_pil(awb_final),
'tungsten': to_pil(tungsten_final),
'shade': to_pil(shade_final)
}
# 5. 根据任务返回结果
if task.lower() == 'awb':
return {'awb': results['awb']}
elif task.lower() == 'editing':
# 编辑任务可能只需要tungsten和shade，用于插值生成中间色温
return {'tungsten': results['tungsten'], 'shade': results['shade']}
else: # 'all'
return results
5.3 实现色温插值编辑
模型直接输出了对应于钨丝灯（约2850K）和阴影（约7500K）两个极端色温的结果。如果我们想要一个中间色温，比如5500K的日光效果，该怎么办？论文提出可以在两个输出之间进行简单的线性插值。
色温（单位：开尔文，K）数值越高，光线越冷（蓝）；数值越低，光线越暖（黄）。我们可以根据目标色温在2850K和7500K之间进行线性插值。
def wb_interpolation(tungsten_img, shade_img, target_kelvin):
"""
在两个基础白平衡图像之间插值，生成目标色温的图像。
Args:
tungsten_img: 钨丝灯白平衡结果图像，PIL Image 或 numpy数组 [0,255]。
shade_img: 阴影白平衡结果图像，格式同tungsten_img。
target_kelvin: 目标色温，介于2850和7500之间。
Returns:
插值后的图像。
"""
# 确保输入为numpy数组，范围[0, 1]
if isinstance(tungsten_img, Image.Image):
tungsten = np.array(tungsten_img).astype(np.float32) / 255.0
shade = np.array(shade_img).astype(np.float32) / 255.0
else:
tungsten = tungsten_img.astype(np.float32) / 255.0
shade = shade_img.astype(np.float32) / 255.0
# 计算插值权重
kelvin_min = 2850.0
kelvin_max = 7500.0
# 将色温映射到[0,1]的权重，0对应tungsten，1对应shade
alpha = (target_kelvin - kelvin_min) / (kelvin_max - kelvin_min)
alpha = np.clip(alpha, 0.0, 1.0)
# 线性插值
interpolated = (1 - alpha) * tungsten + alpha * shade
interpolated = np.clip(interpolated, 0, 1)
# 转换回PIL Image
interpolated_pil = Image.fromarray((interpolated * 255).astype(np.uint8))
return interpolated_pil
# 使用示例
# results = deep_wb_edit('your_photo.jpg', model, device, task='editing')
# daylight_img = wb_interpolation(results['tungsten'], results['shade'], target_kelvin=5500)
# cloudy_img = wb_interpolation(results['tungsten'], results['shade'], target_kelvin=6500)
这个简单的线性插值在实际应用中效果出奇的好，因为它本质上是在两个经过网络“渲染”的、色彩协调的图像之间进行混合，避免了直接在原始颜色空间进行复杂变换。
5.4 模型部署与集成到应用
训练好的模型可以轻松地集成到各种照片编辑应用或服务中。以下是一些部署考虑：
- 模型优化：可以使用PyTorch的
torch.jit.trace或torch.jit.script将模型转换为TorchScript，以提高推理速度并脱离Python环境运行。对于移动端，可以考虑使用ONNX格式导出，然后利用TensorRT、NCNN或MNN等推理引擎进行加速。 - 实时处理：在高端GPU上，处理一张656px的图片可能只需几十毫秒。但对于移动设备或需要处理4K图像的场景，可能需要对网络进行轻量化（如通道剪枝、量化），或者采用更小的输入尺寸（需重新训练或微调）。
- API服务：你可以用Flask或FastAPI快速搭建一个Web服务，接收用户上传的图片和 desired 色温参数，返回处理后的图片。
一个简单的Flask服务示例框架：
from flask import Flask, request, send_file, jsonify
import io
from PIL import Image
app = Flask(__name__)
# 假设 model 已加载
# model = DeepWBModel().load_state_dict(...).eval().to('cuda')
@app.route('/wb_edit', methods=['POST'])
def wb_edit():
if 'image' not in request.files:
return jsonify({'error': 'No image provided'}), 400
file = request.files['image']
target_kelvin = request.form.get('kelvin', type=float, default=5500)
# 处理图像
input_img = Image.open(io.BytesIO(file.read())).convert('RGB')
# ... 调用 deep_wb_edit 和 wb_interpolation ...
# result_img = ...
# 将结果图像转换为字节流返回
img_byte_arr = io.BytesIO()
result_img.save(img_byte_arr, format='JPEG', quality=95)
img_byte_arr.seek(0)
return send_file(img_byte_arr, mimetype='image/jpeg')
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000, debug=False)
6. 超越论文：实战技巧与进阶探索
复现论文只是第一步。在实际项目中，你可能会遇到各种挑战。以下是一些基于经验的技巧和可能的改进方向：
数据与训练相关：
- 数据不足：Rendered WB数据集虽然大，但场景有限。你可以尝试使用数据增广技术，如更复杂的色彩抖动、模拟不同相机响应曲线、添加噪声等，来提升模型的泛化能力。
- 过拟合：如果在小数据集上训练，除了常规的Dropout、权重衰减，可以尝试在编码器部分使用更强的正则化，因为它是共享的，其泛化能力至关重要。
- 损失函数改进：单纯使用L1损失可能导致结果过于平滑。可以尝试结合感知损失（使用预训练的VGG网络提取特征进行对比）来提升视觉质量，或者引入对抗损失（GAN）让结果看起来更“真实”。
模型架构相关：
- 轻量化：如果部署到资源受限环境，可以考虑使用MobileNetV2、EfficientNet等轻量级骨干网络作为编码器，或者对现有网络进行通道剪枝和量化。
- 多尺度训练与测试：训练时使用随机尺寸的patch，测试时使用多尺度推理并融合结果（Test-Time Augmentation, TTA），可以提升效果的稳定性。
- 注意力机制：在编码器和解码器之间引入空间或通道注意力模块，让网络更关注颜色异常的区域，可能有助于处理极端色偏。
应用扩展：
- 视频白平衡：对视频序列应用此模型时，需考虑帧间一致性。可以引入光流或时序循环网络来平滑相邻帧的白平衡调整，避免闪烁。
- 与其他编辑任务结合：将白平衡编辑网络与曝光校正、去噪、超分辨率等网络结合，构建一个端到端的图像增强管线。
- 用户交互：论文方法提供了两个基准点（钨丝灯、阴影）。你可以设计一个交互界面，让用户在这两个结果之间滑动选择色温，甚至允许用户点击图像中的“中性色”（如白色、灰色物体）来进行更个性化的白平衡校正。
在项目实践中，我遇到过的一个典型问题是模型对某些极端场景（如单一颜色背景、极低光照）的处理不佳。这时，在损失函数中加入对图像亮度和对比度的约束，或者在数据集中针对性补充这类困难样本，往往是有效的解决方案。另一个经验是，颜色映射后处理步骤对于最终效果的提升非常明显，不要因为其简单而忽略它，它有效地将网络在低分辨率上学到的颜色变换关系推广到了全分辨率。
通过这套完整的流程——从理解问题、准备数据、构建模型、训练调优到最终部署——你不仅能够复现一篇顶会论文，更能掌握将深度学习应用于具体图像处理任务的系统工程能力。白平衡编辑只是一个起点，这套“编码器-多解码器”的框架思想，可以迁移到许多其他图像到图像的翻译任务中，例如风格迁移、图像修复、HDR合成等。
更多推荐
6
0- 0
已为社区贡献6条内容
所有评论(0)