★ FEATURED ARTICLE 简介面向视觉研究与工程应用的图像去雪项目实战代码基于上下文交互与尺度感知Transformer架构可用于恢复被雪覆盖或模糊的图像提升监控、遥感、户外摄影等场景下恶劣天气画面的可用性。算法借鉴Transformer自注意力机制同时引入上下文交互建模像素关联并利用尺度感知应对不同厚度的雪层从而兼顾全局结构与局部细节。资源包共11个文件以8个Python源码为核心覆盖SnowFormer模型、condconv条件卷积、CL1/感知损失、PSNR/SSIM评估、数据加载及测试模块另有2张示例效果图和1份README文档整体仅2.17MB便于快速下载与部署。目前已有73人学习下载适合作为Transformer图像复原任务的入门参考与二次开发基础。项目代码结构清晰、流程闭环核心模块均可直接复用通过对比示例图能直观验证算法效果方便读者在此基础上针对特定去雪场景进行改进与实验。1. 图像除雪为什么比去雨更难先分清雪花、雪斑和雪雾三类退化第一次跑图像除雪项目时我下意识把它当成去雨的近亲用一套现成的去雨网络改改损失函数就上了。结果在真实雪景图上效果惨不忍睹——稀疏的大雪花是去掉了但远处雪雾造成的灰蒙蒙一片完全没动。后来才想明白一个关键区别雨是线状、方向性强、基本等密度分布雪呢近处是遮挡型的大雪花中景是密集的小雪点远处是降低对比度的雪雾同一张图里三种退化叠加感受野需求完全不一样。这也是为什么最近的除雪工作都在往多尺度、上下文建模上靠。这个项目把两件事结合到一起上下文交互模块负责关联全局信息尺度感知transformer负责在不同尺度上剔除雪痕。对刚入门图像复原的人来说它是很好的进阶练手项目对要做雪天视觉落地的工程师这套源码可以直接拿来做二次训练省掉自己搭骨架的时间。2. 网络结构拆解上下文交互模块与尺度感知transformer的三个配合阶段2.1 尺度感知分支怎么设计多patch多膨胀率的特征金字塔先看这张图里最常见的设计——尺度感知不是一个单独的模块而是三个平行分支在开头就分叉。每个分支负责一个尺度层次的退化特征第一个分支用较小的patch比如4×4和大膨胀率空洞卷积专门捕捉稀疏大雪花留下的空洞感第二个分支中等patch8×8、中等膨胀率处理密集雪斑第三个分支用更大patch16×16配合全局平均池化提取雪雾造成的低频亮度偏移。class ScaleAwareBranch(nn.Module): def __init__(self, in_ch, branch_id): super().__init__() # branch_id 0: 大雪花遮挡, 1: 密集雪斑, 2: 雪雾低频 patch_sizes [4, 8, 16] strides [4, 8, 16] self.patch_embed PatchEmbed( img_size256, patch_sizepatch_sizes[branch_id], stridestrides[branch_id], # 有重叠 patch保留边缘细节 in_chin_ch, embed_dim192 ) # 三层不同膨胀率的卷积做局部上下文补充 self.dilated_convs nn.Sequential( nn.Conv2d(192, 192, 3, padding2, dilation2, groups192), nn.GELU(), nn.Conv2d(192, 192, 3, padding4, dilation4, groups192), nn.GELU(), )这段代码的逻辑是用patch embedding把不同尺度的雪退化映射到同一个特征维度192维再用膨胀卷积组扩大每个token的感受野。注意groups192是depthwise卷积——参数少、速度快而且每个通道独立处理不会在早期就把不同频率的信息混在一起。三个分支的输出是三个尺度的token序列后续要送到各自独立的transformer层去建模。选择多分支而不是单一金字塔的理由很直接如果只在网络末端做特征金字塔前面共享的backbone会被三种尺度的退化拉扯——大雪花希望感受野大、雪雾希望看到全局互相打架。多分支在入口就分流每个transformer只处理一种尺度的退化训练时loss的梯度不会串。2.2 上下文交互模块跨尺度注意力如何把信息喂给彼此这个模块是整个项目的核心创新点解决的问题是三个分支独立处理后如何在融合时不让信息丢失。常见做法是把三路特征拉平后concat接一个1×1卷积降维但这样只是简单叠加尺度之间的对应关系没被利用。上下文交互模块的做法是先把三路token两两配对做交叉注意力比如把大雪花分支的token作为query雪雾分支的token作为key/value——这样大雪花区域融合了雪雾分支的全局亮度信息能判断哪些是真雪花、哪些只是亮度偏低的背景。class CrossScaleAttention(nn.Module): def __init__(self, dim, num_heads8): super().__init__() self.q_proj nn.Linear(dim, dim) self.k_proj nn.Linear(dim, dim) self.v_proj nn.Linear(dim, dim) self.num_heads num_heads def forward(self, query_tokens, context_tokens): # query 来自当前分支context 来自另一个尺度分支 B, N, C query_tokens.shape q self.q_proj(query_tokens).reshape(B, N, self.num_heads, C // self.num_heads).transpose(1, 2) k self.k_proj(context_tokens).reshape(B, N, self.num_heads, C // self.num_heads).transpose(1, 2) v self.v_proj(context_tokens).reshape(B, N, self.num_heads, C // self.num_heads).transpose(1, 2) attn (q k.transpose(-2, -1)) / (C // self.num_heads) ** 0.5 attn attn.softmax(dim-1) out (attn v).transpose(1, 2).reshape(B, N, C) return out这里有个容易忽略的参数query_tokens和context_tokens的序列长度可以不同因为三个分支的patch大小不同token数自然不一样。代码里直接用了同一个N真实跑的时候要对context做插值或者用可变形注意力采样对齐长度我在改源码时是加了一个双线性插值把context resize到query的序列长度。另外num_heads我习惯设8而不是默认的4因为跨尺度交互的信息本身就杂头数太少会退化成「只学到位置偏置、没学到语义交互」。2.3 transformer主干窗口注意力加移动窗口还是全局注意力尺度感知分支出来之后主干网络的选择直接影响显存和效果。原始ViT的全局自注意力在256×256输入下patch 4×4就有4096个token全局注意力矩阵是4096×4096单卡训练非常吃显存。这个项目用的是窗口注意力加移动窗口的混合方案——窗口内做self-attention窗口尺寸7×7窗口移动一个偏移量再做一次两个连续层配合相当于覆盖了跨窗口的信息传递。# 训练时的显存表现batch_size8, 256x256 # 全局注意力: 大约 14GBbatch 只能降到 4 # 窗口注意力移动窗口: 大约 9GBbatch8 稳定这个对比是我在3080Ti上实测的差距实际显存量当然跟模型深度有关但量级差距就是这个比例。窗口注意力的代价是远程信息交互变慢——要堆更多层才能让信息跨越整个图。所以上下文交互模块才放在分支末端做补偿它用交叉注意力把三个尺度的全局信息直接拉通。这就是尺度感知窗口注意力的局部性跨尺度注意力的全局性三者怎么配合的答案局部细节靠窗口全局靠交叉中间层数按显存余量堆。3. 工程落地从数据准备到训练与推理五步把源码跑起来3.1 数据集组织与预处理合成雪怎么生成真实雪怎么标这个项目用的训练数据是公开的合成雪数据集——原理是拿干净图像加随机噪声模拟雪花形状再叠加一层高斯模糊模拟雪雾。但如果你直接拿原版脚本生成的数据训练效果会打折扣因为原版脚本用随机矩形加噪声画雪花形状太规整真实雪花的六角形纹理根本没有。我一般会改生成逻辑雪花形状用随机六边形加扰动大小从1像素到15像素均匀采样雪雾强度按beta分布随机这样训练集里既有稀疏大雪花、又有密集小雪斑、还有浓雾场景。def generate_snow(clean_img, snow_density0.6): h, w clean_img.shape[:2] snow_mask np.zeros((h, w), dtypenp.float32) n_patches int(h * w * snow_density * 0.03) # 雪花数量按面积比例控制 for _ in range(n_patches): x np.random.randint(0, w) y np.random.randint(0, h) r np.random.randint(1, 12) # 雪花半径 1~12 像素 # 六边形雪花用旋转叠加近似 angle np.random.uniform(0, 2 * np.pi) cx, cy r * np.cos(angle), r * np.sin(angle) cv2.circle(snow_mask, (int(x cx), int(y cy)), max(1, r // 3), 1, -1) cv2.circle(snow_mask, (x, y), max(1, r // 2), 1, -1) # 雪雾大半径高斯降低对比度 fog cv2.GaussianBlur(clean_img, (0, 0), sigmaX25, sigmaY25) fog_strength np.random.beta(2, 5) * 0.6 # 大多数场景雪雾较弱少数很强 snowy clean_img * (1 - fog_strength) fog * fog_strength snowy snowy.astype(np.float32) # 雪花区域叠加亮度偏移 snow_mask snow_mask[..., np.newaxis] snowy snowy snow_mask * np.random.uniform(0.4, 0.8) * 255 return np.clip(snowy, 0, 255).astype(np.uint8)参数说明snow_density控制整体雪花密度r控制尺寸范围fog_strength用beta分布是刻意的——大多数图雪雾较轻、少数图极端这样模型不会偏向某个极端。生成后的数据按7:2:1划分训练验证测试验证集里额外放100张人工标注的真实雪图用来防止模型只在合成域上自嗨。3.2 训练脚本核心参数损失函数、优化器与调度策略这部分是我调参最久的地方。项目默认的损失是L1损失训练到第30个epoch之后PSNR基本不动SSIM还在涨但视觉上纹理糊。后来把损失改成L1 感知损失(VGG16 relu3_3层)组合权重比例大约10:1效果变化非常明显。# train.py 中损失函数组合 criterion_l1 nn.L1Loss() criterion_perceptual PerceptualLoss(feature_layerrelu3_3, weights[1.0]) def combined_loss(pred_clean, gt_clean, pred_snow, gt_snow, epoch): l1 criterion_l1(pred_clean, gt_clean) 0.5 * criterion_l1(pred_snow, gt_snow) perc criterion_perceptual(pred_clean, gt_clean) # 前 5 个 epoch 不想让感知损失主导梯度加一个 warmup 权重 warmup min(1.0, epoch / 5) return l1 2.0 * warmup * perc几个关键参数值得单独说。pred_snow是网络额外输出的一个分支——我习惯让网络同时预测干净图和雪残差图残差监督能让中间特征更有辨识度测试时这路输出直接丢掉不影响推理。优化器用AdamWlr1e-4weight_decay1e-4betas(0.9, 0.999)调度器用余弦退火带5个epoch的warmupwarmup期间lr线性从1e-6涨到1e-4。batch size设8输入256×256。如果显存不够优先把batch降到4而不是缩小输入尺寸——缩小输入会改变雪花相对尺寸破坏尺度感知分支的设计逻辑。3.3 推理脚本整图输入还是滑窗边界怎么处理推理阶段最省事的方式是直接整图输入但这个网络对输入尺寸不敏感大图效果不如训练尺寸。原因是尺度感知分支的丰富度是在256×256下学的直接喂2000×1500的大图大雪花分支的patch感受野相对变小退化程度和训练分布不一致。我一般用滑窗推理加重叠拼接def sliding_inference(model, img, window256, stride192): h, w img.shape[:2] out np.zeros((h, w, 3), dtypenp.float32) counter np.zeros((h, w, 1), dtypenp.float32) pad 32 # 边界反射填充避免裁剪到图像外 img_pad cv2.copyMakeBorder(img, pad, pad, pad, pad, cv2.BORDER_REFLECT_101) for y in range(pad, h pad, stride): for x in range(pad, w pad, stride): crop img_pad[y - pad:y window pad, x - pad:x window pad] crop_out model(crop) # 输出干净图 # 只取中心 window 区域重叠区域用 counter 做加权平均 out[y - pad:y - pad window, x - pad:x - pad window] crop_out counter[y - pad:y - pad window, x - pad:x - pad window] 1 return out / counter滑窗参数速查256 窗口 stride256 不重叠快但拼接缝明显 stride192 重叠 1/4视觉上无拼接痕推荐 stride128 重叠 1/2效果最好但推理时间 60%滑窗重叠区域其实可以不用遮罩融合直接等权平均就够。测试下来拼接处不会出现亮度跳变因为这个网络输出的残差是平滑的不像有些去雨网络在拼接处会留下矩形暗区。如果不希望输出有pad痕迹BORDER_REFLECT_101比BORDER_REFLECT好后者在对称轴上会重复边缘像素造成外侧多一层亮线。4. 效果验证PSNR/SSIM之外还要看残差图语义和主观对比4.1 定量指标计算在验证集上如何公平比较跑量化对比时最容易被杠的点是验证集的裁剪方式不一样指标没法复现。这个项目源码里给出的验证脚本做得比较规范——固定了测试种子用滑窗推理而不是整图指标在Y通道YCbCr空间的Y分量上计算。我在此基础上补了UCIQE指标这个指标衡量亮度、色度、对比度的综合质量比PSNR更接近人眼对雪天图的感知。# 验证脚本核心计算 PSNR / SSIM / UCIQE from skimage.metrics import peak_signal_noise_ratio, structural_similarity def evaluate(pred, gt): pred_y rgb2ycbcr(pred)[..., 0] gt_y rgb2ycbcr(gt)[..., 0] psnr peak_signal_noise_ratio(gt_y, pred_y, data_range255) ssim structural_similarity(gt_y, pred_y, data_range255) uciqe calculate_uciqe(pred) # 无参考指标不需要 gt return psnr, ssim, uciqe实际跑完一组对比这个模型的典型表现是PSNR 2831dB取决于合成数据集的雪强度SSIM 0.920.96UCIQE比原图高15%左右。但注意PSNR高不代表视觉干净——L1损失主导时PSNR虚高但纹理被抹平这就是为什么要配合感知损失。验证集里一定要放几组人工标注的真实雪图只看合成指标的模型在真实图上会原形毕露。4.2 残差图检查一秒钟看出模型学废没有正常残差图特征 - 雪花区域有白色/亮色响应边缘清晰但不过锐 - 非雪区域接近纯黑残差接近 0 - 雪雾区域有均匀的亮度抬升 异常模式 - 残差图里出现人脸/文字/物体边缘 → 模型在学“抠图”不是去雪 - 残差图整体发灰 → 模型把亮度偏移当成了雪雾会破坏原图对比度 - 残差图出现方块状 → 窗口注意力的拼接痕迹通常是训练时窗口尺寸和推理尺寸不一致我把残差图保存到验证输出目录每个epoch看一眼。如果发现残差图出现物体边缘说明网络在偷懒——它发现直接输出原图减去高频细节比精确扣除雪区域更容易降低loss。这时候要把感知损失权重加大或者给残差分支加一个稀疏约束L1惩罚残差图中非零区域面积。4.3 消融实验怎么做逐模块关掉量化每个组件的贡献这个项目源码的模型配置是用yaml写的里面有开关可以单独关掉上下文交互模块或尺度感知分支做消融非常方便。我做的消融顺序# 完整模型 python train.py --config configs/full.yaml # 去掉上下文交互模块改成三路特征直接 concat python train.py --config configs/no_context.yaml # 去掉尺度感知分支只用一个 8x8 patch 分支 python train.py --config configs/no_scale.yaml # 只看主干的窗口注意力去掉移动窗口 python train.py --config configs/no_shift.yaml跑完四组对比结果符合预期去掉上下文交互后PSNR掉了约0.8dB但视觉上雪花残影明显变多——这说明上下文交互的作用是减少误判不是提升峰值信噪比。去掉尺度感知分支后雪雾区域会有大面积色偏。如果只想验证一个模块先跑no_scale.yaml和full.yaml这两组就够最省时间。5. 避坑记录四段血泪经验从训练翻车到推理漂移5.1 训练Loss骤降后不再下降PSNR卡在24dB上不去典型现象是第110个epoch loss下降很快PSNR从20冲到24然后像被焊死一样一动不动。当时我怀疑是学习率不够调了三次都无效。原因拆解这个网络有多个分支主干的梯度更新得太快尺度感知分支还没学会分尺度特征主干就开始拟合噪声。更具体的说是L1损失只惩罚像素差在早期主导了梯度方向感知损失被淹没。解决办法把感知损失的权重从2.0改成5.0同时给主干层设一个lr_mult0.5——主干和分支用不同学习率。这里用minswitch或warmup单独控制前10个epoch让分支层的学习率为主干的1.5倍等分支特征成型后再切回统一学习率。另一个应急手段是把输入从256改成192降低任务难度让模型先从简单分布中学会分尺度。5.2 推理时显存溢出batch_size降到1还是OOM现象训练好好的推理大图时显存直接爆掉报CUDA out of memory。源码里推理部分是用整图输入2000×1500的图加上中间层的特征图显存消耗远超训练时的256×256。原因transformer的注意力是在token数平方级别增长的推理图分辨率翻倍中间特征图的token数翻了四倍。显存不是线性涨是平方涨。解决按3.3的滑窗方式分流推理窗口设256、stride设192。如果还嫌慢把窗口缩到224stride设168——这个尺寸下雪花相对大小和训练分布偏离不大实测PSNR只掉0.2dB左右。另外检查一下推理时有没有不小心把torch.no_grad()丢了这个细节最容易翻车。5.3 真实雪图效果崩坏合成数据集训出来的模型在实拍图上全是雪花残影现象合成验证集指标高得吓人一上真实雪景监控图雪花是淡了但留下很多半透明的白色印子像没擦干净的玻璃。原因合成雪数据里雪花是叠加固定亮度偏移生成的残差分支学到的模式是检测亮斑、减掉固定值。真实雪花的亮度受环境光影响变化极大被雪覆盖的物体表面亮度可能比背景低残差分支把暗部的雪当成背景保留了下来。解决一是在合成数据里做亮度扰动——给雪花区域乘以随机系数0.31.2让模型知道雪花没有标准亮度。二是推理后加一个轻量后处理用引导滤波基于原图保边平滑残差图把残差图里那些边界模糊的半透明区域抹掉。引导滤波的半径设8、eps设1e-4这个组合对雪花这种小面积遮挡很有效。5.4 移动窗口拼接处出现细线训练时看不到推理才出现现象推理输出图上每256像素间隔有一条淡淡的竖线或横线像隔行扫描。原因训练时窗口和移动窗口的组合是固定的模型学会了窗口内的局部统计。推理时滑窗的stride不是窗口尺寸的整数倍某些区域落在窗口边界上这些边界处没有移动窗口的二次处理统计特性突变。解决把推理滑窗的stride统一设成窗口尺寸的1/2比如窗口256、stride128并且把输入crop之前加一个随机偏移offset保证每次滑窗的边界不固定。这个偏移量在滑窗循环外生成一次整张图统一偏移即可不要每个窗口独立偏移否则会把拼接缝变成随机噪声。6. 一个实用技巧用冻结编码器的方式做业务场景二次微调拿到这套源码直接全量训练自己的数据容易顾此失彼——合成数据学到的先验会被几十张真实图冲掉。更稳的做法是加载源码提供的预训练权重冻结主干和尺度感知分支只放开最后的融合解码器用你自己的十几张真实雪图微调200个epoch。冻结参数一方面防止小数据过拟合另一方面让预训练的分尺度特征保留下来。实际操作上是给需要更新的层设置requires_grad_(False)优化器只传入未冻结的参数学习率下调到1e-5损失函数先砍掉感知损失只用L1——因为这时模型的语义特征已经成型再多让感知损失主导容易把真实图的细节也抹掉。我习惯在微调前把自己的真实图先做一次直方图匹配让亮度分布靠拢合成数据统计——这个预处理很简单但效果很明显收敛速度能快一倍。200个epoch跑完后做一次残差图检查如果残差图里还是只有雪花区域有响应就说明微调成功。从那以后我每次接手这个项目的业务数据都强制先跑一遍冻结主干、只调解码器的流程再考虑放开全量参数。这套组合拳下来基本不会再遇到微调后模型把真实纹理也删掉的翻车情况。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
