ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

ViT图像去雾实战:从原理到源码复现的完整指南

ViT图像去雾实战:从原理到源码复现的完整指南 简介基于Vision Transformer的图像去雾算法研究与实现源码及文档说明Python高分项目面向计算机视觉方向的研究人员、算法工程师与毕业设计开发者聚焦Transformer架构在图像去雾任务中的数据处理、模型训练与结果评估全流程。资源包含完整的训练数据预处理脚本可将图像切分为256×256小图并附NH-HAZE数据集train/test目录划分说明方便复现去雾实验。包内共338个文件其中204个Python程序涵盖数据生成与训练逻辑16个YAML配置文件用于模型参数设定另有IPython Notebook调试记录、CSV指标统计及图表可视化文件压缩包约156.36MB。已有291人学习下载。读者可借助该资源直接理解Uformer等Vision Transformer去雾模型的工程实现获取数据增强、训练脚本配置与损失景观分析等一手资料适合作为课程设计、论文实验或算法二次开发的参照基线。1. 一件“雾天识别失败”的真实事故和这个去雾项目能帮你什么某园区智慧安防系统在一个夏季清晨的浓雾里把停在路边的卡车识别成了墙体——输入图像被雾层污染检测模型当场失效。图像去雾算法要处理的就是这类退化场景从带雾图像I(x)还原清晰内容J(x)。早期方案依赖大气散射物理模型近几年CNN端到端方法已经成熟而Vision TransformerViT的全局注意力机制进一步解决了卷积感受野不足导致的远距离雾层建模缺失因为雾的浓度分布本质依赖整幅图的全局统计。这个标题里的“源码文档说明python高分项目”说明它是一个能直接跑通的项目包适合课程设计、毕业设计起步或算法方向验证。要把它真正跑赢需要从原理、代码、运行时三个层面去拆。2. 从CNN到ViT图像去雾模型为什么需要全局注意力2.1 CNN去雾的边界局部卷积够不到远距离雾层图像去雾早期是CNN的阵地。DehazeNet用多尺度卷积估计透射率AOD-Net把物理模型折叠进一个轻量网络直接输出清晰图MSCNN则在粗糙尺度上预测透射率再逐步细化。这些结构有个共同特点每个位置的输出只依赖卷积核覆盖的局部邻域虽然可以靠加深网络把理论感受野撑大但实际训练后网络对远端信息的利用效率远低于理论值。这个弱点在雾天场景里被明显放大因为雾的散射过程天然是全局的——远处物体透过雾层传到相机时信息衰减程度和场景深度强相关模型需要知道“远”和“近”才能正确估计不同位置的雾浓度。ViT的切入方式完全不同它把输入图像切成固定大小的patch每个patch展平成token再用自注意力让任意两个token之间直接交互。这样一来任何一个patch的表示都能看到全图所有位置注意力权重由数据动态学出来。对去雾任务来说模型可以自动学会“用天空区域的亮度和颜色来约束近处物体区域的透射率估计”不需要靠卷积层一层层把远距离信息搬运过来。这个全局建模能力在浓雾大范围分布的场景里比CNN更容易形成一致的估计。直观对比来看处理一张1024×1024的雾图CNN在浅层感受野只有几十像素而ViT即使在最浅层每个patch也已经聚合了全图信息。这也是为什么很多去雾评测榜单里ViT系方法在浓雾和深度变化剧烈的图像上普遍比CNN方法高出一截PSNR差距主要来自全局上下文而不是参数量大小。2.2 雾天退化物理模型与数据驱动建模范式去雾问题的数学起点是大气散射模型I(x) J(x) × t(x) A × (1 - t(x))其中I(x)是观测到的带雾图像J(x)是要恢复的清晰图像t(x)是透射率A是全局大气光。传统去雾方法如暗通道先验的思路是先估计t(x)和A再代入上式反解出J(x)。这类方法的强假设是“无雾图像的局部块里至少有一个通道的像素值很低”一旦遇到天空、白墙、浓雾这类违背假设的区域估计出的透射率图就会出错恢复结果出现颜色失真或光晕。数据驱动方法跳过了手工先验直接学习从I(x)到J(x)的映射。训练数据是成对的清晰图和合成雾图模型输入带雾图回归输出清晰图。好处是不再依赖暗通道假设代价是需要大量成对数据并且模型学到的映射高度依赖合成雾图的方式。你这个项目源码走的就是这条路线大概率使用RESIDE数据集或自己写合成脚本生成训练对具体是哪一种看data目录下有没有synth相关函数就能判断。ViT在数据驱动去雾里通常有两种角色一是作为主干网络直接回归清晰图二是作为物理模型参数的估计器先输出透射率图和大气光再用散射模型重建清晰图。直接回归路线训练简单但模型要隐式学会估计透射率和大气光物理模型约束路线可解释性强但两个子任务的误差会叠加需要额外的中间监督。源码里如果解码器有两个head分别输出透射率和清晰图走的就是后一条。2.3 “物理模型ViT”和“纯端到端ViT”这对路线选哪个判断你手里的源码属于哪条路线看模型输出层就行。纯端到端路线最后一层通常是3通道输出直接对应RGB清晰图物理模型路线要么是1通道透射率加3通道大气光要么是透射率图经过一个重建分支得到清晰图。判断维度纯端到端物理模型ViT训练难度较低单任务回归较高多任务需平衡可解释性黑匣子中间结果可解释合成数据泛化依赖数据分布物理约束泛化更好适用场景课程设计、快速验证论文研究、真实场景落地课程设计选纯端到端路线能在更短时间内跑出好看的效果图目标是发论文或真实场景部署物理模型ViT有更多可讲的改进故事也更容易应对合成雾到真实雾的域差距。2.4 位置编码和token空间顺序去雾结果空间一致性的来源ViT的自注意力运算对token的排列是置换等变的——如果完全去掉位置编码模型无法知道每个token来自图像的哪个区域。图像去雾对位置信息非常敏感天空通常在图像上部近景在下部远处雾层和近处雾层的透射率衰减规律也不同。位置编码就是模型区分这些空间语义的唯一依据。源码里位置编码通常是一个可学习的nn.Parameter初始化是随机值需要靠训练拟合。有一个常见坑从ImageNet预训练权重迁移时原位置编码表是按224×224的patch网格学出来的换到你的训练分辨率后如果直接插值resize要确保插值后的网格和patch划分方式对齐。不对齐的表现是模型前几轮训练出奇的慢且去雾效果在不同图像区域出现不对称伪影。3. 解压源码后按这个顺序读代码目录、数据管线与模型核心拿到zip压缩包先别急着跑train.py花20分钟理解目录结构能避免后面90%的报错。3.1 项目目录里每个文件/文件夹的角色一个典型的ViT去雾项目目录长这样├── data/ │ ├── __init__.py │ ├── dataset.py # 数据集类定义负责读图和配对 │ └── transforms.py # 数据增强与合成雾图 ├── models/ │ ├── __init__.py │ ├── vit.py # Vision Transformer主干 │ ├── decoder.py # 解码器从特征回归清晰图 │ └── losses.py # 损失函数集合 ├── utils/ │ ├── metrics.py # PSNR / SSIM 计算 │ └── visualization.py # 结果保存与对比图 ├── configs/ │ └── train.yaml # 训练超参数配置 ├── train.py # 训练入口 ├── test.py # 推理与评估入口 ├── requirements.txt └── README.mdREADME.md里通常会写训练命令、数据集格式和预训练权重的下载方式先看它能省掉一半排错时间。utils/metrics.py负责PSNR和SSIM计算论文里报告的数值就是这里产出的后文会讲它的口径问题。configs目录用yaml管理超参数比直接散落在train.py里更适合做多组对比实验。提示如果项目没有configs目录超参数会集中在train.py顶部改动时要全局搜索一遍避免改了训练集路径却漏了验证集路径。3.2 合成有雾图像的预处理逻辑配对的雾图是怎么来的去雾训练需要成对数据清晰图雾图真实场景几乎拍不到严格配对的图像所以主流做法是用大气散射模型自己合成雾图。合成逻辑的核心就一句用深度图算透射率t再和全局大气光A做线性混合。# data/transforms.py节选 import numpy as np def synthesize_fog(clean_img, depth_map, A220, beta0.15): 按大气散射模型合成雾图 clean_img: HxWx3像素范围[0,1] depth_map: HxW深度值范围[0,1] A: 大气光强度常用[200,240] beta: 散射系数控制雾浓度 t np.exp(-beta * depth_map) # 透射率随深度指数衰减 fog clean_img * t[:, :, None] A * (1 - t[:, :, None]) return np.clip(fog, 0.0, 1.0)逻辑说明透射率从1衰减到接近0深度越大的像素雾越浓这与真实物理过程一致。A取值在210到240之间模拟亮灰色天空beta取值0.1到0.2分别对应薄雾到浓雾。深度图质量直接决定合成雾图质量如果深度图带块状伪影模型会学到伪影而不是去雾能力。3.3 Patch Embedding与Transformer Encoder核心模型代码怎么读ViT主体由两部分组成Patch Embedding把图像切成patch并映射成token序列Transformer Encoder在token之间做自注意力。下面这段精简代码和源码里的vit.py结构一致可以直接对照着读# models/vit.py节选 import torch import torch.nn as nn class PatchEmbed(nn.Module): 把B,C,H,W的图像映射为B,num_patches,embed_dim的token序列 def __init__(self, img_size224, patch_size16, in_channels3, embed_dim512): super().__init__() self.num_patches (img_size // patch_size) ** 2 # 用一个stridepatch_size的卷积同时完成切分和映射 self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): x self.proj(x) # (B, embed_dim, H/p, W/p) x x.flatten(2).transpose(1, 2) # (B, num_patches, embed_dim) return x class TransformerEncoderBlock(nn.Module): def __init__(self, embed_dim512, num_heads8, mlp_ratio4.0, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention(embed_dim, num_heads, dropoutdropout) self.norm2 nn.LayerNorm(embed_dim) self.mlp nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Linear(int(embed_dim * mlp_ratio), embed_dim), nn.Dropout(dropout) ) def forward(self, x): # Pre-LN结构先归一化再做注意力深层训练更稳定 x x self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x x self.mlp(self.norm2(x)) return x逻辑说明PatchEmbed用一个stride等于kernel_size的卷积完成patch切分和线性映射224×224输入、patch16时产出196个token。Encoder Block用的是Pre-LN加残差的结构这是DeiT、Swin等视觉Transformer统一采用的做法。去雾模型的解码器需要把token序列还原成224×224×3的图像常见做法是转置卷积加几个残差卷积块做细节恢复。参数说明embed_dim在去雾任务里常用256到512之间num_heads取8或16必须能整除embed_dimmlp_ratio固定为4.0能使参数量和计算量保持合理配比。patch_size是最关键的旋钮第5章会专门展开。3.4 从数据流形状读代码一个Batch在模型里怎么流动读ViT代码时最容易迷失的是张量形状尤其是从序列重新reshape回二维网格那一步。我用一个极短的验证脚本把形状变化走一遍# 快速验证数据流形状 with torch.no_grad(): x torch.randn(2, 3, 224, 224).cuda() x model.patch_embed(x) # (2, 196, 512) x model.pos_drop(x) # (2, 196, 512) for block in model.blocks: x block(x) # 每层保持 (2, 196, 512) x model.norm(x) # (2, 196, 512) x x.transpose(1, 2) # (2, 512, 196) x x.view(2, 512, 14, 14) # 重新排回二维网格 pred model.decoder(x) # (2, 3, 224, 224)逻辑说明transpose后必须要view回(B, C, H/p, W/p)H/p和W/p的乘积必须等于token数量。如果patch_size和img_size不整除比如224/15这里的view必然报错。遇到shape mismatch时就顺着这段数据流逐行断点排查比看报错信息猜快得多。3.5 损失函数与训练循环的关键代码去雾训练的损失不能只有MSE。只用L1或MSE收敛快但恢复图偏平滑纹理细节容易被平均掉。高分项目里的损失通常是加权组合# models/losses.py节选 import torch.nn.functional as F class DehazeLoss(nn.Module): 组合损失L1 感知损失 边缘保持损失 def __init__(self, l1_weight1.0, percep_weight0.1, edge_weight0.05): super().__init__() self.percep PerceptualLoss() # 常用ImageNet预训练的VGG16前几层 self.l1_weight l1_weight self.percep_weight percep_weight self.edge_weight edge_weight def forward(self, pred, target): l1 F.l1_loss(pred, target) percep self.percep(pred, target) # 边缘损失比较梯度图的L1让边缘保持锐利 pred_grad torch.abs(pred[:, :, :-1, :] - pred[:, :, 1:, :]) target_grad torch.abs(target[:, :, :-1, :] - target[:, :, 1:, :]) edge F.l1_loss(pred_grad, target_grad) return (self.l1_weight * l1 self.percep_weight * percep self.edge_weight * edge)训练循环本身不复杂标准写法里加一个梯度累积就能在显存不足时模拟更大batch# train.py节选 optimizer.zero_grad() for batch_idx, (fog, clear) in enumerate(train_loader): fog, clear fog.cuda(), clear.cuda() pred model(fog) loss criterion(pred, clear) / accumulate_steps # 按累积步数均分 loss.backward() if (batch_idx 1) % accumulate_steps 0: optimizer.step() optimizer.zero_grad()梯度累积的意思是把一个大步拆成多个小步来凑显存accumulate_steps4时等效batch_size放大4倍但学习率也要相应微调。这个技巧在ViT这种大模型上去雾任务里很常用。4. 复现跑通全流程环境配置、训练命令与参数调优代码读通了接下来解决“怎么让它跑起来”。这一章覆盖环境到评估的完整路径。4.1 环境准备Python版本与依赖安装ViT去雾项目依赖PyTorch生态这不是靠纯NumPy能跑的。requirements.txt里通常包含以下关键包pip install torch1.10 pip install torchvision pip install numpy opencv-python pip install pyyaml pip install scikit-image pip install tqdmtorch和torchvision必须先装并且版本要匹配。机器有GPU的话先确认驱动支持的CUDA版本再选对应的安装源conda创建环境是最省事的做法conda create -n dehaze python3.8 conda activate dehaze pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118cu118表示CUDA 11.8的预编译wheel机器是CUDA 12.x就把cu118改成cu121。这一步做错的表现是import torch时报libcuda.so找不到然后整个环境不可用。无GPU的机器去掉--index-url直接装CPU版训练慢一个数量级至少能先验证代码逻辑通畅。4.2 最简训练启动命令与每个超参数的效果环境就绪后最简启动命令是python train.py --config configs/train.yamltrain.yaml里的核心超参数长这样data_root: ./data/RESIDE batch_size: 8 num_workers: 4 epochs: 100 lr: 1e-4 min_lr: 1e-6 warmup_epochs: 5 img_size: 224 patch_size: 16 embed_dim: 512 depth: 12 num_heads: 8img_size与patch_size配合决定token数量224/161414×14196个token。embed_dim512、depth12大约是ViT-Small级别加解码器后完整模型能放进24G显存如果是12G卡batch_size降到4embed_dim降到384比较稳。lr1e-4是AdamW在ViT上的常见起点学习率太高会出现前几个epoch的loss直接跳nan太低则收敛慢。warmup_epochs5让学习率线性升到峰值能极大避免前期震荡——ViT对学习率敏感是出了名的不少翻车都来自跳过warmup。断点续训在后处理时是后悔药一样的存在。保存checkpoint时把优化器状态和best指标一起存torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_psnr: best_psnr, }, checkpoints/best.pth)恢复训练时加载这个字典再手动把model和optimizer的state_dict灌回去。别只存model权重否则中途停了就得从epoch0重跑。4.3 换到自己的数据集上微调至少要改的3处自定义数据集时至少有三处代码必须动第一处是dataset.py的构造函数。把图像按清晰图和雾图两个文件夹放好用文件名排序配对# data/dataset.py节选 class DehazeDataset: def __init__(self, clear_dir, foggy_dir, img_size224, trainTrue): self.clear_list sorted(glob(clear_dir /*.png)) self.foggy_list sorted(glob(foggy_dir /*.png)) self.train train self.img_size img_size文件名必须能一一对应否则训练时模型输入的“清晰图”和“雾图”根本是两幅不同内容loss数值会表现得很怪但不报错这是自定义数据里最容易踩的坑。第二处是数据增强。从“直接裁剪”改成“随机翻转颜色抖动随机裁剪”能提升泛化if self.train: fog_transform Compose([ RandomHorizontalFlip(), RandomCrop(self.img_size), ToTensor(), Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])注意Normalize用的ImageNet统计值模型输出层要对应做反归一化否则恢复图整体偏灰。第三处是评估部分的数据范围。skimage的PSNR/SSIM需要显式指定data_rangefrom skimage.metrics import peak_signal_noise_ratio as psnr from skimage.metrics import structural_similarity as ssim def calc_metrics(pred, target): # 数据在[0,255]范围就必须传data_range255 p psnr(pred, target, data_range255) s ssim(pred, target, channel_axis-1, data_range255) return p, sdata_range用错的话PSNR会虚高或虚低数值完全失真。4.4 用PSNR和SSIM验证结果评估脚本的写法test.py的核心逻辑是逐张推理、算指标、存对比图model.eval() psnr_sum 0.0 with torch.no_grad(): for fog, clear in test_loader: pred model(fog.cuda()).cpu() pred_img pred.clamp(0, 1) * 255 clear_img clear.clamp(0, 1) * 255 p, s calc_metrics(pred_img.numpy(), clear_img.numpy()) psnr_sum p可视化图同样重要PSNR只反映像素级差异雾区残留和颜色偏色必须靠眼睛确认import matplotlib.pyplot as plt fig, axes plt.subplots(1, 3) axes[0].imshow(fog.permute(1, 2, 0)); axes[0].set_title(Foggy) axes[1].imshow(pred.permute(1, 2, 0)); axes[1].set_title(Dehazed) axes[2].imshow(clear.permute(1, 2, 0)); axes[2].set_title(Clear) fig.savefig(fvis/result_{idx}.png, dpi150)评估口径要和论文基准保持一致多数论文对RGB三通道算PSNR、对灰度图算SSIM也有的全部按灰度算。复现出来的数字和论文差很多时先查口径而不是怀疑模型。5. 避坑指南ViT去雾项目里最容易翻车的5个细节5.1 训练Loss不降反升先看数据归一化区间现象loss前几个epoch正常下降第3个epoch后开始抖动上升最终稳定在比初始还高的水平。原因输入图像归一化到[0,1]模型输出层却没有做逆归一化模型和loss在完全不同的尺度空间里计算。更隐蔽的情况是Normalize用了ImageNet的mean/std输出侧却按[0,1]尺度处理。解决在train.py的循环里加一个检查打印model输出和target的实际数值范围确认两者一致。如果预测值均值始终比真值小0.5左右就是反归一化缺失。排查方式是临时在loss后面加一行print(pred.min(), pred.max(), target.min(), target.max())一眼就能看出来。5.2 显存溢出patch_size与batch_size的连锁反应现象训练到第100步左右CUDA OOM把batch_size调小后又跑一会儿再次OOM。原因patch_size从16改成8token数量变成原来的4倍224/82828×28784自注意力的显存开销随token数量近似平方增长。只改了patch_size一个参数整个注意力矩阵的显存占用就翻了4倍batch_size再不降必然溢出。解决改patch_size必须同步调batch_size。起手配置建议patch_size16、batch_size8、embed_dim512在24G卡上能跑换patch_size8时batch_size先降到2还紧张就把embed_dim降到384。用torch.cuda.max_memory_allocated()记录每步峰值亲眼确认瓶颈在哪。5.3 patch_size调大后雾反而去不干净现象patch_size从16调到32重新训练后浓雾区的残留明显变多恢复图细节也发糊。原因patch越大每个token覆盖的像素区域越大空间细节在embedding阶段就被压缩了。去雾任务的最终输出是逐像素清晰图细节恢复高度依赖patch级信息。patch32时每个token覆盖32×32像素薄雾区域的纹理已经没有独立token了。解决保持patch_size16或更小。想要减小序列长度省显存优先降embed_dim或depth不要动patch_size。希望同时获得全局语义和局部细节可以加一条CNN分支处理原分辨率特征再和ViT输出融合这也是Swin Transformer的简化思路。5.4 PSNR虚高评估流程里的三个陷阱现象测试集PSNR报到35dB以上但把图放大看纹理和雾区细节依然不自然。原因一测试图像与训练图像存在裁剪重叠模型记住了重复内容常见于按固定crop位置切分数据集的实现。原因二合成雾图的beta和A取值在训练与测试中完全相同模型学到了固定映射换一种雾浓度分布立刻露馅。原因三计算指标时有实现会偷偷裁剪图像边界因为边缘patch自注意力质量通常偏低裁掉后数值自然更好看。解决训练与测试按图像ID严格分开不允许有重叠裁剪测试合成参数做随机化报告PSNR和SSIM时注明是否带边界裁剪并保持和对比方法同口径。验证自己的评测脚本可以拿一张已知图像算一遍确认数值和参考实现一致。5.5 加载预训练权重时key不匹配现象训练脚本报Missing key(s) in state_dict或某几个权重shape对不上。原因ViT通常在ImageNet上预训练它的分类头是1000类输出去雾模型把输出换成了3通道回归分类头自然加载不了。不同实现里LayerNorm的key命名也可能不一样比如blocks.0.norm1和encoder.layers.0.norm1。解决用strictFalse加载主干权重然后把形状不匹配的层过滤掉。更稳妥的做法是写一个过滤函数只加载key前缀匹配且shape一致的权重# 加载预训练权重时过滤不匹配的key state_dict torch.load(vit_base.pth, map_locationcpu) model_state model.state_dict() matched {k: v for k, v in state_dict.items() if k in model_state and model_state[k].shape v.shape} model.load_state_dict(matched, strictFalse)打印一下漏掉的key列表确认都是分类头和位置编码这类预期中的层。不做这一步前期收敛速度会非常慢这不是学习率问题是初始化问题。6. 从复现到加分三个让项目有论文味的改进方向能稳定跑通训练和测试之后下一步的差异化可以从这三个方向切入。轻量化ViT是一个实用方向。蒸馏或降维的常见做法是把depth从12降到6观察PSNR变化通常会有1到2dB的损失但参数量几乎减半。还可以把embed_dim从512降到384这个改动对全局建模能力的影响比减层更小。操作方式很简单改掉train.yaml里对应字段重新训练画出“参数量-推理帧率-PSNR”曲线这本身就是消融实验素材。多任务学习也值得考虑。去雾模型学到的特征里天然包含场景深度信息可以在Transformer编码器之后引出两个解码分支一个回归清晰图一个做语义分割共享同一套ViT主干。分割任务提供的场景先验能反过来约束去雾特征的语义一致性雾区边界处的效果提升尤其明显。损失函数调整为L_dehaze λ * L_segλ从0.1开始试通常在0.1到0.3之间有效。真实雾图泛化是最后一道坎。合成雾和真实雾的最大差距在于真实雾的深度分布不规则还带传感器噪声和雨滴干扰。拿到少量没有清晰参考的真实雾图后不要直接微调先用合成数据训练好的模型去预测伪清晰图把真实雾图和伪清晰图组成伪训练对迭代两三轮。每轮结束人工筛掉模型输出特别差的样本能明显缓解真实场景的雾区残留。我自己最早做ViT去雾时把patch_size当成一个随便填的数字从8到32试了一遍在16G卡上等了一周换来一张效果不如CNN的对比图。后来才悟到这个项目的核心调试对象是patch_size、embed_dim、depth三个变量而不是盲目调学习率。每改一个变量就把显存占用、指标数值、视觉效果三点同时记录否则实验等于白做。希望帮到你。本文还有配套的精品资源点击获取
返回列表