简介:这是一份面向计算机相关专业学生与深度学习实践者的图像去雾项目源码包,以Vision Transformer为核心网络结构,可用于毕业设计、课程设计或期末大作业等场景,难度适中,适合具备一定Python与深度学习基础的学习者上手研究。压缩包共338个文件,约156.35MB,其中204个py文件构成模型训练与推理主体,另有yaml配置、csv实验记录、ipynb笔记、png与gif可视化结果、md说明文档及pptx汇报材料等,覆盖从数据准备、模型搭建到结果展示的完整流程。项目经本地编译调试,确保可运行,并附有文档说明,便于理解网络结构与实验设计。已有102人学习关注。读者可据此掌握ViT在去雾任务中的实现思路,参考损失曲线与实验数据组织方式,快速搭建自己的实验环境并完成论文或答辩所需的结果复现与展示。
1. 拆开这份 VisionTransformer 图像去雾项目:98 分毕设到底交付了什么
如果你正在做计算机方向的毕业设计,选题卡在“图像去雾”上,又不想用传统暗通道先验那套被做烂的方案,这份基于 VisionTransformer 的图像去雾算法实现值得先看一眼。它不是一篇纯论文,而是一个能跑起来的 Python 工程:模型定义、训练脚本、推理入口、文档说明都在里面,评审分 98 分,导师认可过。图像去雾本身是底层视觉里的经典病态问题——给定一张被雾霾退化的图,反推清晰场景,难点在于透射率估计和大气光耦合,传统方法靠先验,深度学习方法靠数据驱动,而 ViT 的加入是想用全局注意力替代卷积的局部感受野,让模型在浓雾、大面积天空区域也能抓住远距离依赖。这份资源适合三类人:赶毕设进度的本科生、需要课程设计/期末大作业素材的研究生,以及想拿一个完整视觉项目练手的学习者。难度适中,源码本地编译可运行,文档说明配套,拿到手能直接复现训练和推理流程。
2. 环境搭建与数据准备:从零把工程跑起来
2.1 依赖清单与版本对齐
拿到一个深度学习工程,第一件事不是急着python train.py,而是先把依赖锁死。图像去雾项目通常涉及 PyTorch、torchvision、timm(ViT 主干常从这里取)、numpy、opencv-python、Pillow、tqdm、matplotlib。版本不对齐是新手翻车的高发区,尤其是 PyTorch 和 CUDA 的匹配,以及 timm 版本和 ViT 权重加载的兼容性。
我一般会先建独立虚拟环境,避免污染全局:
# 创建并激活虚拟环境,Python 建议 3.8~3.10 python -m venv venv # Windows venv\Scripts\activate # Linux / macOS source venv/bin/activate # 安装核心依赖,torch 版本按你的 CUDA 版本去官网选对应命令 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install timm opencv-python pillow numpy tqdm matplotlib逻辑说明:虚拟环境隔离是为了防止不同项目间的包版本互相打架,这是血泪经验。torch的安装命令里cu118代表 CUDA 11.8,如果你没有 GPU 或只想 CPU 跑通,把整条换成pip install torch torchvision即可,但训练会慢到怀疑人生。timm是 ViT 类模型权重和结构的常用来源,版本差异会导致create_model的参数名变化,建议装较新的稳定版。
参数说明:--index-url指定 PyTorch 官方 wheel 源,比默认源快且版本全。如果你用的是 conda,把 pip 换成 conda 安装也行,但注意 conda 的 torch 和 pip 的 timm 混用偶尔会有 ABI 问题,能统一就统一。
2.2 去雾数据集的组织方式
图像去雾的训练数据一般是成对的:一张有雾图(hazy),一张对应的清晰图(clear/GT)。常见公开数据集有 RESIDE、Haze4K、Dense-Haze 等,工程里通常已经给出数据加载器,你只需要按它约定的目录结构放数据。典型结构长这样:
datasets/ ├── train/ │ ├── hazy/ │ │ ├── 001.png │ │ └── ... │ └── clear/ │ ├── 001.png │ └── ... └── test/ ├── hazy/ └── clear/关键点是 hazy 和 clear 下的文件名必须一一对应,否则配对加载会错位,训练出来的模型学的是“随机映射”,loss 降不下去你还找不到原因。数据加载器里一般用ImageFolder或自定义Dataset,自定义的更常见,因为要同时读两张图。
import os from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class DehazeDataset(Dataset): def __init__(self, root, mode='train', size=256): self.hazy_dir = os.path.join(root, mode, 'hazy') self.clear_dir = os.path.join(root, mode, 'clear') self.names = sorted(os.listdir(self.hazy_dir)) # 排序保证配对一致 self.transform = T.Compose([ T.Resize((size, size)), T.ToTensor(), # 归一化到 [0,1] ]) def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] hazy = Image.open(os.path.join(self.hazy_dir, name)).convert('RGB') clear = Image.open(os.path.join(self.clear_dir, name)).convert('RGB') return self.transform(hazy), self.transform(clear)逻辑说明:sorted是关键,os.listdir返回顺序不保证跨平台一致,不排序会导致 hazy 和 clear 错配。convert('RGB')防止灰度图或带 alpha 通道的图混进来导致通道数不匹配。Resize到固定尺寸是因为 ViT 对输入分辨率敏感,patch 划分要求尺寸能被 patch size 整除。
参数说明:size=256是常见选择,ViT 的 patch size 一般是 16,256/16=16,能整除。如果你显存够,可以上 384 或 512,但 ViT 的注意力计算量随分辨率平方增长,显存吃紧就老实 256。ToTensor已经把像素值缩到 [0,1],后面不要再手动除 255。
提示:如果你的数据集是 RESIDE 的 ITS/OTS 子集,注意它原始尺寸不统一,务必在 Dataset 里统一 resize,否则 batch 拼接时直接报错。
3. ViT 去雾网络结构:patch 嵌入、注意力与解码重建
3.1 为什么用 ViT 而不是纯 CNN
传统去雾网络(如 AOD-Net、GridDehazeNet)以 CNN 为主,卷积核的局部性决定了它擅长提取纹理和边缘,但对全局大气光分布、远距离雾浓度关联的建模能力有限。ViT 把图像切成 patch,每个 patch 线性嵌入成 token,再通过多头自注意力让任意两个 patch 直接交互,理论上能捕捉整幅图的雾分布规律。代价是 ViT 缺少 CNN 的归纳偏置,需要更多数据或预训练权重才能训好,这也是为什么工程里通常会加载 ImageNet 预训练的 ViT 主干。
选型理由落到实操:如果你的数据集只有几百对图,纯 ViT 从零训大概率过拟合或欠拟合;加载预训练权重 + 较小学习率微调是稳妥路线。工程里如果用的是timm的vit_tiny或vit_small,参数量可控,单卡也能跑。
3.2 patch 嵌入与位置编码的实现细节
ViT 的第一步是把 H×W×3 的图切成 N 个 P×P 的 patch,每个 patch 展平后过线性层映射到 embed_dim 维。位置编码是可学习参数,加上去让模型知道 patch 的空间顺序。
import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, img_size=256, patch_size=16, in_chans=3, embed_dim=768): super().__init__() self.num_patches = (img_size // patch_size) ** 2 # 256/16=16, 16*16=256 个 patch self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_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逻辑说明:用Conv2d的 kernel_size 和 stride 都等于 patch_size,等价于不重叠切块加线性映射,比手动 unfold 再 matmul 更高效。flatten(2)把 H/P 和 W/P 两维压成一维,transpose(1,2)把通道维换到最后,得到 transformer 需要的[B, N, C]形状。
参数说明:img_size必须能被patch_size整除,否则最后一行/列 patch 不完整。embed_dim决定 token 维度,vit_tiny 是 192,vit_small 是 384,vit_base 是 768。embed_dim 越大表达力越强但显存和计算量越大,毕设级别 vit_tiny 或 vit_small 足够。
3.3 编码器-解码器与去雾重建头
纯 ViT 输出的是 token 序列,要还原成图像需要解码。常见做法是编码器用 ViT 提全局特征,解码器用轻量卷积上采样逐步恢复分辨率,最后接一个卷积输出 3 通道的清晰图。也有直接预测透射率图 t 和大气光 A,再用大气散射模型 J = (I - A) / t + A 反推清晰图,这种物理引导的路线可解释性更强。
class DehazeViT(nn.Module): def __init__(self, img_size=256, patch_size=16, embed_dim=384, depth=6, num_heads=6): super().__init__() self.patch_embed = PatchEmbed(img_size, patch_size, 3, embed_dim) num_patches = self.patch_embed.num_patches self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, batch_first=True) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth) # 解码:token 序列 reshape 回特征图再上采样 self.decoder = nn.Sequential( nn.Conv2d(embed_dim, 128, 3, padding=1), nn.ReLU(inplace=True), nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), nn.Conv2d(128, 64, 3, padding=1), nn.ReLU(inplace=True), nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), nn.Conv2d(64, 3, 3, padding=1), nn.Sigmoid(), # 输出归一化到 [0,1] ) def forward(self, x): B = x.shape[0] x = self.patch_embed(x) + self.pos_embed x = self.encoder(x) # [B, N, C] -> [B, C, H/P, W/P] h = w = int(x.shape[1] ** 0.5) x = x.transpose(1, 2).reshape(B, -1, h, w) return self.decoder(x)逻辑说明:pos_embed用可学习参数,初始化成零或正态分布都行,训练中会自己学。TransformerEncoderLayer的batch_first=True让输入是[B, N, C],省去手动转置。解码阶段先把 token 序列 reshape 回[B, C, H/P, W/P],再两次 2 倍上采样,16×16 恢复到 64×64,如果你的输入是 256,这里只恢复到 64,需要再加两层上采样或改用插值到原尺寸。
参数说明:depth是 transformer 层数,6 层是轻量配置,vit_base 是 12 层。num_heads必须能整除embed_dim,384/6=64 每头维度。Sigmoid保证输出在 [0,1],和输入归一化范围一致,loss 才能正常收敛。
注意:如果你的解码器上采样次数和 patch 下采样倍数不匹配,输出尺寸会和原图对不上,训练时 loss 计算直接报 shape 错误。256 输入、patch 16,下采样 16 倍,解码要上采样 16 倍,即 4 次 2 倍上采样。
4. 训练、推理与指标评估:把模型跑出可用结果
4.1 损失函数选择与训练循环
去雾任务的损失通常组合 L1/L2 像素损失和感知损失(perceptual loss),有的还加 SSIM 损失。像素损失保证整体亮度颜色接近,感知损失用预训练 VGG 提特征,让重建图在语义上更自然。毕设级别用 L1 + 少量 SSIM 就够,感知损失会增加显存和训练时间。
import torch.nn.functional as F def train_one_epoch(model, loader, optimizer, device): model.train() total_loss = 0 for hazy, clear in loader: hazy, clear = hazy.to(device), clear.to(device) pred = model(hazy) # 若解码输出尺寸和 clear 不一致,插值对齐 if pred.shape[-2:] != clear.shape[-2:]: pred = F.interpolate(pred, size=clear.shape[-2:], mode='bilinear', align_corners=False) loss = F.l1_loss(pred, clear) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader)逻辑说明:F.interpolate是后悔药,防止解码尺寸和 GT 差几个像素导致报错。zero_grad必须在backward前,否则梯度会累加。loss.item()取出标量,避免把计算图带进累计变量导致显存泄漏。
参数说明:优化器一般用 AdamW,学习率 1e-4 到 2e-5,ViT 微调建议小学习率。batch size 受显存限制,256 分辨率下 vit_small 大概能开到 8~16。训练轮数看数据量,几百对图跑 100~200 epoch 常见。
4.2 推理脚本与结果保存
推理就是把测试集的有雾图喂进去,保存输出图,再和清晰图对比算指标。
@torch.no_grad() def inference(model, hazy_path, save_path, device, size=256): model.eval() img = Image.open(hazy_path).convert('RGB') inp = T.Compose([T.Resize((size, size)), T.ToTensor()])(img).unsqueeze(0).to(device) out = model(inp).squeeze(0).cpu() out_img = T.ToPILImage()(out.clamp(0, 1)) out_img.save(save_path)逻辑说明:torch.no_grad()关闭梯度,省显存提速。unsqueeze(0)加 batch 维。clamp(0,1)防止个别像素越界导致ToPILImage异常。
参数说明:推理尺寸可以和训练不同,但差异太大会掉点,建议一致。保存格式 PNG 无损,JPEG 会引入压缩伪影影响指标。
4.3 PSNR 与 SSIM 指标计算
去雾常用 PSNR 和 SSIM 衡量。PSNR 基于 MSE,对亮度敏感;SSIM 看结构相似性,更贴近人眼。
| 指标 | 含义 | 数值范围 | 越高越好 | 注意点 |
|---|---|---|---|---|
| PSNR | 峰值信噪比 | 一般 10~40 dB | 是 | 对像素误差敏感,模糊图也可能高分 |
| SSIM | 结构相似度 | 0~1 | 是 | 更贴近感知,但计算需对齐窗口 |
| MSE | 均方误差 | 0~∞ | 否 | PSNR 的中间量 |
import numpy as np from skimage.metrics import peak_signal_noise_ratio, structural_similarity def evaluate(pred, gt): pred = np.array(pred).astype(np.float32) / 255.0 gt = np.array(gt).astype(np.float32) / 255.0 psnr = peak_signal_noise_ratio(gt, pred, data_range=1.0) ssim = structural_similarity(gt, pred, channel_axis=2, data_range=1.0) return psnr, ssim逻辑说明:data_range=1.0因为已经归一化。channel_axis=2告诉 SSIM 这是彩色图,旧版 skimage 用multichannel=True,版本不同参数名不同,这是常见翻车点。
参数说明:PSNR 到 25 dB 以上算可用,30 dB 以上算不错。SSIM 0.85 以上视觉上就比较干净了。注意别拿训练集算指标,那是自欺欺人。
5. 避坑与排查:那些让工程跑不起来的常见问题
5.1 显存溢出(CUDA out of memory)
现象:训练刚开始或跑几个 batch 就报RuntimeError: CUDA out of memory。 原因:ViT 的注意力矩阵是 N×N,N 是 patch 数,256 分辨率下 N=256,注意力矩阵 256×256,batch 一大就爆;或者输入分辨率设太高。 解决:先把 batch size 降到 4 或 2,再把输入尺寸从 512 降到 256,还不行就换 vit_tiny。用torch.cuda.empty_cache()清理缓存,但根本还是降配置。
5.2 loss 不下降或变成 NaN
现象:训练几个 epoch loss 纹丝不动,或者突然变 NaN。 原因:学习率太大、数据没归一化、损失函数里出现 log(0) 或除零。 解决:学习率降到 1e-5 试,检查ToTensor后像素是否在 [0,1],如果自定义了物理模型反推,透射率 t 要 clamp 到 [0.1, 1.0] 防止除零。梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)也能救急。
5.3 输出图和输入尺寸对不上
现象:F.l1_loss报 shape mismatch,pred 是 [B,3,64,64],clear 是 [B,3,256,256]。 原因:解码器上采样倍数不够,patch 下采样 16 倍但只上采样了 4 倍。 解决:数清楚上采样次数,256 输入需要 4 次 2 倍上采样,或者最后加一个F.interpolate到目标尺寸兜底。
5.4 加载预训练权重报 key 不匹配
现象:load_state_dict报 missing keys 或 unexpected keys。 原因:timm 版本不同导致模型层命名变化,或者你改了网络结构但权重还是旧的。 解决:用strict=False加载,打印 missing/unexpected keys 看是哪些层,位置编码和分类头不匹配是正常的,主干匹配上就行。
5.5 评估指标异常高
现象:PSNR 35+,但肉眼看输出图还是灰蒙蒙。 原因:可能测试集和训练集有重叠,或者评估时用了训练集的清晰图。 解决:严格划分 train/test,评估只在 test 上做,且确认 hazy 和 clear 配对没错位。
6. 进阶技巧:让去雾效果再上一个台阶
跑通基础版本之后,想拿更高分或者让效果更耐看,有几个方向可以试。第一个是加多尺度特征融合,ViT 只取最后一层输出太浪费,把中间几层的 token 也接出来,reshape 成不同分辨率的特征图,用 FPN 那套自顶向下融合,浓雾区域和远景细节能同时兼顾。第二个是引入物理引导分支,让网络除了直接回归清晰图,再额外预测透射率 t,用大气散射模型约束,这样在合成数据集上指标可能略降,但真实雾图上泛化更好,答辩时也更有话讲。
第三个是数据增强,去雾数据集往往不够大,随机裁剪、水平翻转、颜色抖动都能用,但注意颜色抖动别太猛,否则大气光颜色分布被破坏,模型学歪。我一般会把增强后的图和原图一起看一遍,确认雾的浓度和色调还合理。
验证方法上,除了 PSNR/SSIM,建议做一组消融:去掉 ViT 换 CNN、去掉位置编码、去掉感知损失,各跑一遍,把指标列成表。答辩老师最爱问“你为什么这么设计”,有消融数据就是硬通货。
| 消融项 | PSNR | SSIM | 说明 |
|---|---|---|---|
| 完整模型 | 28.5 | 0.91 | 基准 |
| 去 ViT 换 CNN | 26.2 | 0.87 | 全局建模能力下降 |
| 去位置编码 | 27.1 | 0.89 | patch 顺序信息丢失 |
| 去感知损失 | 28.3 | 0.90 | 像素指标略升但视觉偏糊 |
最后说个具体技巧:推理时用滑动窗口或整图推理,取决于显存。整图推理如果尺寸和训练不一致,先 resize 到训练尺寸再推理,输出再 resize 回原尺寸,虽然会损失一点细节,但比直接喂大图导致 patch 数暴增爆显存强。我吃过这个亏,一张 4K 图直接喂进去,注意力矩阵大到显存瞬间清零,从那以后我每次推理前都强制检查输入尺寸和 patch 数的乘积,超过阈值就先 resize。希望帮到你。
本文还有配套的精品资源,点击获取