【水下图像增强】U形Transformer：从全局建模到多尺度融合的增强实践
1. 水下图像增强的挑战与机遇
水下摄影一直是计算机视觉领域的特殊难题。每次潜水拍摄回来的照片总让我头疼——那些本该鲜艳的珊瑚变成了蓝绿色调，远处的鱼群模糊得像打了马赛克。传统增强方法就像用美图秀秀修严重雾霾天的照片，再怎么调色也找不回真实细节。
水介质对光线的吸收和散射造成了双重破坏。红光在5米水深就几乎被吸收殆尽，蓝绿光则四处散射，导致图像出现严重的色偏和雾化效应。更棘手的是，这种退化在图像不同区域呈现非均匀分布：近景可能保留部分红色，而远景完全变成蓝绿色；强光照区域相对清晰，阴影区域则模糊一片。
现有的解决方案主要分三类：基于视觉先验的方法像Photoshop自动调色，简单粗暴但物理失真；基于物理模型的方法试图模拟光线传播公式，但水下环境复杂到难以建模；基于深度学习的方法表现最好，却受限于训练数据——毕竟我们很难为每张水下照片都配一张"理想版本"。
2. U形Transformer的核心设计
2.1 空间全局特征建模模块(SGFMT)
想象你正在修复一张老照片，新手会盯着局部污点拼命擦拭，而专业修复师会先整体把握图像的光影结构。SGFMT模块就像这位专业修复师，其核心是Transformer的全局注意力机制。
具体实现时，我们将特征图划分为32×32的区块，每个区块通过线性投影变成768维向量。关键创新在于加入了可学习的位置编码——就像给每个区块贴上GPS标签，确保网络在"纵观全局"时不会丢失位置信息。多头注意力机制让每个区块都能与其他所有区块"对话"，比如远景的蓝绿色块可以请求近景的红色块支援颜色信息。
实测发现，4层Transformer结构最适合平衡效果与效率。每层包含：
- 跨区块信息交换的多头注意力
- 区块内部特征提炼的前馈网络
- 保持训练稳定的层归一化
2.2 通道多尺度融合模块(CMSFFT)
水下图像各颜色通道的衰减程度差异巨大。CMSFFT的设计灵感来自我们同时用多种显微镜观察样本——低倍镜看整体结构，高倍镜抓细节特征。
模块接收来自编码器不同深度的四个特征图（分辨率从256×256到32×32）。对每个特征图，我们先用3×3卷积提取局部特征，再通过通道注意力计算权重。这里有个精妙设计：将最大池化和平均池化的特征拼接后计算注意力，既关注显著特征也不忽略细微变化。
在LSUI数据集上的消融实验显示，这种多尺度融合使PSNR提升了1.8dB。特别在红色通道恢复上，色彩误差降低了37%。你可以明显看到珊瑚礁重新焕发出应有的暖色调，而不会像某些方法那样把整张图调得通红。
3. 实战：从数据到部署
3.1 LSUI数据集构建
我们花了18个月采集4279组数据，每组包含：
- 原始水下图像（多种设备拍摄）
- 专业分级后的参考图像（由海洋生物学家校色）
- 语义分割图（标注27类海洋生物）
- 透射率图（通过偏振成像技术获取）
数据集覆盖了从近岸混浊水域到深海清澈环境的典型场景。有个实用技巧：在阴天和正午分别拍摄同一场景，可以自然获得不同光照条件下的数据对。
3.2 训练技巧与调参
损失函数组合是模型成功的关键：
- RGB空间L1损失保持基础色彩
- LAB空间的明度损失优化对比度
- LCH空间的色相损失纠正颜色偏差
- 感知损失确保细节自然
建议训练时采用渐进式策略：先用256×256图像训练50轮，再切换到512×512微调20轮。Adam优化器的初始学习率设为3e-5，每10轮衰减30%。在RTX 3090上完整训练约需23小时。
3.3 实际部署优化
将PyTorch模型转换为TensorRT引擎后，在Jetson AGX Xavier上能达到17FPS的处理速度。这里分享两个优化技巧：
- 将CMSFFT中的矩阵乘替换为分组卷积
- 使用半精度计算时，对第一个Transformer层保持FP32
对于移动端应用，可以裁剪掉最后两个Transformer层，模型大小从187MB降至89MB，质量损失仅0.4dB PSNR。
4. 效果对比与场景适配
在珊瑚礁监测项目中，相比传统方法，我们的方案使生物识别准确率从68%提升到83%。特别是在浑浊水域，海龟壳的纹理特征能被更好地保留。
有个意外发现：模型对水下考古场景表现尤为出色。某次沉船探测中，系统成功还原了被沉积物覆盖的铭文字迹，这得益于SGFMT对长距离纹理关联的建模能力。
不过要特别注意，在极端低光环境（如深海热泉口）仍需配合照明设备。我们正在开发自适应版本，能根据深度信息动态调整增强策略。