1. 项目概述
水下图像增强(Underwater Image Enhancement, UIE)是计算机视觉领域的一个重要研究方向,旨在改善水下拍摄图像的质量。由于水下环境的特殊性,光线在水中传播时会受到吸收和散射的影响,导致图像出现颜色失真、对比度降低和模糊等问题。传统的水下图像增强方法主要基于物理模型,如Retinex理论和水下成像模型,但这些方法在复杂场景下效果有限。
近年来,深度学习技术在水下图像增强领域取得了显著进展。然而,现有的基于生成对抗网络(GAN)的方法存在训练不稳定、模式崩塌等问题。本文提出了一种基于Transformer的条件扩散模型(Transformer-based Diffusion Model),通过引入非均匀采样策略(Non-uniform Sampling for Skip Strategy)来提升水下图像增强的效果和效率。
2. 核心原理与技术路线
2.1 扩散模型基础
扩散模型(Diffusion Model)是一种生成模型,包含前向过程和反向过程两个阶段:
- 前向过程:通过逐步添加高斯噪声,将清晰图像逐渐转化为纯噪声图像
- 反向过程:训练神经网络从噪声图像逐步去噪,恢复出清晰图像
数学上,前向过程可以表示为:
code复制q(x_t|x_{t-1}) = N(x_t; √(1-β_t)x_{t-1}, β_tI)
其中β_t控制噪声添加的强度。
2.2 条件扩散模型
传统扩散模型从随机噪声生成图像,输出具有多样性。但在水下图像增强任务中,我们需要生成与输入图像对应的特定增强结果。为此,本文提出了条件扩散模型:
- 条件输入:将水下退化图像作为条件c输入网络
- 训练目标:最小化预测噪声与真实噪声的L1损失:
code复制L_s = ||ε_t - ε_θ(x_t,c,t)||
- 反向采样:通过迭代公式逐步生成增强图像:
code复制x_{t-1} = (x_t - β_t/√(1-α_t)ε_θ(x_t,c,t))/√(1-β_t) + σ_t z
2.3 Transformer去噪网络
本文设计了一个轻量级的基于Transformer的去噪网络,主要特点包括:
- 通道注意力机制:在通道维度进行自注意力计算,降低计算复杂度
- 时间步嵌入:将时间步t通过全连接层编码后与图像特征相加
- 网络结构:仅使用8个Transformer块,显著减少参数量
网络处理流程:
- 拼接噪声图像x_t和条件图像c,得到6通道输入
- 通过卷积层调整通道数
- 时间步特征与图像特征相加
- 经过多个Transformer块处理
- 输出预测的噪声
3. 非均匀采样策略
3.1 跳步采样
扩散模型反向过程通常需要数千次迭代,计算开销大。本文采用跳步采样策略(Skip Sampling Strategy)来减少迭代次数:
- 训练时使用密集时间步(T=2000)
- 推理时只选择部分时间步进行采样(如S=10)
- 通过DDIM方法移除随机项,使采样过程确定性
3.2 分段采样
观察到反向过程的前期比后期更重要,本文提出分段采样:
- 将时间序列分为两段:[a,c]和[c,b]
- 前段使用较小步长d1,后段使用较大步长d2
- 公式表示:
code复制τ_s = {
τ_0 + (s-1)·d1, τ_s ∈ [a,c]
τ_0 + (s-1)·d2, τ_s ∈ [c,b]
}
3.3 进化算法搜索
为进一步优化采样序列,本文采用进化算法:
- 将采样序列视为"基因序列"
- 通过变异和交叉生成新序列
- 使用PSNR作为适应度函数评估序列质量
- 保留高质量序列,淘汰低质量序列
- 经过多轮迭代得到最优采样序列
4. 实验与结果分析
4.1 实验设置
数据集:
- UIEB数据集:890组成对图像,使用90张测试
- LSUI数据集:5004组成对图像,使用4500张训练,504张测试
评价指标:
- PSNR(峰值信噪比):衡量内容相似性
- SSIM(结构相似性):衡量结构相似性
实现细节:
- 框架:PyTorch
- 优化器:Adam,学习率1e-4
- Batch size:8
- 图像尺寸:训练128×128,测试256×256
- 时间步:T=2000
- 硬件:NVIDIA RTX 3080
4.2 对比实验
与Fusion、MMLE、HLRP、WaterNet、Ucolor等方法的对比结果显示:
-
定量结果:
- PSNR:本文方法最优,比Ushape提升3.49dB
- SSIM:具有竞争力,略低于部分方法
- 参数量:10M
- 推理时间:0.13s/张
-
定性结果:
- 颜色校正效果显著优于对比方法
- 细节恢复更好,伪影更少
- 避免了Ucolor的绿色残留和Ushape的红色伪影
4.3 消融实验
-
去噪网络结构:
- Transformer网络比UNet提升显著
- 参数量更少,运行更高效
-
跳步采样步数:
- S=40/20/10时PSNR/SSIM保持较高水平
- S=10时速度显著提升
-
采样策略:
- 分段采样(Ours_p)和进化搜索采样(Ours_s)均优于均匀采样
- Ours_s略优于Ours_p
5. 实操经验与注意事项
5.1 实现技巧
-
条件图像处理:
- 将水下图像与噪声图像在通道维度拼接
- 保持图像尺寸一致
- 像素值归一化到[-1,1]
-
时间步嵌入:
- 使用全连接层编码时间步
- reshape为特征图后与图像特征相加
- 确保时间信息有效融入网络
-
通道注意力:
- 使用全局平均池化获取通道统计量
- 通过1D卷积学习通道权重
- 使用sigmoid激活函数限制权重范围
5.2 调参建议
-
噪声调度:
- β_t线性采样于[10^-6,10^-2]
- 初始小噪声,后期逐渐增大
-
进化算法:
- 初始种群规模建议50-100
- 变异率0.1-0.3
- 交叉率0.6-0.8
- 迭代轮数10-20
-
分段采样:
- 前段占比60-70%
- 前段步长为后段的1/3-1/2
5.3 常见问题与解决
-
颜色失真:
- 检查通道注意力模块
- 增加颜色相关数据增强
- 调整损失函数权重
-
细节模糊:
- 增加Transformer块数
- 调整跳步采样策略
- 尝试不同的噪声调度
-
训练不稳定:
- 减小学习率
- 增大batch size
- 使用梯度裁剪
6. 扩展与应用
本文方法可扩展到其他图像增强任务:
- 低照度增强:将低照度图像作为条件输入
- 图像去雾:调整网络结构适应雾图特性
- 医学图像增强:修改损失函数适应医学图像特点
未来改进方向:
- 无监督学习:减少对成对数据的依赖
- 实时应用:进一步优化推理速度
- 多任务学习:联合处理多种退化类型
在实际部署时,建议:
- 根据具体场景调整采样策略
- 对特定水下环境进行微调
- 考虑模型轻量化以适应边缘设备