深度学习水下图像增强实战：从U-Net训练到YOLO检测

简介面向深度学习与计算机视觉方向毕业设计、课程设计及人工智能初学者这套系统针对水下图像低对比度、低亮度与色彩失真问题基于卷积神经网络与U-net等结构通过数据驱动方式实现自动增强。资源共43个文件以Python源码为主10个py辅以模型权重meta、index、data-00000-of-00001、checkpoint、测试效果图png、依赖清单及README文档压缩包仅6.04MB轻量便于快速搭建与学习。当前已有101人学习下载。内容包含完整可运行系统app.py为程序入口UWCNN与WaterNet两个子目录分别存放模型架构、训练及评估脚本data文件夹提供训练测试数据集requirements.txt列明依赖版本README.md指导环境配置与运行方法。读者可借此掌握从数据预处理、模型训练到指标评估如PSNR、SSIM的完整流程也可直接替换或扩展模型用于科研探索。1. 水下图像增强为什么普通算法在这套数据上全部失灵拿到《基于深度学习的水下图像增强系统.zip》这个工程时我第一反应是先跑一组常见的陆上图像增强算法对比——直方图均衡、伽马校正、Retinex全部来一遍。结果很符合预期直方图均衡把偏绿的水下照片拉成了灰紫色Retinex 把暗部噪点放大了三倍。水下图像增强的核心难点不是“对比度低”而是蓝绿通道严重失衡、雾化散射、以及深度带来的非线性衰减这三者叠加在一起传统算法只能靠手工调参数硬撑换个场景就失效。这套资源属于典型的“深度学习 毕业设计/课程设计”组合用卷积网络学习水下退化图像到清晰图像的映射配套完整的数据处理、训练、评估和演示脚本工程结构贴近一个能直接答辩演示的系统不是那种只有推理 demo 的代码包。它的目标用户很明确——正在做毕业设计、课程设计或者想把水下增强作为项目经历写进简历的从业者你不需要自己搭一套完整的深度学习工程只需要弄懂关键模块怎么改、参数怎么调、坑在哪儿。2. 数据与预处理水下的光学衰减和两类数据集清洗方案2.1 为什么普通图像增强模型在水下翻车在陆上图像退化主要是光照不足或噪点在水下光在水体中传播时会被吸收和散射红光衰减最快蓝绿光衰减最慢这就是水下照片普遍偏蓝绿的原因。深度学习模型学的是输入到输出的映射如果输入分布不稳定——有的图偏绿、有的图偏蓝、有的图带严重的前向散射雾——模型就会困惑同一个增强目标输入特征完全不同。我一般会先用一个简单实验验证这种困惑把 UIEB 数据集里的图像按色偏程度分组分别统计 RGB 三个通道的像素均值。偏绿的图通常 G 通道均值比 R 通道高出 40% 以上偏蓝的图是 B 通道占优。用这些统计做分组的依据后续训练和评估才有章法。2.2 数据集的目录组织与清洗脚本这套工程默认的数据组织方式是双目录配对input/放退化图target/放对应的清晰参考图。水下增强数据集大多依赖合成配对或真实场景采集真实配对数据量少合成数据又存在域差异所以目录组织要预留足够的扩展空间。# 数据集目录结构推荐兼容 PyTorch ImageFolder 和自写 Dataset water_dataset/ ├── input/ # 退化图像命名如 001_in.jpg │ ├── train/ │ ├── val/ │ └── test/ ├── target/ # 清晰参考图命名如 001_gt.jpg │ ├── train/ │ ├── val/ │ └── test/ ├── uieb_clean.py # 清洗脚本剔除失配、损坏图片 └── stats.json # 通道均值/方差统计供预处理使用清洗脚本是我拿到新数据集后第一件做的事。水下数据集里常见两类问题一是配对错位input/0001.jpg和target/0001.jpg根本不是同一张图二是参考图本身质量差比如仍带轻微色偏或模糊。前者会导致网络学到错误的映射关系后者会把模型的输出拉向“中等偏绿”因为网络学会了在退化图和参考图之间取折中。# uieb_clean.py —— 配对检查与基础筛选 import cv2 import numpy as np import os, json input_dir water_dataset/input/train target_dir water_dataset/target/train drop_list [] for fname in sorted(os.listdir(input_dir)): in_path os.path.join(input_dir, fname) gt_path os.path.join(target_dir, fname) # 假设同名配对 if not os.path.exists(gt_path): drop_list.append(fname) continue img_in cv2.imread(in_path) img_gt cv2.imread(gt_path) if img_in is None or img_gt is None: drop_list.append(fname) continue if img_in.shape ! img_gt.shape: drop_list.append(fname) # 尺寸不一致也剔除 print(f待剔除 {len(drop_list)} 张:, drop_list[:10])这段脚本做的是最基础的清洗配对缺失、无法读取、尺寸不一致的样本直接剔除。注意img_in.shape ! img_gt.shape这个判断水下数据集中不同相机拍摄的样本很容易在分辨率上不一致如果你的网络包含跳跃连接或者需要固定尺寸输入失配样本会在训练时报错或者产生隐形的空间错位。2.3 白平衡与归一化两个影响训练效果的关键参数水下图像增强里面有一个容易被忽略的预处理——通道白平衡。虽然深度网络理论上能自己学习色偏校正但实际训练中输入色偏过于严重时网络会把大部分容量花在去色偏上而不是学习细节恢复。一个低成本的做法是先做灰度世界白平衡再进网络。# preprocess.py —— 灰度世界白平衡 归一化 def gray_world_balance(img): 灰度世界假设三个通道的均值应趋于一致 result img.astype(np.float32) avg_b np.mean(result[:, :, 0]) avg_g np.mean(result[:, :, 1]) avg_r np.mean(result[:, :, 2]) avg_gray (avg_b avg_g avg_r) / 3.0 # 按比例缩放各通道避免溢出 scale_b avg_gray / (avg_b 1e-6) scale_g avg_gray / (avg_g 1e-6) scale_r avg_gray / (avg_r 1e-6) result[:, :, 0] np.clip(result[:, :, 0] * scale_b, 0, 255) result[:, :, 1] np.clip(result[:, :, 1] * scale_g, 0, 255) result[:, :, 2] np.clip(result[:, :, 2] * scale_r, 0, 255) return result.astype(np.uint8)归一化这里有个血泪经验很多人直接把图像除以 255然后用transforms.Normalize做标准化时使用 ImageNet 的均值方差。这在自然图像分类上没问题但水下图像的通道分布和 ImageNet 差异极大直接套用会导致网络训练初期梯度方向偏斜。我一般会先用stats.json统计训练集的通道均值方差再写入 Normalize 参数。3. 模型选型与训练配置U-Net 基线到带对抗损失的增强3.1 用 U-Net 做回归而不是分类网络结构怎么定水下图像增强本质是像素级回归/映射任务输入一张退化图输出一张清晰图所以模型选择基于 U-Net 的编码器-解码器结构最稳妥。为什么不直接用 ResNetResNet 做分类表现好但做逐像素回归时会丢失空间细节尤其是水下图像的边缘和纹理信息。U-Net 的跳跃连接可以把编码器的浅层特征直接传给解码器保证输出图像的细节接近输入。# 简化版 U-Net核心是跳连结构 import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels3, out_channels3, base64): super().__init__() # 编码器 self.enc1 DoubleConv(in_channels, base) self.enc2 DoubleConv(base, base * 2) self.enc3 DoubleConv(base * 2, base * 4) self.pool nn.MaxPool2d(2) # 瓶颈 self.bottleneck DoubleConv(base * 4, base * 8) # 解码器通过跳连拼接特征 self.up3 nn.ConvTranspose2d(base * 8, base * 4, 2, stride2) self.dec3 DoubleConv(base * 8, base * 4) self.up2 nn.ConvTranspose2d(base * 4, base * 2, 2, stride2) self.dec2 DoubleConv(base * 4, base * 2) self.up1 nn.ConvTranspose2d(base * 2, base, 2, stride2) self.dec1 DoubleConv(base * 2, base) self.out nn.Conv2d(base, 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], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)这个结构里torch.cat把编码器特征和解码器上采样特征在通道维度拼接是 U-Net 的灵魂。base64是基础通道数显存不够时降到 32效果会损失一些但仍然可用。注意输出层用的是1x1卷积把通道数压回 RGB 三通道。3.2 损失函数组合MSE SSIM 还是增加对抗分支很多教程直接把 MSEL2损失当作所有图像生成任务的首选但单独用 L2 会让输出图像变“糊”——因为 L2 假设像素独立高斯分布它对多模态输出取均值落到水下增强场景就是纹理被抹平。我建议至少组合一个感知或结构损失最省事的是 L1 SSIM 的组合。# loss.py —— 组合损失L1 SSIM import torch import torch.nn.functional as F from pytorch_msssim import SSIM class CombinedLoss(nn.Module): def __init__(self, alpha0.8, window_size11): super().__init__() self.alpha alpha self.ssim SSIM(data_range1.0, size_averageTrue, channel3) def forward(self, pred, target): l1 F.l1_loss(pred, target) # 逐像素重建损失 ssim 1 - self.ssim(pred, target) # SSIM 越接近 1 越好所以取反 return self.alpha * l1 (1 - self.alpha) * ssimalpha0.8意味着重建损失占大头SSIM 作为结构约束。实践里 alpha 在 0.7~0.9 之间都可以如果你的输出图整体偏模糊把 alpha 调小、增加 SSIM 权重如果输出图颜色失真但纹理清晰把 alpha 调大。这属于深度学习里面“玄学”最多的环节但至少比盯着单一项 MSE 曲线靠谱。带对抗分支GAN的做法比如 UGAN 的变体会增加一个判别器来区分增强图和真实清晰图。在课程设计这种体量下GAN 训练稳定性是个大问题不建议作为主力方案如果要做创新点可以把它作为 U-Net 的后续微调阶段——先生成预训练模型再挂判别器微调 2000 个 iteration不要从头训。3.3 train.py 的完整参数表与单卡训练命令训练脚本核心参数如下表这些参数决定了一个模型能不能从“跑通”进化到“能用”。参数推荐值说明坑点image_size256x256输入分辨率过大会爆显存过小细节恢复差batch_size8单卡常用值显存不够先降 batch不要先降分辨率lr1e-4Adam 初始学习率1e-3 容易训练直接发散epochs100水下增强一般 80~150用早停按 val loss 截断weight_decay1e-5L2 正则设太大模型欠拟合num_workers4DataLoader 进程数Windows 下超过 8 可能报错# 单卡训练命令工程自带 train.py 的常见用法 python train.py \ --input_dir water_dataset/input/train \ --target_dir water_dataset/target/train \ --val_dir water_dataset/input/val \ --image_size 256 \ --batch_size 8 \ --lr 1e-4 \ --epochs 100 \ --model unet \ --loss l1_ssim \ --output_dir checkpoints/训练过程中我会在每个 epoch 结束时保存一次 checkpoint同时在 val 集上算一次 PSNR。这里有个注意点val 集必须和训练集来自同一分布但不要直接用 UIEB 的标准测试集做验证否则你会在评估阶段看到异常高的 PSNR因为训练时模型已经见过类似分布的图像。审计的时候也要把 val 指标和测试集指标分开记录。3.4 训练过程中怎么判断模型没跑偏光看训练 loss 是不够的。我见过太多训练 loss 平稳下降、val 指标忽高忽低的情况。一个有效率的方法是每个 epoch 把 val 集的增强结果拼成网格图保存到vis/目录肉眼检查色偏、纹理、过曝情况。深度学习训练是一个黑匣子但可视化输出至少能告诉你“黑匣子现在的偏好是什么”。另一个判断点是验证集上的 SSIM 和 PSNR 曲线的相关性。正常情况下两者在最后 20 个 epoch 应该同步上升如果 SSIM 升但 PSNR 降大概率是输出图开始过度平滑这时候你应该调整损失函数权重而不是继续训下去。4. 评估指标与 YOLO 检测衔接PSNR/SSIM/UIQM 之外还要看什么4.1 三个指标的计算脚本与读法水下图像增强的评估指标学术界常用 PSNR、SSIM还有一个水下专属的 UIQM。PSNR 是像素级误差指标SSIM 是结构相似度UIQM 不需要参考图直接从色彩、清晰度、对比度三个子项评分。毕业设计里如果你只用 PSNR 和 SSIM评委通常会追问“没有参考图怎么办”或者“这个分数高说明什么”。# evaluate.py —— 批量计算 PSNR 和 SSIM import cv2 import numpy as np from skimage.metrics import peak_signal_noise_ratio, structural_similarity def calc_psnr_ssim(pred_path, gt_path): pred cv2.imread(pred_path) gt cv2.imread(gt_path) # 尺寸不一致时先对齐 if pred.shape ! gt.shape: pred cv2.resize(pred, (gt.shape[1], gt.shape[0])) psnr peak_signal_noise_ratio(gt, pred) ssim structural_similarity(gt, pred, channel_axis2) return psnr, ssim注意channel_axis2这个参数OpenCV 读进来的图像是 HWC 格式通道维度在第三个轴。很多人用 skimage 计算 SSIM 时默认通道轴是最后一个写成channel_axis-1也没问题但如果你用灰度图模式multichannelFalse计算的 SSIM 和彩色图的结果会差很多对比不同模型时要保证评估代码一致。UIQM 的实现各版本权重略有差异工程里如果附带了 uiqm.py建议先跑一张标准测试图确认数值范围在 0~10 之间。我之前见过一个版本算出来的 UIQM 都在 100 以上后来发现是色彩维度归一化漏了除以图像最大值导致整个指标失去了可比性。4.2 把增强结果喂给 YOLO 做目标检测的衔接脚本水下增强通常不是终点很多毕业设计会加一个下游任务对增强后的图像做目标检测。这时候你的 U-Net 输出要能无缝接入 YOLO 的预测接口。常见的做法是把增强图保存成临时文件再用 YOLO 的检测脚本跑工程里如果提供这种衔接模块核心逻辑一般是这样的# detect_pipeline.py —— 先增强再检测 import cv2 import torch import numpy as np from model import UNet def enhance_then_detect(img_path, enhance_model, detect_model): img cv2.imread(img_path) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 归一化到 [0,1]转 tensor tensor torch.from_numpy(img_rgb / 255.0).permute(2, 0, 1).unsqueeze(0).float() with torch.no_grad(): enhanced enhance_model(tensor) # 输出范围 0~1 enhanced enhanced.squeeze(0).permute(1, 2, 0).numpy() enhanced np.clip(enhanced * 255.0, 0, 255).astype(np.uint8) # 交给 YOLO 检测 results detect_model(enhanced) return enhanced, results这里最容易出问题的衔接点是通道顺序和数值范围。U-Net 训练时如果用 RGB 顺序、输出范围 0~1推理时就要保证输入前做同样的转换YOLO 如果内部用 BGR你的增强图是 RGB直接喂进去会导致检测效果大幅下降。我一般会在代码里显式打印增强图的通道均值和类型确认格式一致再做检测。4.3 毕业设计答辩时的实验对比表怎么整理对比实验的整理思路决定了答辩效果。不要只放一张大表要分成三个层次传统算法基线直方图均衡、CLAHE、Retinex、纯 CNN 增强、增强后接检测的结果。表格格式建议如下方法PSNR(dB)SSIMUIQM检测 mAP原图直出无参考无参考3.1242.1%CLAHE 增强18.70.723.8545.3%U-Net(L1)24.30.844.3251.6%U-Net(L1SSIM)25.10.874.4153.8%检测 mAP 这一列是加分项它说明增强不是自嗨而是对下游任务真的有帮助。整理表格时注意把训练集分布说明白这一点很多人漏了如果你的模型在 UIEB 训练集上做的评估那 PSNR 再高也没有说服力评委一看数据集重叠就会追问。5. 避坑记录训练、指标、部署环节五个常见翻车点5.1 现象一训练 loss 在前 500 个 iteration 直接发散训练刚开始 loss 不降反升几轮后直接 NaN。原因通常是学习率过大或者输入数据没有归一化到 [0,1]导致梯度数值过大。解决先确认输入 tensor 的数值范围再用 lr1e-4 重跑。如果仍然发散检查是否有 NaN 像素混入数据——有些水下图片在保存时损坏读进来是空值。5.2 现象二增强结果整体偏紫红色输出图像的红色通道明显过强整体偏紫。原因是参考图本身存在色偏白平衡预处理只应用在了输入图上没有同步到参考图或者 loss 中 L1 权重过大模型倾向于复制参考图的颜色分布。解决在预处理阶段对 input 和 target 同时做白平衡并在验证阶段检查 target 的通道均值是否接近灰色。5.3 现象三同一个 checkpoint 在不同机器上评估结果不一致同一份测试集两台机器跑出来的 PSNR 差 0.5dB 以上。原因是 OpenCV 版本不同导致的读取差异——新版 OpenCV 默认处理 EXIF 方向老版本不处理图像被翻转或旋转后像素级评估自然对不上。解决在评估脚本里统一用cv2.IMREAD_UNCHANGED并且显式禁用 EXIF 方向变换同时在 README 里写清楚依赖版本。5.4 现象四YOLO 检测时提示维度错误增强模型的输出是 4D tensorYOLO 需要 HWC numpy 数组。原因是permute的顺序写错常见的是把(0, 2, 3, 1)写成了(0, 3, 2, 1)。解决按squeeze - permute(1, 2, 0) - numpy() - *255 - uint8的顺序处理每一步后打印 shape 确认。5.5 现象五训练速度越来越慢GPU 利用率不到 50%数据加载成了瓶颈。水下图像分辨率高读取和 resize 占用了大量 CPU 时间。解决把图片提前 resize 成 256x256 再存到磁盘或者在DataLoader里把num_workers调大到 8并设置pin_memoryTrue。还有一个省事技巧把整个训练集打包成一个.npy或.h5文件避免训练时频繁随机读取小文件。6. 从脚本到演示系统Tkinter 界面与模型导出6.1 做一个能选图、增强、保存的桌面演示毕业设计现场演示时命令行推理的效果远不如一个可视化界面来得直观。用 Tkinter 做前后端分离的最小演示能覆盖“选择图片 → 增强 → 对比展示 → 保存结果”四个核心动作。# app.py —— 精简版 Tkinter 演示界面 import tkinter as tk from tkinter import filedialog, Label, Button from PIL import Image, ImageTk import torch import cv2 import numpy as np def select_and_enhance(): path filedialog.askopenfilename(filetypes[(图像, *.jpg *.png)]) if not path: return img cv2.imread(path) # 同训练一致的白平衡 归一化省略细节 tensor preprocess(img) with torch.no_grad(): output model(tensor) result postprocess(output) # 转 HWC uint8 cv2.imwrite(enhanced_result.jpg, result) show_preview(img, result) def show_preview(orig, result): # 用 PIL 缩放后显示在 Label 上注意销毁旧图片避免内存泄漏 pass这个演示系统的注意点是界面刷新每点击一次按钮就要把两张图并排显示PIL 的ImageTk.PhotoImage不能完全依赖局部变量要保存引用否则图片显示一次后会变成空白——这是我做演示时翻车最多的地方现场演示变白屏非常尴尬。6.2 导出 ONNX 让模型脱离 PyTorch 环境如果你想在答辩现场用另一台没有 PyTorch 环境的机器做演示把模型导出成 ONNX 是一个稳妥路径。ONNX Runtime 的部署依赖轻CPU 上跑 256x256 的 U-Net 推理大约在 300ms 左右完全够用。# export_onnx.py —— 导出 ONNX 并验证输出 import torch from model import UNet model UNet() model.load_state_dict(torch.load(checkpoints/best.pth, map_locationcpu)) model.eval() dummy torch.randn(1, 3, 256, 256) torch.onnx.export( model, dummy, unet_water.onnx, input_names[input], output_names[output], opset_version11, dynamic_axes{input: {0: batch}, output: {0: batch}} )dynamic_axes设为 batch 维度动态这样部署时一次可以喂一张图也可以同时喂一批图。导出的 ONNX 用onnxruntime跑一遍对比 PyTorch 推理的输出差异超过 1e-4 就要检查是否有不支持的操作。从那以后我每次拿到一个新的水下模型都会强制走一遍完整链路清洗配对数据、用固定统计参数做预处理、训练勾选可视化、评估对齐读图方式、部署导出 ONNX。这套流程帮我避开了至少十次答辩现场翻车也节省了大量重复排查的时间。希望这篇笔记能帮你在复现或改造这套工程时少走点弯路把力气花在真正值得花的地方。本文还有配套的精品资源点击获取