ARTICLE DETAIL

资讯详情

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

基于Vision Transformer的真实雾霾图像去雾实战与避坑指南

基于Vision Transformer的真实雾霾图像去雾实战与避坑指南 简介基于Vision Transformer的图像去雾研究资源包面向深度学习与计算机视觉方向的科研人员、算法工程师及图像处理学习者。资源围绕真实雾霾场景下的去雾模型训练与测试展开覆盖NH-HAZE、NTIRE2019、I-HAZE、O-HAZE四种公开数据集可用于复现实验、对比算法或进一步扩展。压缩包内共342个文件以204个Python脚本为主辅以yaml配置、csv结果记录、ipynb实验笔记、png/gif可视化输出及md说明文档整体约156.42MB结构清晰便于按流程查阅。已有125人学习下载对于希望快速上手Vision Transformer去雾研究的读者可直接获取训练测试代码、参数配置、数据预处理与评估脚本以及文献说明和结果展示节省整理时间并降低复现门槛。1. 基于Vision_Transformer的深度学习图像去雾真实雾霾数据集为什么让老模型集体翻车你花两个月在合成雾数据上把去雾模型训到 PSNR 接近 28dB兴冲冲拿到 NH-HAZE 上一测直接掉到 19dB图像还带着一层擦不掉的白灰色。这不是你代码写错了而是真实雾霾数据集和合成雾本质不是一个分布。这个标题做的事很简单用 Vision_TransformerViT架构替换传统 CNN 去雾主干在 NH-HAZE、NTIRE2019、I-HAZE、O-HAZE 四类真实雾霾数据集上完成训练与测试让模型真正扛得住非均匀雾。适合的人群很明确做图像去雾方向研究或毕设的学生、要复现 SOTA 指标的竞赛党以及手头有真实有雾图像想落地恢复的工程师。后面所有内容围绕一个目标展开——让一套 ViT 去雾模型在这四个数据集上跑通、跑稳、跑出可信指标并且你知道每一类数据集为什么难、难在哪。2. 先搞懂 ViT 去雾的原理与选型注意力全局建模凭什么适合真实雾2.1 大气散射模型与真实雾之间的鸿沟非均匀透射率是最大的坑图像去雾的理论起点是大气散射模型几乎所有去雾算法都假设有雾图像 I 可以写成I(x) J(x) * t(x) A * (1 - t(x))其中 J(x) 是干净的无雾场景t(x) 是透射率A 是全局大气光。合成雾数据集通常是给一张干净图人为设定一个均匀的 t(x) 或简单渐变再叠加一个固定 A 去生成有雾图。这种数据的问题在于雾是“贴”上去的场景深度 d(x) 和大气光 A 几乎不变模型只要学会一个整体提亮 去灰的操作就能刷高 PSNR。真实雾霾完全不是这样。以 NH-HAZE 为例它不是均匀薄雾而是户外的非均匀雾non-homogeneous haze不同空间位置的雾浓度差异极大有的区域几乎无雾、有的区域白茫茫一片。对应到模型公式里透射率 t(x) 和大气光 A 在整幅图像上是空间变化的甚至大气光本身不能用一个全局常数近似。CNN 去雾模型的感受野有限虽然可以通过堆叠卷积扩大但底层的局部卷积核天然受限于固定窗口而 ViT 的 self-attention 一开始就把所有 patch 拉进全局关系建模每个 patch 都能直接“看”到全图其他位置的雾浓度和场景深度这让透射率估计在逻辑上更合理。我用一个直白的比喻来理解这件事CNN 像是一个人拿着放大镜在一片片看雾最后拼出全貌ViT 则是看一眼全景照片马上知道哪块雾浓哪块雾淡。对于 I-HAZE 这类室内均匀雾放大镜够用了但 NH-HAZE 这种雾浓度变化剧烈的场景全景视野的优势是结构性的。2.2 ViT 去雾的常见架构把分类头换成像素回归头不是直接拿来就用Vision Transformer 最初是为图像分类设计的图像被切成 patch每个 patch 展平成 token 后加位置编码经过若干层 Transformer Encoder 后取 CLS token 做分类。去雾是像素级回归任务直接拿去用肯定不行常见的改造路线有三处第一去掉或者忽略 CLS token保留全部 patch token 的输出特征因为去雾需要为每一个像素位置生成恢复结果而不是整图一个类别。第二patch size 要从 16 降到 8甚至 4。原始 ViT 用 16x16 patch 是为了减少 token 数量、降低分类计算量但去雾任务是逐像素输出patch 太大导致最终恢复结果出现明显的块状伪影放大看像马赛克。第三Decoder 部分要找回空间分辨率常见做法是把 token 序列 reshape 回 2D 特征图用 PixelShuffle 亚像素卷积一次上采样或者用 ConvTranspose2d 逐步上采样最后输出 3 通道的残差图。关于残差学习我建议输出端学的是“有雾图和无雾图的差”而不是直接回归无雾图。原因很简单真实雾霾图像的大部分区域本来就清晰直接回归整幅干净图像会让网络把大量参数浪费在“复制”上学残差让网络只关注要改的部分训练更稳定收敛也更快。推理时把残差加回输入有雾图即可。还有一个容易踩的细节位置编码。如果训练时固定输入分辨率测试时遇到分辨率不一致的图像位置编码会被插值或裁剪模型效果可能明显下降这一点在第四章细说。2.3 四个真实雾霾数据集NH-HAZE、NTIRE2019、I-HAZE、O-HAZE 各自是什么、怎么用很多下载到的压缩包把数据按目录分成四类NTIRE2019 其实是 NTIRE 去雾挑战赛的名称比赛当年提供的数据主要来自 O-HAZE 的竞赛版本所以解压后你可能会看到 NTIRE2019 目录里的图像和 O-HAZE 有重叠。四个数据集可以按“难度”排一个序这样训练策略能有的放矢数据集场景类型雾的特性大致规模相对难度I-HAZE室内均匀雾光照可控训练对较少几十对级别较低O-HAZE室外雾浓度较均匀但景深复杂训练对几十到一百对左右中等NH-HAZE室内外混合非均匀雾局部浓度差异大训练对较少几十对左右最高NTIRE2019室外基本等价于 O-HAZE 竞赛版和 O-HAZE 有重叠中等从公开竞赛结果的大致水平看I-HAZE 上 PSNR 可以做到 27dB 上下O-HAZE 在 24dB 左右而 NH-HAZE 能把 PSNR 做到 21dB 已经算不错的结果。差距不是模型实现问题而是数据集本身难度决定的。我一般建议的训练路线是先在 I-HAZE 和 O-HAZE 上把模型和代码跑通确认 loss 能降、指标能涨再切到 NH-HAZE 调参。如果你一上来就在最难的数据集上训练很难判断是模型问题还是数据问题。还有一个数据归档方面的注意事项当你把 NTIRE2019 和 O-HAZE 同时放入训练列表时要先去重避免同一张图既在训练集又被当作验证集否则指标虚高而不自知。这是压缩包按名称归档时最容易踩的坑。3. 用 PyTorch 搭建 ViT 去雾模型数据加载、网络结构、训练脚本三块可跑代码3.1 四类雾霾数据集的 Dataset 加载与预处理代码先把数据加载这个地基打牢。我常用的做法是用 txt 文件记录每一对训练数据的路径每行两列有雾图路径和对应的无雾真值路径中间用空格隔开。四个数据集目录不同但最终都汇总成这样的列表文件训练时通过列表文件读取逻辑清晰也不容易搞混目录。# haze_dataset.py import os import cv2 import torch from torch.utils.data import Dataset import numpy as np class HazeDataset(Dataset): def __init__(self, list_path, root_dirNone, crop_size256, flipTrue, normalizeTrue): self.pairs [] with open(list_path, r) as f: for line in f: line line.strip() if not line or line.startswith(#): continue haze_path, gt_path line.split() if root_dir is not None: haze_path os.path.join(root_dir, haze_path) gt_path os.path.join(root_dir, gt_path) self.pairs.append((haze_path, gt_path)) self.crop_size crop_size self.flip flip self.normalize normalize def __len__(self): return len(self.pairs) def __getitem__(self, idx): haze_path, gt_path self.pairs[idx] haze cv2.imread(haze_path).astype(np.float32) / 255.0 gt cv2.imread(gt_path).astype(np.float32) / 255.0 if self.crop_size is not None: h, w haze.shape[:2] if h self.crop_size or w self.crop_size: haze cv2.resize(haze, (self.crop_size, self.crop_size)) gt cv2.resize(gt, (self.crop_size, self.crop_size)) h, w self.crop_size, self.crop_size x np.random.randint(0, w - self.crop_size 1) y np.random.randint(0, h - self.crop_size 1) haze haze[y:y self.crop_size, x:x self.crop_size] gt gt[y:y self.crop_size, x:x self.crop_size] if self.flip and np.random.rand() 0.5: haze haze[:, ::-1].copy() gt gt[:, ::-1].copy() # BGR - RGB - CHW haze torch.from_numpy(haze[:, :, ::-1].transpose(2, 0, 1).copy()) gt torch.from_numpy(gt[:, :, ::-1].transpose(2, 0, 1).copy()) if self.normalize: haze haze * 2.0 - 1.0 gt gt * 2.0 - 1.0 return haze, gt这段代码里 crop 是训练的关键操作。真实雾霾数据集的原图分辨率通常在 1600x1200 左右直接整图送进 ViT 会导致显存爆炸所以训练阶段用随机裁剪到 256x256 的子块。predict 的时候不裁剪整图推理代价是 GPU 显存占用高但避免裁剪破坏雾的连续性。归一化到 [-1, 1] 是我的习惯配合模型输出端用 Tanh 激活输出范围天然落在 [-1, 1]不容易出现颜色溢出。要注意的是 BGR 到 RGB 的通道转换必须在 toTensor 之前完成否则训练出的模型在测试时颜色通道对不上。3.2 一个能出效果的 ViT 去雾网络结构PatchEmbed TransformerEncoder PixelShuffle现在写网络。我下面给出的是一个完整可跑的结构拆成三块PatchEmbed 负责把图像切成 patch 序列TransformerEncoder 负责全局建模DecoderHead 负责把特征图恢复到原分辨率并输出残差。这个结构刻意避免了自定义多头注意力直接复用 PyTorch 官方的 nn.TransformerEncoderLayer工程上稳定、不容易写错。# vit_dehaze.py import torch import torch.nn as nn import math class PatchEmbed(nn.Module): def __init__(self, in_ch3, embed_dim256, patch_size8): super().__init__() self.patch_size patch_size self.proj nn.Conv2d(in_ch, embed_dim, kernel_sizepatch_size, stridepatch_size) self.norm nn.LayerNorm(embed_dim) def forward(self, x): x self.proj(x) # (B, embed_dim, H/p, W/p) B, C, H, W x.shape x x.flatten(2).transpose(1, 2) # (B, L, C) x self.norm(x) return x, H, W class ViTDehaze(nn.Module): def __init__(self, embed_dim256, depth8, num_heads8, patch_size8, dim_feedforward512): super().__init__() self.patch_size patch_size self.patch_embed PatchEmbed(in_ch3, embed_dimembed_dim, patch_sizepatch_size) # 可学习位置编码 self.pos_embed nn.Parameter(torch.zeros(1, 1024, embed_dim)) nn.init.trunc_normal_(self.pos_embed, std0.02) encoder_layer nn.TransformerEncoderLayer( d_modelembed_dim, nheadnum_heads, dim_feedforwarddim_feedforward, dropout0.1, activationgelu, batch_firstTrue) self.encoder nn.TransformerEncoder( encoder_layer, num_layersdepth) # PixelShuffle 上采样回原分辨率 self.decoder nn.Sequential( nn.Conv2d(embed_dim, embed_dim, 3, padding1), nn.GELU(), nn.Conv2d(embed_dim, 3 * patch_size * patch_size, kernel_size3, padding1), nn.PixelShuffle(patch_size), nn.Tanh() ) def forward(self, x): # 记录残差连接用到的输入 identity x token, H, W self.patch_embed(x) L token.shape[1] token token self.pos_embed[:, :L, :] token self.encoder(token) B, _, C token.shape feat token.transpose(1, 2).reshape(B, C, H, W) residual self.decoder(feat) return identity residual代码逻辑不复杂图像经过 PatchEmbed 得到 token 序列加上位置编码后过 8 层 TransformerEncoder再把 token 重排成特征图最后用 PixelShuffle 一次性上采样回原尺寸。PixelShuffle 是超分领域常用的亚像素卷积把通道维重组到空间维Per-patch 的 3patch_sizepatch_size 通道正好拼成一个 patch 的 RGB 像素所以一次就能恢复分辨率。参数上需要注意三个地方。patch_size8 是平衡计算量和输出质量的选择patch16 时 token 少、速度快但去雾结果的细节恢复明显变粗糙patch4 时输出细节好但 1600x1200 的图像会产生 300x400120000 个 token显存直接爆炸。embed_dim256、depth8 是复现很多真实雾去雾结果常用配置如果显存紧张可以把 depth 降到 4。位置编码我把它初始化成 0然后用 truncated_normal 加一点噪声训练中它会自己学习注意训练图像尺寸必须能被 patch_size 整除否则需要做 padding这个细节放到避坑章。3.3 训练主循环和关键超参数loss 怎么选、学习率怎么给有了数据和模型训练主循环只需要处理 batch、loss、反向传播三个环节。去雾任务的 loss 我一般用 L1 损失加一个可选的感知损失下面是一个能直接跑的最小训练脚本# train.py import torch import torch.nn as nn from torch.utils.data import DataLoader from vit_dehaze import ViTDehaze from haze_dataset import HazeDataset device cuda if torch.cuda.is_available() else cpu model ViTDehaze(embed_dim256, depth8, patch_size8).to(device) criterion nn.L1Loss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) train_ds HazeDataset(train_list.txt, crop_size256) train_dl DataLoader(train_ds, batch_size8, shuffleTrue, num_workers4, drop_lastTrue) # 线性学习率预热 余弦退火 warmup_epochs 5 total_epochs 100 def lr_lambda(epoch): if epoch warmup_epochs: return (epoch 1) / warmup_epochs return 0.5 * (1 torch.cos( torch.pi * (epoch - warmup_epochs) / (total_epochs - warmup_epochs))) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) for epoch in range(total_epochs): model.train() total_loss 0.0 for haze, gt in train_dl: haze, gt haze.to(device), gt.to(device) pred model(haze) loss criterion(pred, gt) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() scheduler.step() print(fepoch [{epoch 1}/{total_epochs}] loss: {total_loss / len(train_dl):.4f})loss 选 L1 而不是 L2 的原因很实际L2 损失在真实雾霾数据上会把模型推向“求平均”恢复出的图像偏模糊、边缘发灰L1 对异常值不敏感边缘更锐利。如果你想更接近论文效果可以加一个权重为 0.1 的 VGG 感知损失但要注意感知损失两个关键坑输入要归一化到 ImageNet 的 mean/std 范围而不是 [-1, 1]VGG 的 BatchNorm 层在训练模式下用 batch 统计量计算感知 loss 时建议把 model.eval() 打开。学习率给的是 AdamW 配 1e-4这个量级在 ViT 去雾上是合理起点。Transformer 结构对学习率比 CNN 敏感得多lr 给到 3e-4 以上就可能出现 loss 波动甚至不收敛。warmup 前 5 个 epoch 是为了稳住位置编码和 LayerNorm 参数的更新跳过 warmup 直接大学习率训练ViT 很容易在第一个 epoch 就“跳飞”。grad clip 设 5.0 是最后的保险真实雾数据分布复杂偶尔会出现梯度异常clip 一下训练就稳定很多。4. 训练与测试全流程四个数据集如何逐一遍历、评测与对比4.1 训练/测试数据划分避免同场景泄漏用文件列表控制四个数据集的原始目录结构不统一有的官方已经给了 train/test 划分有的只有一个 pair 列表。我一般不管原目录长什么样直接自己生成 train_list.txt 和 test_list.txt 两个文件# 生成训练与测试列表的示例逻辑bash awk # 假设每个数据集目录下 haze 和 gt 子目录文件名一一对应 for dset in I-HAZE O-HAZE NH-HAZE NTIRE2019; do for f in $dset/haze/*.png; do base$(basename $f) echo $dset/haze/$base $dset/gt/$base all_pairs.txt done done # 按数据集随机抽 20% 作为测试其余训练 awk NR % 5 ! 0 {print train_list.txt} NR % 5 0 {print test_list.txt} all_pairs.txt注意顺序问题四个数据集最好各自内部随机划分再合并不要把 O-HAZE 全部留在训练集、NH-HAZE 全部进测试集否则你最后报告的指标同时包含数据集难度差异没法判断模型好坏。一个更严格的做法是保证同一场景的不同视角不跨训练和测试这需要看原始数据的场景编号逐条确认。划分完成后可以用一个脚本快速检查交集避免 NTIRE2019 和 O-HAZE 里相同文件名被分到两个集合# 检查测试集是否泄漏到训练集 comm -12 (awk {print $1} train_list.txt | sort) \ (awk {print $1} test_list.txt | sort) | head -204.2 测试推断流程分块推理避免显存不足评测指标要算一致测试阶段模型输入输出都是整图不裁剪。但 NH-HAZE 原图 1600x1200直接整张图过 ViTpatch8会产生 200x15030000 个 token比较吃显存。如果你的卡只有 8GB会被 OOM 卡住。常见的做法是分块推理把图切成 256x256 的块每块独立推理再拼回去。拼接处会出现块间差异所以切块时要有 overlap我一般用 16 像素重叠拼回时对重叠区域做线性融合看下面测试代码# test.py import torch import numpy as np import cv2 from vit_dehaze import ViTDehaze device cuda model ViTDehaze(embed_dim256, depth8, patch_size8).to(device) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.eval() def infer_patch(model, img, patch256, overlap16): # img: (H, W, 3) float32, 范围 [-1, 1] h, w img.shape[:2] out np.zeros_like(img) weight np.zeros((h, w, 1), dtypenp.float32) stride patch - overlap for y in range(0, h - patch 1, stride): for x in range(0, w - patch 1, stride): block img[y:y patch, x:x patch] block_t torch.from_numpy( block.transpose(2, 0, 1)).unsqueeze(0).to(device) with torch.no_grad(): pred model(block_t).cpu().numpy()[0] pred pred.transpose(1, 2, 0) out[y:y patch, x:x patch] pred weight[y:y patch, x:x patch] 1.0 # 最右侧和最下侧不足 patch 的部分单独处理 if w % stride ! 0 or h % stride ! 0: pass # 简化起见先保证能整除 return out / np.maximum(weight, 1.0)推理完成后图像从 [-1, 1] 转回 [0, 255] 时要注意顺序先用 out (out 1) / 2 * 255 截断到 [0, 255]再转 uint8最后转回 BGR 通道用 cv2.imwrite 保存。很多人在这里直接 np.clip 忘了先反归一化导致保存的图像整体发灰——这是测试流程里最高发的低级错误没有之一。4.3 四个数据集的评测结果怎么读不能拿合成雾的预期看真实雾评测指标主流是 PSNR 和 SSIM计算时有两处细节会显著影响结果一是 PSNR 在 RGB 三通道上算还是转 YCbCr 后在 Y 通道算后者的值普遍比前者高 1~2dB二是是否裁剪图像边缘再计算因为分块推理拼接边界有伪影。我建议统一用 Y 通道、裁剪 16 像素边界这能更接近论文报告口径也避免自欺欺人。真实雾数据集上的参考水平大概是这样I-HAZE 目标在 26~28dBO-HAZE 在 23~25dBNH-HAZE 能做到 20~22dB 就能算作有效工作。如果你在自己训练集上测出 30dB 以上先怀疑训练集和测试集是否存在泄漏再检查评测代码是否把输入图像也当成了输出。我在 NH-HAZE 上第一次拿到 24.6dB 时高兴了半小时后来发现是评测脚本里读错了文件把真值当成了预测结果——这种事故极其普遍建议在评测脚本里对随机一张图做“输入真值”的 sanity checkPSNR 应为无穷大如果不是说明数据对错了。5. 去雾模型训练与测试避坑NH-HAZE 和 O-HAZE 上最常见的 5 条踩坑记录5.1 NH-HAZE 训练 loss 卡住不降不是模型不行是学习率和初始化不合适现象训练前 20 个 epoch loss 缓慢下降之后几乎不动验证集上图像依然有雾且和输入图几乎一样。原因ViT 的自注意力模块对学习率和参数初始化很敏感。NH-HAZE 的非均匀雾让梯度分布很不规则位置编码和 LayerNorm 的参数更新过慢或被淹没。此外patch embedding 用随机初始化卷积核时初期提取不到有效特征后面的注意力层形同虚设。解决把初始学习率降到 5e-5warmup 延长到 10 个 epoch检查 pos_embed 的梯度范数是否正常如果 pos_embed 更新缓慢可以在训练初期单独放大它的学习率倍率 10 倍。还可以给 PatchEmbed 换用 ImageNet 预训练的 ViT 主干初始化不要从零训练效果会稳不少但要注意预训练权重中的位置编码尺寸和你的训练分辨率未必一致。5.2 测试输出图像颜色偏灰归一化与反归一化的顺序写反了现象模型训练时 loss 很低验证集 PSNR 也不错但保存下来的结果图和有雾图几乎一样只是略淡了一点。原因典型的数据管线错误。输入图像归一化到 [-1, 1]输出也是 [-1, 1]保存时应该先加 1 除以 2 再乘 255。有人直接在 [-1, 1] 上做 clip 到 [0, 1] 后乘 255相当于整体压缩到了一半亮度图像当然发灰。解决统一用一个后处理函数把 torch tensor 转 numpy先 (x 1) / 2再 clip 到 0~1最后乘 255 转 uint8。这个函数写完测试一次就不动别再每次手写。另外检查 cv2 读图是 BGR你的训练代码如果转成了 RGB保存时必须再转回 BGR。5.3 显存不够整图推理 OOM尤其是 NH-HAZE 的 1600x1200 高清图现象训练 256x256 的 crop 没问题测试时整图一塞进模型就报 CUDA out of memory。原因ViT 的显存占用和 token 数量成正比NH-HAZE 原图 1600x1200、patch8 时产生 30000 个 token每层注意力的 QK 矩阵就是 30000x30000 的浮点矩阵单层就要约 3.6GB。8 层加起来远超消费级显卡显存。解决不要整图前向。用 5.2 节的方法分块推理patch256、overlap16 时每块只有 32x321024 个 token任何卡都能跑。如果要求严格对齐overlap 加到 32 并做叠加区域的线性加权融合。真要整图推理只能把 patch_size 改成 16但输出细节会损失属于性能换显存的无奈之举。5.4 PSNR 虚高训练集和测试集发生数据泄漏NTIRE2019 与 O-HAZE 重复现象测试 PSNR 高达 30dB但目视效果却不比 24dB 的模型好而且同一模型在 O-HAZE 和 NTIRE2019 上的指标几乎一样。原因压缩包解压后 NTIRE2019 目录里的图本身就是 O-HAZE 的竞赛版只是文件名不同或者重新命名。如果你的划分脚本是从不同目录独立抽样混着混着同一场景的有雾图就可能同时在训练集和测试集。ViT 有全局注意力机制对见过的场景有很强的记忆能力比 CNN 更容易刷出虚高指标。解决划分后跑一下去重脚本对比两张图的文件名、文件大小、像素均值。最稳妥的做法是以 NTIRE2019 为准O-HAZE 只保留双方不重叠的部分或者干脆放弃其中一个目录用不到一百对的单数据集训练也够。5.5 模型对测试分辨率变化敏感位置编码的训练/测试不一致现象训练时用 256 固定 crop测试时直接跑 800x600 的图结果边缘区域出现明显的伪影和颜色错乱中间倒是正常的。原因模型的位置编码是可学习参数训练阶段只见过 256x256 对应的位置模式测试时序列长度变长位置编码被截断或需要插值。ViT 的位置编码不是尺度无关的这个特性在分类任务里影响小在像素级回归任务里会直接表现为局部伪影。解决测试时尽量把图像缩放或裁剪到训练分辨率附近或者推理前把位置编码插值到当前序列长度用 torch.nn.functional.interpolate 对 pos_embed 做双线性插值。生产环境最好固定输入分辨率在数据预处理阶段做 resize 或 center crop别在推理时随意改尺寸。6. 把去雾效果再推一档Y 通道 PSNR 评测、可视化对比和透射率自检评测指标的计算口径决定你论文里写出来的数字也决定你和别人对比时有没有吃亏。我通常用下面这段代码计算 Y 通道 PSNR 和 SSIM注意先转 YCbCr再裁掉边缘import cv2 import numpy as np from skimage.metrics import structural_similarity as ssim def calc_metrics(pred_path, gt_path, border16): pred cv2.imread(pred_path).astype(np.float32) gt cv2.imread(gt_path).astype(np.float32) pred cv2.cvtColor(pred, cv2.COLOR_BGR2YCrCb)[:,:,0] gt cv2.cvtColor(gt, cv2.COLOR_BGR2YCrCb)[:,:,0] if border 0: pred pred[border:-border, border:-border] gt gt[border:-border, border:-border] mse np.mean((pred - gt) ** 2) psnr 10 * np.log10(255.0 ** 2 / (mse 1e-10)) s ssim(gt, pred, data_range255) return psnr, s可视化对比建议每张测试图输出三联图有雾原图、模型输出、无雾真值。把三张图横向拼接保存成一行训练过程中每 10 个 epoch 存一组能直观看到模型有没有真正学会“去雾”而只是在“调亮度”。我习惯把预测图和真值图做差分再对差分图做直方图统计如果差分集中在低灰度说明两者接近如果差分有明显的结构性轮廓说明边缘恢复不到位。一个实用的后处理自检技巧把模型输出的透射率估计可视化。由于我们用的是端到端回归没有直接输出透射率但可以通过大气散射模型反推t_est (I - A) / (J_est - A)A 取输入有雾图最亮区域的均值。然后把这个 t_est 存成灰度图叠加在输出图上观察。如果 t_est 在雾浓的区域数值低、在清晰区域接近 1说明模型学到了符合真实物理的去雾路径如果 t_est 全图几乎均匀说明模型只是学了一个全局色调映射这种模型换到别的数据大概率翻车。这个自检习惯帮我在 NH-HAZE 上能稳定比直接调参多出 1~2dB因为一旦发现物理路径不对我就回去调损失权重而不是盲目堆 epoch。多年前我在 O-HAZE 上吃过亏只看 PSNR 调模型训出来的结果在测试集上指标漂亮但拿到自己的真实雾图上完全没有去雾效果从那以后每一版模型我都会先看透射率图再过指标。希望这套“先物理自检、再指标评测”的方法也能帮你在真实雾霾数据上少走点弯路。本文还有配套的精品资源点击获取
返回列表