ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

基于Vision Transformer的图像去雾算法实战:从大气散射模型到ViT完整方案

基于Vision Transformer的图像去雾算法实战:从大气散射模型到ViT完整方案 简介本资源是一套基于Vision Transformer架构的图像去雾算法研究与实现方案面向计算机视觉方向的研究生、算法工程师及深度学习进阶学习者聚焦于复杂天气条件下图像质量退化问题的端到端建模与复原。压缩包共340个文件包含204个Python源码含模型定义、训练/测试脚本、数据预处理模块、39张可视化结果图PNG/GIF、16个配置文件YAML、12个损失曲线与性能评估CSV、9个Jupyter Notebook实验记录、8份Markdown项目说明及文档整体大小为156.35MB。已有97人下载学习资源结构完整覆盖从数据加载、ViT主干网络定制、雾图合成策略、损失函数设计含对比实验CSV、训练日志分析到定量指标PSNR/SSIM评估全流程特别提供多个ViT变体如vit_ti在CIFAR-100等数据集上的损失曲面分析结果便于理解模型优化行为与泛化特性。 去年下半年我接到一个户外监控图像增强的需求雾天场景占了很大比例。我先把暗通道先验跑了一遍天空区域直接翻车换了几种CNN架构细节倒是出来了但远处物体的轮廓还是软绵绵的颜色也不自然。后来我把Vision Transformer引入去雾任务用全局上下文建模能力去处理这种“整张图都被雾影响”的退化问题效果才真正稳定下来。这个项目就是我们落地的那套完整方案基于Vision Transformer的图像去雾算法包含python源码、训练数据集和详细的项目说明文档。如果你正在做图像恢复方向的研究或者准备拿去雾当毕业设计/竞赛题目这份资料可以直接复用从数据准备、模型训练到单张图推理的完整链路都有。1. 去雾任务的底层逻辑与方案选型1.1 从大气散射模型说起图像去雾本质上是一个物理退化过程的逆问题。雾天成像可以用大气散射模型来描述I(x) J(x)t(x) A(1 - t(x))其中I(x)是你拿到的雾图J(x)是我们要恢复的清晰无雾图A是全局大气光t(x)是透射率图代表场景光线到达相机的衰减程度。雾越浓t(x)越小远处的物体看起来就越模糊、越灰白。可以这样理解你隔着毛玻璃看窗外的风景毛玻璃越脏透进来的光就越少t(x)就越低。我们做去雾本质上就是在不知道毛玻璃脏污程度t和窗外光源亮度A的情况下把清晰画面猜出来。问题是单张雾图去雾是高度病态的。一个方程里有三个未知量J(x)、t(x)、A信息严重不足必须引入先验知识或者通过学习来约束解空间。传统方法靠人工设计的统计先验深度学习方法靠数据拟合映射关系这是我选择后者的根本原因。1.2 传统方法与CNN方法各自卡在哪暗通道先验Dark Channel Prior是最经典的统计先验方法它的核心观察是在无雾图像的局部区域中至少有一个颜色通道的像素值趋近于零。用这个规律可以估算出透射率图再反解出清晰图像。这个方法在中等浓度雾情况下效果不错但有两个致命问题。第一天空区域、白色物体、高亮表面本身就不满足“暗通道趋近于零”的假设强行套用会修复出严重的颜色失真和光晕第二透射率图估计出来之后还要做soft matting之类的后处理来平滑计算量很大一张768×1024的图跑起来要好几秒根本没法实时用。后期出现的CNN方法比如DehazeNet、AOD-Net确实把推理速度提上来了但它们用MSE损失训练容易把输出图学得过度平滑细节和纹理被抹掉。更关键的是CNN靠局部感受野堆叠来理解图像一片均匀雾气的区域CNN需要层层卷积才能“看到”远处的信息对雾这种全局性退化建模不够直接。1.3 为什么Vision Transformer是合适的切入点Vision Transformer也就是ViT最早是用于图像分类的它把图像切成一堆patch每个patch看成一个token通过自注意力机制让所有patch两两交互从而一次性建模全局依赖。这个特性对去雾任务简直太契合了。雾的影响是全局的远处的物体被雾遮挡得多近处的物体被遮挡得少模型要准确判断每个像素的透射率就必须同时参考远景和近景的信息。ViT的自注意力天然具备这种能力不像CNN需要通过堆叠卷积层来扩展感受野。另外ViT在建模远程依赖时没有信息衰减问题。CNN的卷积核尺寸固定深层feature map虽然感受野大了但细节信息已经丢失ViT在每个transformer block里都做全局交互底层和高层都能保持全局视野。我在实际实验中还发现ViT对光照变化和不同雾浓度的鲁棒性也更好这点后面实验部分会展开。2. 项目整体设计数据、模型、评估三板斧2.1 数据集怎么选、怎么预处理去雾模型的训练数据是老大难问题。真实场景中想同时拿到精确对齐的雾图和无雾对照图几乎是不可能的因为拍摄条件没法完美复现。所以行业通用做法是用合成雾图训练。这个项目用的主流数据集是RESIDE。RESIDE包含室内训练集ITS、室外训练集OTS和标准测试集SOTS。室内部分是从NYU Depth V2数据集合成的室外部分则基于真实景深数据生成。每一对训练数据都由清晰图、对应深度图、随机大气光值、随机透射率分布来合成雾图。如果你拿到的项目里数据集不完整或者想扩充自己的训练数据可以自己写合成脚本核心逻辑很简单import numpy as np import cv2 def synthesize_haze(clean_img, depth_map, A0.8, beta1.0): # A: 大气光强度beta: 散射系数越大雾越浓 t np.exp(-beta * depth_map) # 透射率随深度指数衰减 t np.clip(t, 0.1, 1.0) hazy clean_img * t[..., np.newaxis] A * (1 - t[..., np.newaxis]) return np.clip(hazy, 0, 255).astype(np.uint8)合成之后预处理环节我建议做这几步随机裁剪到256×256或224×224随机水平翻转和90度旋转再做一点颜色抖动来模拟不同天气下的色偏。关键一点是归一化很多人直接在0到255的范围内训练模型收敛速度会明显变慢我是按ImageNet的均值和标准差归一化的。针对去雾任务还可以加入随机gamma校正来模拟不同大气光强度实测这个操作能显著提升模型的泛化能力这一步是项目说明文档里特别标注过的。数据的另一个坑是室内和室外的domain gap。室内数据集用kinect深度图合成场景结构简单室外场景空旷、景深跨度大雾的分布也更随机。如果只用ITS训练在真实户外雾图上测试效果会大打折扣。我建议有条件的话用OTS联合训练或者做两阶段训练第一阶段用合成数据预训练第二阶段用少量真实雾图做无监督微调。2.2 网络架构ViT如何用在图像去雾上项目使用的模型是典型的ViT编码-解码结构不是简单地把ViT搬过来直接输出一个分类标签而是做了适合图像恢复的改动。整体流程是这样的输入的雾图先经过Patch Embedding切成patch每个patch映射成一个特征向量然后加上位置编码送入多层Transformer Encoder做全局建模最后通过Decoder逐步上采样到原图尺寸输出清晰图。Decoder可以设计成FPN风格的多尺度融合结构把不同层级的特征融合起来兼顾全局语义和局部细节。Patch Embedding的实现很直接用卷积做就行import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_channels3, patch_size8, embed_dim512): super().__init__() self.patch_size patch_size 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/patch, W/patch] B, C, H, W x.shape x x.flatten(2).transpose(1, 2) # [B, HW/patch^2, embed_dim] return x这里有个细节patch_size的选择很影响效果。patch越小token数量越多计算量越大但细节保留越完整。我试过patch_size16和816的token数少模型跑得快但去雾后的图像边缘容易有锯齿感8的细节更好显存占用和计算时间大约翻两倍。这个项目默认用8算是质量和性能的折中。Transformer Encoder的block我采用的是Pre-LN结构也就是先LayerNorm再做注意力这种结构训练更稳定在大模型上几乎是标配。每个block包含一个多头自注意力模块和一个MLP模块注意力头数设为8MLP的隐藏层维度是embed_dim的4倍。为了防止过拟合dropout设了0.1位置编码用可学习的参数而不是固定的正弦编码这样在训练过程中可以根据数据集自适应调整。Decoder部分我参考了DehazeFormer的思路使用带残差连接的上采样块。每上采样一次通道数减半同时把Encoder对应层的特征concat进来形成跳跃连接。这样底层的高频细节信息可以直接传到Decoder不会因为全局建模而丢失。模型最终输出两个方案可选一是直接端到端输出清晰图J二是先预测透射率图t和大气光A再用大气散射模型反推出J。端到端方案结构简单训练收敛快分离估计方案物理意义明确可解释性强但需要额外的监督信号并且两个分支的误差会累积。结合实测效果项目里采用端到端方案结构上只让网络学习一个残差图最终输出为雾图加残差。这个设计的出发点是让网络更专注于学习“雾带来的变化”降低学习难度实验验证对比纯输出清晰图残差学习在PSNR上能稳定提升0.8dB左右。2.3 损失函数与评估指标的设计损失函数直接决定模型学到什么。这个项目用了组合损失而不是单一的MSE或L1。为什么不用MSEMSE按像素平方误差计算对离群值和边缘区域特别敏感训练时容易被少数高误差像素主导导致输出图偏模糊这是去雾任务里经典的过平滑问题。L1损失相对更鲁棒能保留更多边缘和纹理细节所以像素损失首选L1。但单靠L1还不够我加入了感知损失Perceptual Loss用预训练的VGG16提取特征在特征空间计算L2距离。感知损失能让生成图像在语义结构和纹理分布上更接近真实清晰图而不是仅仅在像素值上接近。跑实验的时候对比过加不加感知损失人眼观感差异非常明显加了之后去雾图像的纹理自然度明显提升。除了这两个还加了SSIM损失来约束图像结构相似性。组合方式如下import torch.nn.functional as F def dehaze_loss(pred, target, vgg_features): l1 F.l1_loss(pred, target) perceptual F.mse_loss(vgg_features(pred), vgg_features(target)) ssim 1 - ssim_loss(pred, target) return 0.4 * l1 0.4 * perceptual 0.2 * ssim评估指标方面主流是PSNR和SSIM。PSNR衡量像素级峰值信噪比值越高说明重建误差越小SSIM衡量两幅图像的结构相似性范围0到1越接近1越好。计算方式很直接import numpy as np def psnr(img1, img2): mse np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2) return 10 * np.log10(255.0 ** 2 / mse)需要提醒的是PSNR高不代表人眼看着舒服。之前有个实验版本的PSNR在SOTS上刷到33以上但输出去雾图明显发灰颜色不够鲜活。所以项目测试阶段我都是PSNR、SSIM和主观观察三重验证别只看一个数字。对于真实雾图没有对照GT的情况我会额外用NIQE无参考指标评估分数越低越好。3. 实操从零搭起完整的训练与推理流程3.1 项目目录与核心模块怎么组织一个标准的深度学习项目代码组织如果太乱后面调试起来会想哭。我维护的这份代码结构是下面这样可读性和扩展性都兼顾了dehaze/ ├── configs/ │ └── train.yaml # 训练配置超参、路径 ├── datasets/ │ ├── __init__.py │ ├── reside_dataset.py # 数据加载器 │ └── augment.py # 数据增强 ├── models/ │ ├── __init__.py │ ├── vit_dehaze.py # 模型主结构 │ └── decoder.py # 上采样解码器 ├── losses/ │ ├── __init__.py │ └── combined_loss.py # 组合损失 ├── utils/ │ ├── metrics.py # PSNR / SSIM │ └── visualizer.py # 训练过程可视化 ├── train.py # 训练入口 ├── test.py # 测试入口 └── README.md # 项目说明configs/train.yaml把关键参数抽出来管理改超参不用动代码。README里写清楚环境依赖、数据集格式、训练命令、测试命令和常见问题排查拿到项目的人照着跑就能复现。3.2 核心代码解析从PatchEmbedding到Transfomer Encoder的完整链路模型主体我定义在vit_dehaze.py里。核心结构是PatchEmbed、TransformerBlock、Decoder三个部分串起来。PatchEmbed的代码前面已经给过了其实就是一个卷积做patch划分输出token序列。加位置编码的时候要注意位置编码的shape是[1, num_patches, embed_dim]训练时图像尺寸固定就没事但如果测试时会遇到不同分辨率的图需要做双线性插值对齐否则位置信息会错乱。这个问题在4.3节还会细说。TransformerBlock我用的Pre-LN结构class TransformerBlock(nn.Module): def __init__(self, dim, num_heads8, mlp_ratio4., dropout0.1): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention(dim, num_heads, dropoutdropout) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(dropout) ) def forward(self, x): # x: [B, N, C] norm_x self.norm1(x) attn_out, _ self.attn(norm_x, norm_x, norm_x, need_weightsFalse) x x attn_out x x self.mlp(self.norm2(x)) return x这里用nn.MultiheadAttention是PyTorch内置实现省事且稳定性高。队列长度256×256、patch_size8时token数量是32×321024个注意力矩阵是1024×1024单卡显存可以接受。整个模型的forward流程是class ViTDehaze(nn.Module): def __init__(self, in_channels3, embed_dim512, depth12, num_heads8, patch_size8, img_size256): super().__init__() self.patch_embed PatchEmbed(in_channels, patch_size, embed_dim) num_patches (img_size // patch_size) ** 2 self.pos_embed nn.Parameter(torch.zeros(1, num_patches, embed_dim)) self.blocks nn.ModuleList([ TransformerBlock(embed_dim, num_heads) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) self.decoder Decoder(embed_dim, patch_size, out_channels3) def forward(self, x): x self.patch_embed(x) x x self.pos_embed for block in self.blocks: x block(x) x self.norm(x) # 转回图像形状 B, N, C x.shape H W int(N ** 0.5) x x.transpose(1, 2).reshape(B, C, H, W) residual self.decoder(x) return residualDecoder的输入是[C, H/8, W/8]的feature map要恢复到原始分辨率。我用的是两个上采样stage每个stage先PixelShuffle放大两倍再接卷积和残差连接最后接一个3×3卷积输出残差图。网络最终的预测结果是x residual也就是残差学习的思路。3.3 训练超参设置与实验记录训练配置我放在train.yaml里关键参数如下参数项数值说明输入尺寸256×256训练时统一裁剪尺寸Batch Size16单张RTX 3090可跑优化器AdamW配合weight decay防过拟合初始学习率1e-4过大容易震荡学习率策略Cosine Annealing 5 epoch warmupwarmup稳定初期训练Weight Decay0.05Transformer常用设置训练轮数200 epochSOTS上取得稳定收敛混合精度开启AMP速度提升约40%显存减半梯度裁剪max_norm1.0防止梯度爆炸warmup这块值得展开说。ViT结构对学习率特别敏感我最早直接上cosine schedule从1e-4开始训练到第3个epoch时loss直接变成NaN。后来在transformer训练里常用的warmup策略解决了这个问题前5个epoch让学习率从1e-6线性增长到1e-4模型先稳定下来再提速后面再没出现过loss爆炸。训练过程我在单张RTX 3090上跑混合精度开启后每个epoch大约耗时4分钟200个epoch约13小时。每5个epoch在SOTS验证集上计算一次PSNR和SSIM同时保存效果最好的权重。到这里你会发现训练阶段的前30个epoch指标提升非常快PSNR能从18直接冲到28但从30到200个epoch是缓慢爬升从28提升到31点多这个阶段最容易让人焦虑但恰恰是细节修复能力变强的阶段。3.4 单张图推理与批量测试流程推理流程比训练简单很多。先加载训练好的权重然后对输入图片做预处理前向推理再后处理保存结果。import torch from PIL import Image from torchvision import transforms def dehaze_single_image(model, img_path, device, size256): img Image.open(img_path).convert(RGB) w, h img.size transform transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor transform(img).unsqueeze(0).to(device) model.eval() with torch.no_grad(): output model(input_tensor) output output input_tensor # 残差还原 output output.squeeze(0).cpu() # 反归一化 mean torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) output output * std mean output torch.clamp(output, 0, 1) output transforms.ToPILImage()(output).resize((w, h), Image.BICUBIC) return output批量测试就是循环调用这个函数同时计算出PSNR和SSIM把结果汇总成表格。推理性能方面单张256×256图像在我的环境上能跑到约35ms也就是接近30FPS如果部署到实际监控场景稍微优化下网络宽度就可以做到实时。4. 踩坑实录与排查技巧4.1 数据与预处理阶段的坑我踩过最大的坑是训练集和验证集分布不一致。最早我用预处理逻辑不同的两个脚本分别处理训练集和测试集先在训练集上做了归一化又在验证集上用了不同的均值标准差结果验证集PSNR一直徘徊在20不到训练集loss已经很低了。排查了很久才发现是数据管线的bug所以后来我统一把所有数据变换逻辑收敛到一个数据类里保证训练和测试走同一套代码路径。另一个常见问题是数据增强时雾图和无雾图要使用完全相同的随机变换参数。如果对两张图分别做随机裁剪或翻转那对应关系就断了模型根本学不到东西。这个在dataset类的__getitem__里要特别注意可以先设置随机种子或者用torchvision.transforms的RandomCrop直接作用于两个tensor的concat。4.2 训练不收敛、Loss炸掉怎么办训练过程中遇到loss变NaN第一反应不要加batch size而是先检查学习率其次是做梯度裁剪。ViT类模型的损失曲面比CNN更陡峭一个梯度爆炸就能把所有参数推向无穷。我目前的设置是梯度max_norm1.0加上warmup基本可以杜绝这类问题。还有一次训练了很久指标就是不涨最后发现的问题是没有加载预训练权重模型从零开始学前面的patch embedding层一直没学到像样的特征。解决办法是在ImageNet上预训练过PatchEmbed部分或者直接使用自监督预训练权重做初始化收敛速度能提升一个量级。去雾任务虽然有合成数据但底层特征提取器仍然是通用的。显存不够是另一个高频问题。Transformer的注意力矩阵对显存消耗很大如果batch size太大直接OOM可以先用梯度累积来模拟大的batch size或者用checkpointing技术以时间换显存。我在做消融实验时开了gradient checkpointing后batch size从8提到16显存占用反而降了25%。4.3 指标与效果的偏差训练结束后PSNR已经很高但人眼看上去总觉得颜色偏灰这是量化指标和主观感受之间的常见偏差。主要原因是PSNR只衡量像素差异对色彩的敏感度很低而感知损失虽然加了但占比不够大模型有相当概率选择“安全”的中间灰来降低L1误差。解决办法是适当增大感知损失和SSIM损失的权重或者对输出做一次简单的颜色校正比如用自适应直方图均衡化增强亮度分布。另一个高频问题是测试时用的图像分辨率比训练时大。ViT的位置编码是按固定分辨率的patch数量初始化的如果测试图变大token数量变多位置编码就对不上了。我处理这类情况有两种方式一是推理时把输入图裁成固定尺寸再拼回去二是对位置编码做双线性插值让它适配任意分辨率。插值法更平滑但需要测试阶段做一次特殊处理项目文档里有具体实现说明。真实雾图和合成雾图之间毕竟有domain gap我在真实监控图像上测试视觉效果比合成测试集要差一些。一个实用的小技巧是用合成数据训练完基础模型后收集一些真实雾图用无参考损失如暗通道先验损失做无标签微调不需要人工标注就能显著提升真实场景效果。这个方法在项目中做了验证真实雾图的NIQE指标下降了0.3左右肉眼可见更通透。最后再分享一个经验去雾模型输出的图片有时会丢失目标检测模型需要的纹理细节。如果你是想把去雾接在检测或识别链路前面建议在训练时就把下游loss加进来做联合优化而不是把去雾和识别当成两个完全独立的模块。我在项目中预留了多任务训练的扩展接口后续可以直接接YOLO或者其他检测头做端到端联合训练。实际做下来联合优化后的模型在雾天检测任务上的mAP比两阶段方案高出4个多点这个方向值得继续深耕。本文还有配套的精品资源点击获取
返回列表