基于VisionTransformer的图像去雾算法Python源码实战:从原理到部署

发布时间:2026/10/10 12:20:55
基于VisionTransformer的图像去雾算法Python源码实战:从原理到部署 简介本资源为基于VisionTransformer的图像去雾算法研究与实现项目面向计算机相关专业做毕业设计、课程设计或期末大作业的学生以及需要图像去雾实战练习的学习者。项目经导师指导并认可评审分98分源码均经本地编译调试可稳定运行难度适中。压缩包共338个文件约156.35MB以204个Python源码为核心辅以yaml配置、csv实验数据、ipynb笔记、png与gif效果图、md说明文档及pptx等覆盖模型定义、训练配置、结果记录与可视化全流程。已有102人学习下载。读者可获得完整的VisionTransformer去雾实现方案、可复现的训练与评估脚本、实验数据与损失景观记录以及配套文档说明便于快速理解算法结构、复现实验并在此基础上完成二次开发或论文撰写。1. 从一张雾天照片说起这套 VisionTransformer 去雾源码到底能跑出什么去年帮一个做安防的朋友处理监控截图画面里雾蒙蒙一片车牌和人脸全糊成一团。他问我有没有办法用深度学习把雾去掉我第一反应是找现成的开源方案结果翻了一圈发现大部分去雾项目要么是纯 CNN 的老架构要么代码跑不起来、依赖缺东少西。后来在整理毕业设计资源时翻到这套基于 VisionTransformer 的图像去雾算法 Python 源码结构完整、文档齐全本地编译过能直接跑才算是把这个需求落地了。这套资源的核心就是用 ViTVision Transformer做图像去雾和传统 CNN 去雾最大的区别在于CNN 靠卷积核的局部感受野逐层堆叠来捕捉雾的分布特征而 ViT 通过自注意力机制直接建模全局像素之间的关系对大面积均匀雾的去除效果更自然。资源里包含完整的训练脚本、推理脚本、预训练权重加载逻辑和一份文档说明适合正在做计算机相关毕业设计的学生、需要项目实战练习的学习者也能直接作为课程设计或期末大作业的底稿。难度适中不需要你从零推导 Transformer 的数学公式但得能看懂 PyTorch 的基本训练流程。2. VisionTransformer 去雾的原理拆解与网络结构选型2.1 为什么去雾任务适合用 ViT 而不是纯 CNN图像去雾的本质是估计一张透射率图transmission map和大气光值atmospheric light然后根据大气散射模型反推出无雾图像。传统方法比如暗通道先验依赖手工设计的先验知识遇到天空区域或者白色物体就容易翻车。CNN 方法虽然能端到端学习但卷积核的感受野有限要覆盖整张图的雾分布需要堆很深的层训练成本高且容易过拟合。ViT 的思路不一样。它把图像切成固定大小的 patch每个 patch 展平后加位置编码然后送进 Transformer Encoder。自注意力机制让每个 patch 都能直接和图中任意其他 patch 计算关联权重这意味着模型在去雾时能同时考虑近处的纹理细节和远处的雾浓度分布。常见做法是采用混合架构浅层用几个卷积层做初步特征提取和降采样后面接 Transformer Block 做全局建模最后用卷积层做上采样恢复分辨率。这套源码走的就是这个路线编码器部分用 ViT-Tiny 级别的配置解码器用轻量卷积堆叠整体参数量控制在可接受范围内。选型上还有一个现实考量纯 ViT 对训练数据量要求极高ImageNet 级别才喂得饱。去雾数据集通常只有几千到几万对图像所以源码里大概率用了预训练权重加载或者混合 CNN-Transformer 结构来降低对数据量的依赖。你在看代码时会发现model.py里定义了PatchEmbed、Attention、Block这些模块拼起来就是一个标准的 ViT 骨架然后在头部接了一个去雾专用的重建头。2.2 源码目录结构与关键文件说明拿到压缩包解压后先别急着跑train.py。花五分钟把目录结构过一遍能省掉后面很多瞎找的时间。典型的目录布局是这样的dehaze-vit/ ├── data/ │ ├── train/ │ │ ├── hazy/ # 有雾图像 │ │ └── clear/ # 对应无雾图像GT │ └── test/ │ ├── hazy/ │ └── clear/ ├── models/ │ ├── vit_dehaze.py # ViT 去雾网络定义 │ └── blocks.py # Attention、PatchEmbed 等基础模块 ├── utils/ │ ├── dataset.py # 数据加载与预处理 │ ├── metrics.py # PSNR、SSIM 计算 │ └── visualize.py # 训练过程可视化 ├── configs/ │ └── default.yaml # 超参数配置文件 ├── train.py # 训练入口 ├── inference.py # 单张/批量推理脚本 ├── requirements.txt # 依赖清单 └── README.md # 文档说明models/vit_dehaze.py是核心里面定义了完整的ViTDehaze类包含__init__里的模块初始化和forward里的前向传播逻辑。utils/dataset.py负责把 hazy 和 clear 图像配对加载通常会做随机裁剪、翻转增强和归一化。configs/default.yaml里集中管理学习率、batch size、patch size、embed dim 这些关键参数改配置比改代码安全得多。注意不同版本的源码目录名可能有差异比如models可能叫networkutils可能叫tools。以 README 里的说明为准别硬套上面的结构。2.3 环境配置与依赖安装的实操步骤环境这块是新手最容易卡住的地方。源码大概率是在 PyTorch 1.10 到 2.0 之间的版本上调试通过的Python 版本建议 3.8 或 3.9。先建虚拟环境别在 base 环境里瞎装# 创建并激活虚拟环境 python -m venv venv_dehaze source venv_dehaze/bin/activate # Windows 用 venv_dehaze\Scripts\activate # 安装 PyTorch根据你的 CUDA 版本选对应命令 # CUDA 11.8 的情况 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 安装其他依赖 pip install -r requirements.txtrequirements.txt里通常包含numpy、opencv-python、pillow、tqdm、pyyaml、matplotlib这些。如果安装opencv-python时报错换成opencv-python-headless试试服务器环境没有图形界面时用 headless 版本更稳。装完之后验证一下import torch print(torch.__version__) print(torch.cuda.is_available()) # 有 GPU 应该返回 True如果cuda.is_available()返回 False先检查显卡驱动和 CUDA 版本是否匹配。没有 GPU 也能跑但训练速度会慢很多建议至少用一张 8GB 显存的卡batch size 调到 4 或 8 先跑通再说。3. 训练自己的去雾模型数据准备、参数配置与训练脚本3.1 去雾数据集的获取与配对整理去雾任务需要成对的训练数据同一场景的有雾图和对应的无雾图GT。常用的公开数据集有 RESIDE、Haze4K、Dense-Haze 等源码文档里一般会注明推荐用哪个。如果你手头没有现成数据集可以用合成的方式生成拿一张清晰图像根据大气散射模型人为加雾公式是I(x) J(x) * t(x) A * (1 - t(x))其中I是有雾图J是无雾图t是透射率A是大气光值。用 Python 写个脚本就能批量生成配对数据import cv2 import numpy as np import os def add_haze(img, beta1.5, A0.85): 给清晰图像加雾 beta: 散射系数越大雾越浓 A: 大气光值通常 0.7~1.0 h, w img.shape[:2] # 生成深度图这里用简单的线性深度模拟 depth np.linspace(0.1, 1.0, h).reshape(-1, 1) depth np.repeat(depth, w, axis1) # 透射率 t np.exp(-beta * depth) t np.clip(t, 0.1, 1.0) # 加雾 haze img * t[..., np.newaxis] A * (1 - t[..., np.newaxis]) return np.clip(haze * 255, 0, 255).astype(np.uint8) # 批量处理 clear_dir data/clear hazy_dir data/hazy os.makedirs(hazy_dir, exist_okTrue) for fname in os.listdir(clear_dir): img cv2.imread(os.path.join(clear_dir, fname)) if img is None: continue hazy add_haze(img, beta1.2, A0.9) cv2.imwrite(os.path.join(hazy_dir, fname), hazy)beta控制雾的浓度A控制大气光的亮度。实际训练时建议生成多档beta值的数据让模型见过不同浓度的雾泛化能力会好很多。数据整理好之后确保hazy和clear两个文件夹里的文件名一一对应dataset.py里通常是按文件名匹配的对不上会直接报错。3.2 配置文件里的关键参数怎么调configs/default.yaml是训练的控制中心几个核心参数直接决定模型能不能收敛、显存够不够用# 数据相关 data: train_dir: data/train val_dir: data/test patch_size: 16 # ViT 的 patch 大小16 是标准配置 img_size: 256 # 训练时裁剪的图像尺寸 batch_size: 8 # 显存不够就降到 4 或 2 # 模型相关 model: embed_dim: 384 # ViT-Tiny 的嵌入维度 depth: 6 # Transformer Block 层数 num_heads: 6 # 多头注意力的头数 mlp_ratio: 4.0 # MLP 隐藏层倍率 # 训练相关 train: epochs: 200 lr: 1e-4 # 初始学习率 weight_decay: 1e-5 scheduler: cosine # 余弦退火 warmup_epochs: 10 # 预热轮数patch_size设成 16 意味着 256x256 的图会被切成 16x16 个 patch序列长度是 256。如果显存吃紧把img_size降到 128序列长度变成 64显存占用会大幅下降。embed_dim和depth决定模型容量ViT-Tiny 级别384/6在去雾任务上通常够用想刷更高 PSNR 可以往上加到 768/12但训练时间和显存需求也会翻倍。学习率用1e-4配合 cosine 调度是 Transformer 类模型的常见做法。如果训练 loss 震荡厉害先检查 batch size 是不是太小再考虑把 lr 降到5e-5。warmup 阶段很重要前 10 个 epoch 让学习率从 0 线性升到初始值能避免训练初期梯度爆炸。3.3 启动训练与监控 loss 曲线配置改好之后直接跑训练脚本python train.py --config configs/default.yaml --gpu 0训练过程中终端会打印每个 epoch 的 loss 和验证集 PSNR。同时utils/visualize.py通常会把 loss 曲线和 PSNR 曲线画出来存到logs/目录下。重点关注两件事训练 loss 是否稳定下降验证 PSNR 是否在上升。如果训练 loss 降但验证 PSNR 不涨甚至下降说明过拟合了需要加数据增强或者加 dropout。# 典型的训练循环核心逻辑简化版 for epoch in range(epochs): model.train() for hazy, clear in train_loader: hazy, clear hazy.cuda(), clear.cuda() pred model(hazy) loss criterion(pred, clear) optimizer.zero_grad() loss.backward() optimizer.step() # 验证 model.eval() with torch.no_grad(): psnr evaluate(model, val_loader) print(fEpoch {epoch}: loss{loss.item():.4f}, PSNR{psnr:.2f})criterion一般是 L1 loss 或者 L1 SSIM 的混合损失。纯 L1 收敛快但图像可能偏模糊加 SSIM 能提升结构相似度但训练初期不稳定。源码里如果默认用 L1你可以自己改成0.8 * L1 0.2 * (1 - SSIM)试试效果。提示训练前先跑一个 epoch 看看显存占用用nvidia-smi盯着。如果 OOM 了优先降 batch size其次降 img_size最后才动模型结构。4. 推理与效果验证单张图像去雾和指标计算4.1 用训练好的权重跑单张推理训练完成后权重会保存在checkpoints/目录下通常是best.pth或last.pth。推理脚本的用法python inference.py --weights checkpoints/best.pth --input test.jpg --output result.jpg推理脚本的核心逻辑是加载模型、读取图像、做同样的预处理归一化、resize 到训练时的尺寸、前向传播、后处理还原到 0-255 范围并保存。注意预处理必须和训练时一致否则效果会差很多。常见做法是把推理封装成一个函数方便调用import torch import cv2 import numpy as np from models.vit_dehaze import ViTDehaze def dehaze_image(model, img_path, devicecuda): 对单张图像去雾 model.eval() # 读取并预处理 img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_tensor torch.from_numpy(img).float() / 255.0 img_tensor img_tensor.permute(2, 0, 1).unsqueeze(0).to(device) # 推理 with torch.no_grad(): output model(img_tensor) # 后处理 output output.squeeze(0).permute(1, 2, 0).cpu().numpy() output np.clip(output * 255, 0, 255).astype(np.uint8) output cv2.cvtColor(output, cv2.COLOR_RGB2BGR) return output # 使用 model ViTDehaze().to(cuda) model.load_state_dict(torch.load(checkpoints/best.pth)) result dehaze_image(model, test.jpg) cv2.imwrite(result.jpg, result)如果推理结果偏暗或者偏亮检查一下训练时用的归一化参数是不是mean[0.5,0.5,0.5], std[0.5,0.5,0.5]推理时要对应做同样的归一化和反归一化。4.2 PSNR 和 SSIM 指标的计算与解读去雾效果不能光靠肉眼看得有量化指标。PSNR峰值信噪比和 SSIM结构相似度是两个最常用的。源码的utils/metrics.py里一般已经实现了import numpy as np from skimage.metrics import peak_signal_noise_ratio, structural_similarity def compute_psnr_ssim(pred, gt): 计算 PSNR 和 SSIM pred, gt: numpy array, shape (H, W, 3), dtype uint8 psnr peak_signal_noise_ratio(gt, pred, data_range255) ssim structural_similarity(gt, pred, multichannelTrue, data_range255) return psnr, ssimPSNR 超过 25dB 说明去雾效果已经不错了超过 30dB 属于比较好的水平。SSIM 越接近 1 越好0.9 以上说明结构保留得很完整。但要注意PSNR 高不代表视觉效果好有时候 PSNR 略低但看起来更自然的模型反而更实用。建议在测试集上跑完指标后再随机抽几张图肉眼对比一下。4.3 用损失地形 CSV 文件辅助分析训练稳定性资源正文里提到了几个 CSV 文件比如cifar100_resnet_dnn_50_losslandscape.csv、cifar100_vit_ti_losslandscape.csv这些是损失地形loss landscape的可视化数据。损失地形能直观展示模型在参数空间里的收敛情况如果地形平坦且最低点清晰说明训练稳定如果地形崎岖、有多个局部极小值说明模型可能陷入次优解。虽然这些 CSV 是针对 CIFAR 分类任务生成的但分析方法可以迁移到去雾任务上。你可以用类似的方式记录训练过程中不同 epoch 的 loss 值然后用matplotlib画出来import pandas as pd import matplotlib.pyplot as plt # 读取 loss 记录 df pd.read_csv(logs/train_loss.csv) plt.plot(df[epoch], df[loss], labeltrain) plt.plot(df[epoch], df[val_loss], labelval) plt.xlabel(Epoch) plt.ylabel(Loss) plt.legend() plt.savefig(loss_curve.png)如果训练 loss 和验证 loss 的 gap 越来越大说明过拟合如果两条线都震荡说明学习率可能偏大或者 batch size 太小。这些 CSV 文件的价值在于提供了一个标准化的分析模板你可以照着这个格式记录自己模型的训练数据。5. 避坑与排查跑这套去雾源码时最容易翻车的五个地方5.1 显存不够导致训练直接崩掉现象运行train.py后报RuntimeError: CUDA out of memory训练还没开始就挂了。原因ViT 的自注意力计算复杂度是序列长度的平方256x256 的图切成 16x16 的 patch 后序列长度是 256注意力矩阵是 256x256加上多头和 batch 维度显存占用比同参数量的 CNN 大不少。batch size 设成 8 在 8GB 卡上很容易爆。解决先把 batch size 降到 2 或 4同时把img_size从 256 降到 128。如果还不行用梯度累积模拟大 batch在训练循环里每算 4 个 batch 才optimizer.step()一次效果接近 batch size 翻倍。5.2 数据路径写错导致 DataLoader 返回空现象训练脚本跑起来但每个 epoch 瞬间结束loss 是 nan 或者根本不打印。原因configs/default.yaml里的train_dir路径写错了或者hazy和clear文件夹里的文件名没有一一对应dataset.py在配对时全部跳过返回了空列表。解决在dataset.py的__init__里加一行print(fLoaded {len(self.pairs)} pairs)跑之前先确认配对数量对不对。文件名匹配时注意扩展名大小写.jpg和.JPG在 Linux 下是区分大小写的。5.3 预训练权重加载时报 key 不匹配现象model.load_state_dict(torch.load(vit_pretrain.pth))报Missing key(s)或Unexpected key(s)。原因预训练权重是在 ImageNet 分类任务上训的分类头是 1000 类而去雾模型的输出是 3 通道图像分类头的 key 对不上。另外 ViT 的pos_embed如果预训练时的图像尺寸和当前不一致也会不匹配。解决用strictFalse加载跳过不匹配的层model.load_state_dict(state_dict, strictFalse)。然后在训练时让这些新加的层用更大的学习率单独训练。如果pos_embed尺寸对不上可以用插值的方式调整import torch.nn.functional as F def interpolate_pos_embed(model, state_dict): if pos_embed in state_dict: pos_embed state_dict[pos_embed] if pos_embed.shape ! model.pos_embed.shape: pos_embed F.interpolate( pos_embed.permute(0, 2, 1), sizemodel.pos_embed.shape[1], modelinear, align_cornersFalse ).permute(0, 2, 1) state_dict[pos_embed] pos_embed return state_dict5.4 训练 loss 不降反升或者变成 nan现象前几个 epoch loss 正常下降突然跳到 nan 或者持续上升。原因学习率太大导致梯度爆炸或者数据里有损坏的图像全黑、全白、尺寸异常导致 loss 计算出问题。混合精度训练时如果没做梯度裁剪也容易出 nan。解决先把学习率降到1e-5试试然后在训练循环里加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)同时检查数据集里有没有异常图像用脚本扫一遍from PIL import Image import os for f in os.listdir(data/train/hazy): try: img Image.open(os.path.join(data/train/hazy, f)) img.verify() except Exception as e: print(f损坏文件: {f}, 错误: {e})5.5 推理结果全黑或者全白现象推理脚本跑完保存的图像要么全黑要么全白完全看不出内容。原因预处理和后处理不匹配。训练时如果做了img / 255.0归一化推理时输出是 0-1 范围直接保存成图像就是全黑因为像素值接近 0。或者归一化用了 ImageNet 的 mean/std推理时没做对应的反归一化。解决确认训练时的预处理流程推理时严格对应。如果训练用了transforms.Normalize(mean[0.485,0.456,0.406], std[0.229,0.224,0.225])推理输出后要反归一化mean np.array([0.485, 0.456, 0.406]) std np.array([0.229, 0.224, 0.225]) output output * std mean # 反归一化 output np.clip(output * 255, 0, 255).astype(np.uint8)6. 进阶技巧用混合损失和注意力可视化把去雾效果再提一档训练跑通之后如果想把 PSNR 从 25dB 推到 28dB 以上光靠调学习率是不够的。我一般会从损失函数和模型可解释性两个方向下手。先说损失函数。纯 L1 loss 训出来的结果在数值指标上还行但放大看细节会发现纹理偏平滑边缘不够锐利。换成 L1 SSIM 感知损失的组合视觉效果会明显改善。感知损失用预训练 VGG 提取特征计算去雾结果和 GT 在特征空间的 L2 距离import torch import torch.nn as nn import torchvision.models as models class PerceptualLoss(nn.Module): def __init__(self): super().__init__() vgg models.vgg16(pretrainedTrue).features[:16].eval() for p in vgg.parameters(): p.requires_grad False self.vgg vgg def forward(self, pred, gt): # 归一化到 VGG 期望的输入范围 mean torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1).to(pred.device) std torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1).to(pred.device) pred_norm (pred - mean) / std gt_norm (gt - mean) / std feat_pred self.vgg(pred_norm) feat_gt self.vgg(gt_norm) return nn.functional.l1_loss(feat_pred, feat_gt) # 组合损失 criterion lambda pred, gt: 0.6 * nn.functional.l1_loss(pred, gt) \ 0.3 * (1 - ssim(pred, gt)) \ 0.1 * perceptual_loss(pred, gt)权重分配上L1 占大头保证像素级准确SSIM 提升结构相似度感知损失负责纹理细节。三个权重的比例可以根据你的数据集特点微调我一般从 0.6/0.3/0.1 开始试。另一个技巧是注意力可视化。ViT 的自注意力权重能告诉你模型在去雾时关注了哪些区域。把最后一层 Transformer Block 的注意力图拿出来叠加到原图上能直观看到模型是不是在雾浓的区域分配了更多注意力def visualize_attention(model, img_tensor): 提取并可视化最后一层的注意力权重 attentions [] def hook_fn(module, input, output): # 保存注意力权重 attentions.append(output[1]) # 假设 output 是 (x, attn) # 注册 hook 到最后一个 Block 的 Attention 模块 handle model.blocks[-1].attn.register_forward_hook(hook_fn) with torch.no_grad(): model(img_tensor) handle.remove() # attentions[0] shape: (B, num_heads, N, N) attn attentions[0][0].mean(dim0) # 对多头取平均 # 取 CLS token 对其他 patch 的注意力 cls_attn attn[0, 1:].reshape(16, 16) # 假设 patch 数是 256 return cls_attn.cpu().numpy()如果注意力图显示模型在天空区域和远处景物上分配了高权重说明它学到了雾的分布规律如果注意力分散且无规律可能需要更多的训练数据或者更长的训练轮数。还有一个容易被忽略的点测试时增强TTA。对同一张测试图做水平翻转、垂直翻转分别推理后再把结果翻转回来取平均PSNR 通常能涨 0.2 到 0.5dB代价只是推理时间翻几倍。对于毕业设计答辩来说这点提升足够让指标好看不少。从那以后我每次跑去雾实验都会先把损失地形 CSV 的格式套用到自己的训练日志上每 10 个 epoch 存一次 loss 和 PSNR最后画在一张图里对比不同配置的收敛曲线。这个习惯帮我省了很多次盲目调参的时间也让我在写论文时能直接拿出有说服力的训练过程证据。希望帮到你。本文还有配套的精品资源点击获取

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询