简介面向计算机视觉与深度学习方向的图像去雾研究者这份压缩包提供基于VisionTransformer的图像去雾算法Python实现、项目说明与配套数据集覆盖从数据预处理到模型训练的主要环节尤其适合毕业设计、课程设计或课题复现。源码包含Uformer架构的完整训练流程并给出将NH-HAZE训练数据切分为256×256小图的预处理脚本以及My_train.py训练入口、yaml配置、csv结果记录和ipynb分析笔记项目说明中还对数据存放路径、训练批次大小与轮数等参数作了交代便于按自己环境调整从数据准备、参数设置到模型训练与效果对比均可逐步跟进。资源共340个文件以204个py脚本为主另有png/gif图像示例、sh辅助脚本、md/txt说明文档等包体约156.42MB目录按train/test、脚本、配置与说明区分结构清晰。已有343人学习下载借助其中的预处理与训练代码可省去大量数据整理和工程搭建时间是一份完整可跑的图像去雾参考资料。1. 基于VisionTransformer的图像去雾为什么说它比CNN更适合吃透雾气图像去雾这个任务第一眼看过去像个“图像增强”问题真正做起来才发现它最难的在于“从局部里判断全局”。雾霾在图像里是平滑渐变的左上角和右下角常常被同一层雾气罩着局部卷积核看到的是白茫茫一片根本分不清是雾还是白墙。VisionTransformerViT的全局自注意力天然就是把整张图的上下文拉在一起算所以用基于VisionTransformer的图像去雾算法做研究、写Python源码、配数据集做训练这条路是站得住的。这篇文章我会按从原理到落地的顺序把大气散射模型、数据集整理、模型搭建、训练调试和指标验证这五步讲完。适合正在做毕业设计的学生也适合手头有雾图数据但用CNN效果不满意的一线工程师。2. 去雾任务建模与ViT选型从大气散射模型到全局感受野2.1 大气散射模型与去雾的全局特性绝大多数去雾算法都建立在Koschmieder大气散射模型上表达式是I(x) J(x) * t(x) A * (1 - t(x))其中I(x)是有雾图J(x)是清晰图A是全局大气光t(x)是透射率。做去雾本质上是已知I(x)估计出J(x)。但估计过程会被两个全局量卡住A是全图共享的一个值t(x)虽然逐像素变化但它的空间分布非常平滑——因为雾浓度在物理上是连续变化的不会像物体边缘那样剧烈突变。这意味着模型必须有能力把远处像素的信息传到近处来。CNN理论上能用大卷积核或空洞卷积扩大感受野但深层卷积的感受野是“近似”全局的实际训练中很难真正覆盖256x256甚至更大尺寸的全部位置。ViT就不一样它在第一层就把图像切成patch通过self-attention直接计算任意两个patch之间的相关性。大气光的估计、透射率图的平滑恢复以及被雾覆盖区域的纹理重建都需要这种“远距离参照”能力。所以用ViT做去雾不是赶时髦而是任务本身的物理特性决定的。你会发现很多去雾论文里都在强调“全局上下文信息”这就是ViT的结构优势所在。2.2 ViT结构如何匹配去雾特征VisionTransformer最早是给图像分类设计的它把图像分成N个patch拉平成token序列经过若干层Transformer Encoder后取class token或平均池化做分类。直接拿来去做像素级回归任务会遇到两个问题一是分类模型最后只输出一个向量去雾需要输出一张和输入尺寸相同的RGB图。常见做法是砍掉ViT的分类头把Encoder输出的token序列重新reshape成特征图再接一个解码器上采样回原分辨率。这个解码器可以是简单的转置卷积堆叠也可以是UNet式的跳跃连接结构。二是ViT没有卷积的局部先验对高频细节比如纹理、边缘的恢复能力偏弱。我在实际项目里一般会给ViT编码器配一个轻量的卷积解码器或者在进入Transformer之前先用一层卷积把低层特征拎出来这样既保留全局建模能力又不会让画面发糊。还有一个更贴近去雾的做法不直接回归清晰图而是让ViT预测透射率t和大气光A再套用大气散射模型反解出J(x)。这样物理约束更强模型需要学习的映射更简单但是t和A的估计误差会直接放大到输出上调试起来比较头疼。多数人第一次做的时候还是直接回归清晰图更省事先跑通再优化。2.3 输入输出设计从分类头改造成去雾头的关键代码如果你第一次把ViT接到去雾任务上我建议先不要上现成的大模型库而是自己写一个最小可运行的PatchEmbedding和Encoder这样参数在哪、显存花在哪心里都有数。import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, in_channels3, embed_dim256, patch_size16, img_size256): super().__init__() self.patch_size patch_size self.num_patches (img_size // patch_size) ** 2 # 用卷积做patch切分比手动reshape更干净 self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: [B, 3, H, W] x self.proj(x) # [B, embed_dim, H/p, W/p] x x.flatten(2) # [B, embed_dim, num_patches] x x.transpose(1, 2) # [B, num_patches, embed_dim] return x这段代码的关键点是patch_size决定了序列长度。输入256x256patch_size16序列长度就是256个tokenembed_dim是每个token的特征维度。patch_size越大序列越短显存越小但空间细节丢失越严重patch_size越小序列越长注意力计算量呈平方增长。得到token序列后后面再接标准的TransformerEncoderLayer即可。去雾头只需要在原模型最后加一个上采样解码器把序列还原成图像尺寸。常见做法是先把token序列reshape成特征图再用PixelShuffle或转置卷积放大到原图分辨率。提示第一次调试时不要追求大模型。embed_dim128Transformer层数3层在256x256输入上先把流程跑通再谈优化。3. 数据集准备RESIDE之外你还需要做什么3.1 选择哪种数据集与目录结构做去雾研究最常用的是RESIDE系列包含合成雾图OTS和ITS也有真实雾图RTTS。合成雾的方式是拿清晰图和深度图按大气散射模型人工加雾所以有标准的GTGround Truth标签。如果你手头拿到的压缩包里有train/hazy、train/clean这样的目录那就是最常见的配对结构如果没有需要自己整理成这种格式。我一般会把数据集分成三个目录dataset/ train/ hazy/ # 有雾输入图.png或.jpg clean/ # 对应的清晰图文件名一一对应 val/ hazy/ clean/ test/ hazy/ clean/为什么强调文件名一一对应因为加载时要按名字配对。很多开源数据集里hazy和clean的文件名并不完全相同有的多了后缀有的用了不同扩展名。所以第一步不是写模型而是先把数据对整理干净否则后面Dataset怎么写都会出问题。真实雾数据集NH-HAZE和Dense-Haze也值得关注它们的雾是真实场景拍摄的没有合成雾那种“假平滑”模型在真实雾上的泛化能力一看便知。但这类数据集通常量小单张分辨率还很高直接训练容易过拟合建议先在RESIDE上预训练再在真实雾图上微调。3.2 数据划分、读取与增强脚本下面这个脚本做两件事把hazy和clean配对按比例划分训练集和验证集然后写成一个Dataset类供PyTorch加载。import random from pathlib import Path from PIL import Image from torch.utils.data import Dataset def build_pairs(hazy_dir, clean_dir, suffix.png): hazy_files {p.stem: p for p in Path(hazy_dir).glob(f*{suffix})} clean_files {p.stem: p for p in Path(clean_dir).glob(f*{suffix})} common sorted(hazy_files.keys() clean_files.keys()) pairs [(hazy_files[k], clean_files[k]) for k in common] random.shuffle(pairs) return pairs class DehazeDataset(Dataset): def __init__(self, pairs, size256, trainTrue): self.pairs pairs self.size size self.train train def __len__(self): return len(self.pairs) def __getitem__(self, idx): hazy_path, clean_path self.pairs[idx] hazy Image.open(hazy_path).convert(RGB) clean Image.open(clean_path).convert(RGB) # 随机裁剪和翻转增强空间不变性 if self.train: i random.randint(0, hazy.width - self.size) j random.randint(0, hazy.height - self.size) hazy hazy.crop((i, j, i self.size, j self.size)) clean clean.crop((i, j, i self.size, j self.size)) if random.random() 0.5: hazy hazy.transpose(Image.FLIP_LEFT_RIGHT) clean clean.transpose(Image.FLIP_LEFT_RIGHT) else: hazy hazy.resize((self.size, self.size)) clean clean.resize((self.size, self.size)) return (transforms.ToTensor()(hazy), transforms.ToTensor()(clean))逻辑说明build_pairs用的是stem也就是去掉扩展名后的文件名防止.jpg和.png混在一起时配不上对。shuffle是为了划分前打乱顺序。数据增强用了随机裁剪和水平翻转这两个对去雾任务非常有效因为雾的空间分布和图像内容解耦翻转不会破坏物理模型。参数上size256是输入分辨率。ViT的position embedding是按固定分辨率计算的训练时用256推理时如果换更大分辨率位置编码需要插值否则效果会明显下降。这一点在后文坑位里会专门讲。3.3 训练集大小、图像尺寸和patch size的匹配ViT是数据饥饿型模型它的归纳偏置弱在几千张小图上训练很容易欠拟合。如果没有ImageNet预训练权重建议把patch size调到8可以让模型在有限数据下看到更多局部结构。我自己常用的参数组合如下参数推荐值说明输入图像尺寸256x256过大显存压力大过小细节丢失patch size8或168适合小数据集16适合预训练模型embed_dim256再大训练速度明显下降Transformer层数4-6超过8层在小数据集上容易过拟合数据量最少2000对低于这个量建议用ImageNet预训练如果你拿到的数据集只有几百对雾图不要硬从头训练。常见的做法是加载一个ImageNet上预训练的ViT模型作为编码器只训练解码器部分或者全模型低学习率微调。这个技巧能让收敛速度快很多损失曲线也稳定得多。4. 基于VisionTransformer的去雾网络实现核心模块与最小训练闭环4.1 从零搭建ViT去雾模型下面是一个最小可用的ViT去雾模型包含PatchEmbedding、TransformerEncoder、解码器三个部分。省去了位置编码插值等边界处理先跑通为主。import torch import torch.nn as nn class VitDehaze(nn.Module): def __init__(self, img_size256, patch_size16, in_channels3, embed_dim256, depth4, num_heads8, decoder_dim128): super().__init__() self.patch_embed PatchEmbedding(in_channels, embed_dim, patch_size, img_size) self.pos_embed nn.Parameter( torch.randn(1, (img_size // patch_size) ** 2, embed_dim)) encoder_layer nn.TransformerEncoderLayer( d_modelembed_dim, nheadnum_heads, dim_feedforwardembed_dim * 4, activationgelu) self.encoder nn.TransformerEncoder(encoder_layer, num_layersdepth) # 解码器把token序列还原成特征图再上采样 self.decoder nn.Sequential( nn.ConvTranspose2d(embed_dim, decoder_dim, kernel_sizepatch_size, stridepatch_size), nn.ReLU(), nn.Conv2d(decoder_dim, decoder_dim, 3, padding1), nn.ReLU(), nn.Conv2d(decoder_dim, 3, 3, padding1), nn.Sigmoid() ) def forward(self, x): B, C, H, W x.shape x self.patch_embed(x) # [B, num_patches, embed_dim] x x self.pos_embed x self.encoder(x) # [B, num_patches, embed_dim] # reshape回二维特征图 patch_h patch_w int(x.shape[1] ** 0.5) x x.transpose(1, 2).reshape(B, -1, patch_h, patch_w) x self.decoder(x) return x这里的TransformerEncoderLayer是PyTorch自带的不用自己实现多头注意力。关键参数是d_model必须和embed_dim一致dim_feedforward控制FFN宽度一般取embed_dim的4倍。前向传播里reshape这一步要注意patch_h和patch_w必须相等所以输入图像的长宽要一致否则序列长度开方不是整数代码直接报错。解码器最前面用了一个转置卷积kernel_size和stride都等于patch_size正好把patch_size分辨率的特征图上采样回原始输入尺寸。最后接Sigmoid是因为输入图像归一化到0-1区间输出也限制在0-1区间比较稳妥。4.2 训练循环与损失函数组合去雾的损失函数不能只用L2否则输出会偏向平均值产生“雾没去干净但整体颜色变淡”的问题。我一般会组合三种损失import torch.nn.functional as F def charbonnier_loss(pred, target, eps1e-6): return torch.mean(torch.sqrt((pred - target) ** 2 eps)) def dehaze_loss(pred, target, ssimNone): l1 F.l1_loss(pred, target) char charbonnier_loss(pred, target) # ssim模块需要单独安装没有就先注释掉 loss l1 * 0.6 char * 0.4 if ssim is not None: loss loss ssim(pred, target) * 0.2 return lossL1损失对边缘比L2更友好Charbonnier是L1的平滑版本能减少梯度震荡。SSIM损失看结构相似度但计算慢显存占用高初学时可以先不加。训练循环建议每次迭代只喂一个batch记下loss和当前学习率。如果发现loss在某个值附近反复震荡优先考虑降低学习率而不是加loss权重。4.3 验证与生成去雾结果验证阶段不要只看loss一定要把去雾结果保存成图片肉眼对比。下面是保存结果的最小代码def save_dehaze_result(model, hazy, output_path): model.eval() with torch.no_grad(): pred model(hazy.unsqueeze(0)) pred pred.squeeze(0).permute(1, 2, 0).cpu().numpy() pred (pred * 255).astype(uint8) Image.fromarray(pred).save(output_path)注意模型输出是Sigmoid后的0-1张量所以转回0-255时要先乘255不要直接直接用ToPILImage。很多人在这里翻车保存出来的图是一张全黑的。验证集要固定不要每次重新shuffle否则不同epoch之间的指标没有可比性。5. 训练避坑与常见问题排查显存爆炸、NaN loss和去雾发灰5.1 显存爆炸patch size、batch和位置编码插值现象在2080Ti或更小的卡上batch_size8、输入512x512、patch_size8一跑训练就显存溢出。原因ViT的注意力矩阵大小是num_patches乘以num_patchespatch_size从16降到8序列长度变成原来的4倍注意力显存变成16倍。再叠加batch维度显存爆炸是必然的。解决先把patch_size调回16batch_size降到2输入尺寸降到256跑通一个step后再逐步往上加。如果必须用8的patch_size考虑把Transformer层的head数降到4或者用梯度累积模拟大batch。另一个常见做法是关掉混合精度里的AMP缓存虽然影响不大但有时能挤出几百兆。注意显存不够时优先降序列长度不要无脑减embed_dim因为去雾对特征容量有要求embed_dim太小会让图发糊。5.2 Loss突然变NaN学习率、归一化与混合精度现象训练前几步loss正常跑到第几百步突然变成NaN后续一直在NaN附近不再恢复。原因最常见的是初始学习率太大ViT对学习率比CNN敏感很多。另外如果输入图像没有归一化到0-1直接输入0-255像素值Transformer内部的LayerNorm会承受不住。如果用了AMP混合精度GELU激活函数在fp16下容易溢出。解决把学习率先降到1e-4甚至1e-5测试。输入图像用transforms.ToTensor()归一化不要用0-255的Tensor。AMP的GELU问题上可以在forward里把x转成float32再做激活或者干脆关掉AMP用纯fp32训练。5.3 输出发灰或偏色损失权重和输出层激活现象训练后去雾结果整体发灰远处天空偏白的地方变得脏兮兮彩色物体饱和度下降。原因直接回归清晰图时如果L2或L1损失权重过高模型为了降低平均误差会把输出逼近输入和GT的折中导致雾去除不彻底且颜色变淡。另一个原因是在输出层用了Tanh激活而GT是0-1输出范围是-1到1模型需要重新学习平移。解决输出层换Sigmoid损失里把感知损失或SSIM损失加上结构相似度会强迫模型保留边缘和色彩结构。还可以把天空区域在Loss中降低权重因为天空本身是低纹理区域模型会把整片天空的微小色差放大成斑块。5.4 去雾结果不干净数据增强与全局注意力现象同一张图部分区域去雾效果很好但某个角落残留明显白雾。原因如果patch embedding的stride等于patch_sizepatch之间没有重叠边界处容易被截断注意力很难跨越patch边界感知到连续的低频雾层。另一个原因可能是训练数据里没有覆盖类似的雾浓度分布。解决一是把patch切分的stride设为patch_size的一半让patch有重叠类似OverlappingPatchEmbed这样token数量会变多但信息连续性更好。二是数据增强时加入随机gamma变换和亮度扰动模拟不同雾浓度。三是检查是不是位置编码没有插值如果你训练时用256验证时用512位置编码对不上模型基本失效。6. 进阶验证用PSNR/SSIM和下游任务判断你的去雾模型是否真的能用6.1 指标计算脚本与评价口径PSNR和SSIM要统一口径一般用skimage计算即可from skimage.metrics import peak_signal_noise_ratio, structural_similarity psnr peak_signal_noise_ratio(clean, pred, data_range1.0) ssim structural_similarity(clean, pred, channel_axis2, data_range1.0)注意data_range一定要传否则np.uint8和float之间的结果会差十万八千里。6.2 把去雾结果接到目标检测上验证我在实际项目里发现PSNR高不代表下游任务效果好。后来养成了一个习惯把去雾后的图喂给一个现成的检测模型比如YOLOv8看检测AP有没有提升。如果去雾之后AP反而下降说明模型把纹理磨平了这对去雾应用而言是负优化。这个习惯比单纯看指标更能说明方案价值。6.3 一个习惯固定随机种子和保存中间结果训练前固定所有随机种子包括Python、NumPy、PyTorch和CUDA的否则复现结果全靠运气。每个epoch保存一次去雾输出图这样能直观看到训练过程中模型是先恢复整体亮度还是先恢复细节出了问题也能定位到具体epoch。以前我总以为把PSNR刷上去就万事大吉后来遇到一次真实雾图测试翻车才发现指标再好也可能只是过拟合了合成雾的分布。现在我做去雾项目一定会在真实雾图上做目视检查并且跑一个下游检测任务来验证。希望这个思路对你有帮助。本文还有配套的精品资源点击获取