第一次把图像去雨模型跑通是在一个凌晨。我盯着输出图看了很久雨纹确实淡了但树叶纹理全糊成一团像上了一层磨砂滤镜。那时候我才意识到图像去雨这活儿远不是找个网络、丢数据、train这么简单。后来我花了一整周时间把合成数据、网络结构、损失函数和评测方式全部返工才算摸清门道。网上关于图像去雨的PyTorch实现代码不少但大多分两类一类是直接搬公开仓库跑个实验就完事另一类是站在很高的角度讲论文里的数学推导真到自己写代码时还是会卡住。这篇博文不一样我把自己完整跑通的实现过程全部写出来从合成带雨训练集开始到设计一个参数少但可用的卷积网络再到训练、评测、推理部署每一段代码都带详细注释。整套流程全部使用PyTorch实现照着敲就能跑适合刚入门底层图像恢复任务、以及想做去雨方向课程设计或毕设的同学参考。1. 图像去雨到底在解决什么问题以及为什么用PyTorch重写实现1.1 雨纹干扰的本质它不是普通噪声而是有方向的结构信号图像去雨属于底层图像恢复任务和去噪、去模糊、去雾是同一个大类。很多人把雨纹当成一种密度较高的椒盐噪声来处理这是最大的误区。雨纹在成像上的表现是细长的、有一定方向和透明度的高亮线条它们在空间上稀疏但结构感极强且不同景深位置的雨纹模糊程度还不一样。正是这种结构化特征让通用去噪模型在雨图上表现很差。高斯去噪模型假设噪声在像素间是独立的而雨纹的像素点之间有强烈的空间相关性——相邻像素在方向和亮度上高度一致。换句话说通用去噪会把雨纹当成需要保留的边缘细节结果越去雨雨纹越清晰或者把真正的物体边缘给磨掉。这也是为什么需要专门设计面向雨纹结构的小网络而不是套一个大号的通用CNN。另外一个容易被忽略的点是雨纹往往叠加在背景纹理之上这给网络出了道难题网络要区分高亮的雨线和高亮的物体边缘两者在局部形态上可能很接近。解决这个问题的思路之一是扩大感受野让网络能在更大的上下文里判断一个像素是否属于雨线。具体怎么做会在第三章的网络设计里展开。1.2 开发环境与依赖版本说明我选PyTorch来实现这套代码原因很实际动态图调试太方便了网络中间任何一层的输出都能直接print出来看而且数据加载、GPU训练、模型导出这套生态非常成熟社区里找任何一个小模块都有现成参考。我的环境配置如下仅供对照参考Python 3.9PyTorch 1.12.1CUDA 11.3版本CPU也能跑只是慢torchvision 0.13.1opencv-python 4.6.0numpy 1.23.0tqdm 4.64.0import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader import cv2 import numpy as np import os import random from tqdm import tqdm # 固定随机种子保证每次实验结果可复现 def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42)注意OpenCV在读取图片时返回的是BGR通道顺序而PyTorch训练时一般用RGB。这个坑会导致训练时颜色错乱、推理时输出偏色所以我在代码里统一用cv2.cvtColor(img, cv2.COLOR_BGR2RGB)做转换。2. 训练数据从零合成模拟雨纹背后的成像逻辑2.1 为什么没直接下载公开数据集而是自己合成雨图公开的图像去雨数据集确实不少比如Rain100L、Rain100H、DID-MDN配套数据但实际用起来有几个烦人的问题文件格式不统一有的是.mat文件有的是多张图打包压缩每张图的尺寸差异大需要额外写脚本统一处理部分数据集下载源不稳定。更关键的是自己合成雨图能精确控制雨纹密度、方向、长度和透明度后面做消融实验时可以单独验证模型对斜向雨的鲁棒性或者雨纹太密集时效果如何下降。这种控制力是现成数据集给不了的。所以我采用了最直观的合成方案在干净图像上随机绘制多方向、多长度的半透明线条来模拟雨纹。这个方案不完全等同于真实雨图成像但作为训练数据已经足够让模型学到什么是雨纹、什么不是雨纹。2.2 雨纹生成算法与带注释的完整代码雨纹模拟的核心逻辑分三步生成随机方向与长度的线段、将线段绘制到掩码图上、对掩码做高斯模糊后与原图融合。高斯模糊这一步很关键直接画硬线条会得到边缘锐利的假雨纹网络学到的特征和真实雨纹不匹配。def add_rain(image, num_drops30, max_length30, angle_range(-30, 30)): 在干净图像上合成雨纹 image: HWC格式RGB图像取值范围0-255dtypeuint8 num_drops: 雨线数量 max_length: 雨线最大像素长度 angle_range: 雨线角度范围[-30, 30]模拟自然降雨的倾角 h, w, _ image.shape # 雨线掩码图先画线再做高斯模糊 rain_mask np.zeros((h, w, 3), dtypenp.float32) for _ in range(num_drops): # 随机选择雨线的起点 x0 np.random.randint(0, w) y0 np.random.randint(0, h) # 随机雨线长度控制在max_length的30%-100%之间 length np.random.randint(int(max_length * 0.3), max_length 1) # 随机雨线倾角在给定的范围内 angle np.deg2rad(np.random.uniform(angle_range[0], angle_range[1])) # 计算终点雨线通常向下倾斜所以dx和dy都取正值 dx int(length * np.sin(angle)) dy int(length * np.cos(angle)) x1 min(w - 1, max(0, x0 dx)) y1 min(h - 1, max(0, y0 dy)) # 在掩码图上画一条粗度为1的白色线段 # 后续通过高斯模糊把线变成有厚度的雨带 cv2.line(rain_mask, (x0, y0), (x1, y1), (1.0, 1.0, 1.0), 1) # 高斯模糊让雨线变得柔和模拟雨滴在镜头中的散射效果 rain_mask cv2.GaussianBlur(rain_mask, (5, 5), 1.0) # 雨纹是亮度增强所以是加法而非乘法 # alpha是雨纹强度系数取0.6-1.0之间的随机值 alpha np.random.uniform(0.6, 1.0) rainy_image image.astype(np.float32) alpha * rain_mask * 255.0 rainy_image np.clip(rainy_image, 0, 255).astype(np.uint8) return rainy_image合成雨图的几个参数需要说下num_drops控制雨的密度太少了模型学不到东西太多了模型会把背景细节当成雨纹一并抹掉我实测在30左右比较平衡max_length控制雨线长度这是和真实雨图差别最大的参数真实图中近处雨滴的轨迹长且清晰远处则短而模糊但作为基础版本固定长度范围已经够用了。2.3 Dataset与DataLoader实现数据增强的细节决定训练稳定性有了单张图合成雨纹的函数接下来把它们串成PyTorch标准的Dataset接口。这一步很多人会忽略数据增强的作用这里做了随机裁剪和随机水平翻转理由很简单如果让网络每轮都看到完整的原始尺寸图它很容易记住某一批图的位置特征导致验证集效果虚高泛化能力却很一般。class RainDataset(Dataset): 从干净图像目录读取图片在训练时动态生成雨图 def __init__(self, clean_img_dir, patch_size224, is_trainTrue): self.clean_img_dir clean_img_dir self.image_paths [ os.path.join(clean_img_dir, f) for f in os.listdir(clean_img_dir) if f.endswith((.png, .jpg, .jpeg, .bmp)) ] self.patch_size patch_size self.is_train is_train def __len__(self): return len(self.image_paths) def __getitem__(self, idx): # 读取干净图像并转为RGB img_path self.image_paths[idx] clean cv2.imread(img_path, cv2.IMREAD_COLOR) if clean is None: raise ValueError(f图片读取失败: {img_path}) clean cv2.cvtColor(clean, cv2.COLOR_BGR2RGB) if self.is_train: # 训练时随机裁剪到patch_size并随机水平翻转 h, w, _ clean.shape if h self.patch_size or w self.patch_size: clean cv2.resize(clean, (self.patch_size, self.patch_size), interpolationcv2.INTER_AREA) else: x np.random.randint(0, w - self.patch_size 1) y np.random.randint(0, h - self.patch_size 1) clean clean[y:y self.patch_size, x:x self.patch_size] if np.random.rand() 0.5: clean cv2.flip(clean, 1) # 合成雨图 rainy add_rain(clean) # BGR顺序转RGB这里已经转好了 # 转成Tensor并归一化到[0, 1]区间 clean_tensor torch.from_numpy(clean.transpose(2, 0, 1)).float() / 255.0 rainy_tensor torch.from_numpy(rainy.transpose(2, 0, 1)).float() / 255.0 return rainy_tensor, clean_tensor# 使用方式 train_dataset RainDataset(clean_img_dir./data/clean_images/train, patch_size224, is_trainTrue) val_dataset RainDataset(clean_img_dir./data/clean_images/val, patch_size224, is_trainFalse) train_loader DataLoader(train_dataset, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size8, shuffleFalse, num_workers2, pin_memoryTrue)提示验证时我没有做随机裁剪而是直接resize到224x224是为了保证每次验证的输出尺寸一致从而计算稳定的PSNR和SSIM。如果你希望验证结果更严谨也可以使用滑窗裁剪再把所有结果拼回来。3. 网络模型设计让网络学习“雨纹残差”而不是直接生成背景3.1 残差学习的出发点擦黑板思维网络结构上我选择残差学习的思路。直接让网络输出干净背景图当然可以但这对网络要求太高它需要同时完成识别雨纹区域和生成被遮挡的背景纹理两件事尤其是背景生成部分本质上是一个图像重建问题需要网络有很强的生成能力计算开销和过拟合风险都会上升。残差学习的思路完全不同。雨图可以近似看成干净背景 雨纹叠加那么从雨图中减去雨纹就得到干净背景。网络只需要预测雨纹部分最后用输入减去预测结果即可。这就像擦黑板我们不需要完整重画黑板上被粉笔字覆盖的内容只需要把粉笔字擦掉底下的内容自然就露出来了。好处是显而易见的网络要预测的目标从整张真实图像变成一个稀疏的、结构化的雨纹图训练难度大幅下降。这个思路在图像去噪中同样有效——预测噪声残差比预测干净图像更容易收敛。3.2 网络结构与带注释的PyTorch代码网络主体是一个4层卷积堆叠的小网络。我没有选择很深的结构原因有两个去雨任务中雨纹是局部特征感受野不需要覆盖整幅图像层数太多容易丢失高频细节而高频细节里有大量的背景纹理。class DerainNet(nn.Module): 一个用于图像去雨的简单卷积基线网络 输入: 3通道雨图 (B, 3, H, W) 输出: 3通道去雨图 (B, 3, H, W) def __init__(self): super(DerainNet, self).__init__() # 四层卷积通道数从3 - 32 - 32 - 32 - 3 # 卷积核大小统一用3x3padding1保证特征图尺寸不变 self.net nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(32, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(32, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(32, 3, kernel_size3, padding1), ) # kaiming初始化能保证前向传播时激活值不会快速消失或膨胀 for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.constant_(m.bias, 0) def forward(self, x): # 网络只学习雨纹残差 residual self.net(x) # 输入减去残差得到去雨结果 return x - residual感受野计算一下三层3x3卷积堆叠后最远能看到7x7的区域这个大小对捕获一条宽1-2像素、长十几像素的雨线来说刚刚够用。如果你发现模型对长雨线效果不好可以再加一层卷积感受野会扩大到9x9。注意网络最后没有任何激活函数。因为残差值既可能是正数雨纹增亮区域也可能是负数如果雨纹带轻微阴影效果线性输出比ReLU更合理。4. 训练循环与损失设计精度提升的关键细节4.1 损失函数从MSE到MSE加SSIM的组合策略训练图像去雨模型最常用的损失函数是MSE均方误差。原因很直接MSE和PSNR在数学上是对应关系PSNR越高代表像素级误差越小而MSE越小PSNR越高。但只用MSE有个明显问题它假设每个像素是独立的不考虑像素间的结构关系。训练后期你会发现PSNR在涨但视觉上图像变油了——细节被抹平边缘不够锐利。这是因为MSE对局部结构的轻微错位惩罚不够大。我的建议是加一个简单的SSIM损失作为辅助SSIM衡量的是两个图像块在亮度、对比度和结构三个维度的相似度能迫使网络保留更多结构信息。不过SSIM计算代价略高而且在不同实现里公式细节有差异初学者可以先只用MSE跑通基线再逐步加入结构损失。4.2 训练主循环与Checkpoint保存每个环节都有代码注释# 损失函数MSE这里不直接定义而是在训练循环里调用 criterion nn.MSELoss() model DerainNet().cuda() optimizer optim.Adam(model.parameters(), lr1e-3) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50, eta_min1e-6) def train_one_epoch(model, loader, criterion, optimizer, epoch): model.train() total_loss 0.0 pbar tqdm(loader, descfEpoch {epoch1}) for rainy, clean in pbar: rainy rainy.cuda() clean clean.cuda() # 前向传播 output model(rainy) # 计算MSE损失 loss criterion(output, clean) # 梯度清零反向传播更新权重 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() pbar.set_postfix(lossloss.item()) return total_loss / len(loader) def validate(model, loader): 验证函数计算模型在验证集上的PSNR 这里不开梯度节省显存并加速 model.eval() total_psnr 0.0 with torch.no_grad(): for rainy, clean in loader: rainy rainy.cuda() clean clean.cuda() output model(rainy) # PSNR计算是基于0-1范围的均方误差 mse torch.mean((output - clean) ** 2, dim[1, 2, 3]) psnr 10 * torch.log10(1.0 / (mse 1e-10)) total_psnr psnr.sum().item() return total_psnr / len(loader.dataset)然后是完整的训练脚本我把模型保存、学习率调整都写进去epochs 60 best_psnr 0.0 for epoch in range(epochs): train_loss train_one_epoch(model, train_loader, criterion, optimizer, epoch) val_psnr validate(model, val_loader) # 学习率按照余弦退火调度调整 scheduler.step() print(fEpoch {epoch1}, Loss: {train_loss:.6f}, Val PSNR: {val_psnr:.2f}) # 保存验证集上PSNR最高的模型 if val_psnr best_psnr: best_psnr val_psnr torch.save(model.state_dict(), ./best_model.pth) print(f保存新模型PSNR: {val_psnr:.2f})4.3 调参实战经验几个影响训练质量的隐性因素第一batch size的选择。我用8再大当然可以但8GB显存的显卡跑224x224输入时batch size建议不超过16。如果显存不够优先降低batch size而不是降低图像分辨率因为裁切尺寸变小会让网络看到的信息变少泛化性下降。第二学习率。用Adam时初始学习率1e-3是比较稳的但到了40个epoch之后如果发现loss在缓慢震荡而PSNR不再上升通常不是模型问题而是学习率太大了。我用的余弦退火调度器会平滑地把学习率降到1e-6整体训练曲线很稳定。第三如果训练集只有几十张图网络很容易过拟合。我的经验是每个epoch做随机裁剪和翻转等效于把训练集扩大了很多倍但如果源图太少还是需要换用ImageNet预训练权重做迁移学习单纯靠加大增强力度效果有限。提示检查过拟合的方法很简单打印训练集和验证集的PSNR。如果训练集PSNR持续涨、验证集不涨甚至下跌基本就是过拟合优先考虑增加数据增强力度或者缩小网络容量。5. 评测指标与结果可视化PSNR高不等于效果好5.1 PSNR和SSIM的纯PyTorch实现避开sklearn依赖很多评测代码依赖skimage.metrics.peak_signal_noise_ratio和structural_similarity但我更推荐自己写一个轻量版本一来不依赖额外库二来能彻底理解指标的计算原理。PSNR的定义是峰值信噪比对于0-1范围的数据公式是PSNR 10 * log10(1 / MSE)def compute_psnr(img1, img2): 计算两张0-1范围内图像的PSNR img1, img2: PyTorch Tensor, 形状为 (B, 3, H, W) mse torch.mean((img1 - img2) ** 2) if mse.item() 1e-10: return 100.0 # 完全一致时给出一个很大的值 psnr 10.0 * torch.log10(1.0 / mse) return psnr.item()SSIM稍微复杂一点我不想贴一个过于臃肿的实现这里给出一个简化版本它用高斯窗口计算局部均值和方差再合成结构相似度def gaussian_kernel(size11, sigma1.5): 生成一维高斯核并外积成二维核用于SSIM计算 coords torch.arange(size, dtypetorch.float32) - size // 2 g torch.exp(-(coords ** 2) / (2 * sigma ** 2)) g / g.sum() return g.outer(g).view(1, 1, size, size) def compute_ssim(img1, img2, window_size11, sigma1.5): 简化版SSIM忽略通道合并逐通道计算后取平均 img1, img2: 0-1范围的Tensor形状(B, 3, H, W) device img1.device window gaussian_kernel(window_size, sigma).to(device) # 常数C1, C2防止除零 C1 0.01 ** 2 C2 0.03 ** 2 # 计算均值 mu1 torch.nn.functional.conv2d(img1, window, paddingwindow_size // 2) mu2 torch.nn.functional.conv2d(img2, window, paddingwindow_size // 2) mu1_sq, mu2_sq mu1 ** 2, mu2 ** 2 mu1_mu2 mu1 * mu2 # 计算方差和协方差 sigma1_sq torch.nn.functional.conv2d(img1 * img1, window, paddingwindow_size // 2) - mu1_sq sigma2_sq torch.nn.functional.conv2d(img2 * img2, window, paddingwindow_size // 2) - mu2_sq sigma12 torch.nn.functional.conv2d(img1 * img2, window, paddingwindow_size // 2) - mu1_mu2 ssim_map ((2 * mu1_mu2 C1) * (2 * sigma12 C2)) / \ ((mu1_sq mu2_sq C1) * (sigma1_sq sigma2_sq C2)) return ssim_map.mean().item()5.2 结果可视化把雨图、去雨图、干净图画一起比只看数值更靠谱评估一个去雨模型不能只盯PSNR。我自己吃过这个亏模型PSNR到了32dB但肉眼一看背景细节已经糊成水彩画了。PSNR只衡量像素差异的全局平均值它不能反映局部结构是否被抹平。所以我每次验证都做一张三连图左边是雨图输入中间是模型输出右边是干净原图。代码很简单def save_comparison(rainy, output, clean, save_path): 将输入雨图、模型输出、干净原图横向拼接保存 rainy/output/clean: Tensor (3, H, W)0-1范围 # 转成numpy并还原到0-255范围 rainy_np rainy.cpu().numpy().transpose(1, 2, 0) * 255.0 output_np output.cpu().numpy().transpose(1, 2, 0) * 255.0 clean_np clean.cpu().numpy().transpose(1, 2, 0) * 255.0 combo np.hstack([rainy_np, output_np, clean_np]).astype(np.uint8) # 保存为RGB图注意要转回BGR给cv2 combo_bgr cv2.cvtColor(combo, cv2.COLOR_RGB2BGR) cv2.imwrite(save_path, combo_bgr)然后每个epoch结束时挑几张验证图跑一次保存我就能直观看到模型有没有出现过度平滑的现象。如果输出图像像被水洗过一样那说明要调整损失函数比如加大SSIM的权重。6. 推理部署与常见问题排查从训练到真正能用的一段路6.1 单张图片推理全流程注意归一化必须和训练一致模型训练好了最终要拿去处理任意一张图。推理流程和训练时有几个细节不同我在这里完整写一遍同时也解释为什么有些细节容易出错。def inference_single_image(model, image_path, output_path, use_gpuTrue): 对单张雨图做去雨推理 device torch.device(cuda if use_gpu and torch.cuda.is_available() else cpu) model model.to(device) model.eval() # 1. 读取图像并转为RGB img cv2.imread(image_path) if img is None: raise ValueError(f无法读取图片: {image_path}) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 2. 保存原图尺寸推理后要把输出resize回原尺寸 h, w, _ img.shape img_resized cv2.resize(img_rgb, (224, 224), interpolationcv2.INTER_AREA) # 3. 转Tensor并归一化到0-1范围增加batch维度 input_tensor torch.from_numpy(img_resized.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0 input_tensor input_tensor.to(device) # 4. 推理并禁止计算梯度 with torch.no_grad(): output_tensor model(input_tensor) # 5. 处理后保存先转回numpyclamp到0-1区间再resize回原始尺寸 output_np output_tensor.squeeze(0).cpu().numpy().transpose(1, 2, 0) output_np np.clip(output_np, 0.0, 1.0) * 255.0 output_np output_np.astype(np.uint8) output_rgb cv2.resize(output_np, (w, h), interpolationcv2.INTER_LINEAR) # 6. RGB转BGR存盘 output_bgr cv2.cvtColor(output_rgb, cv2.COLOR_RGB2BGR) cv2.imwrite(output_path, output_bgr)这里有个容易踩的坑如果训练时用了归一化中常见的mean/std操作推理时也必须用完全相同的mean/std。我这个代码训练时只除以了255没有做mean/std归一化所以推理时也保持一致的归一化方式。如果两边不一致模型输出的颜色会偏移而且很难发现原因。6.2 推理阶段常踩的几个坑输入尺寸不一致的问题。我在上面用了resize到224x224再推理的方法简单但会丢失细节尤其在大图上效果不好。更好的做法是滑窗推理把大图切成若干224x224的patch分别推理最后拼回原尺寸。代价是推理时间变长但对细节保留非常明显。另一个坑是输出图像的暗部偏色。如果训练时合成雨纹的alpha值总是偏大模型会在真实无雨图片上产生去雨过猛的问题把暗部细节也当成雨纹去掉。这就需要在训练数据里混入一些不含雨纹的干净图片让模型学会没有雨纹时保持原样。6.3 从基线到更好效果的三条改进路径这套基线模型的作用是跑通全流程如果要做更深入的课题我有三条建议第一把普通卷积块换成残差块或者密集连接块训练稳定性会有明显提升。第二在卷积之后加一个简单的通道注意力模块让网络学会重点处理雨纹出现概率高的通道。第三把MSE损失换成Charbonnier损失配合SSIM损失对离群点的抗干扰性更强。不过我必须强调一点在追求更先进结构之前先把数据合成质量和评测流程做扎实。很多人在小改动上纠结半天结果发现是合成雨纹的程序写错了这亏我吃过。写在最后关于这套代码的一些个人心得整套流程跑下来我最想分享的体会是图像去雨的提升瓶颈往往不在网络结构而在看不见的细节里。数据合成时雨纹密度和长度的分布决定了模型的上限训练时损失函数的结构约束决定了输出是不是耐看评测时如果不看主观效果很容易被PSNR欺骗。我建议每一个做这个方向的同学都从合成数据开始亲手搭一遍这整套代码训练过程中每10个epoch保存一次验证输出图盯着图像看趋势比只看曲线靠谱得多。跑通之后你会对残差学习、感受野、损失函数设计这些概念有完全不同的理解。最后分享一个小技巧训练接近收敛时把验证集中效果最差的几张图挑出来看。分析这些图面的共同点往往能准确指出当前模型的短板——是长雨纹处理不了还是密集雨纹区域的背景纹理被抹掉。针对短板去调整数据或网络比盲目堆参数高效得多。
PyTorch图像去雨实战：从合成数据到模型部署全流程解析
第一次把图像去雨模型跑通是在一个凌晨。我盯着输出图看了很久雨纹确实淡了但树叶纹理全糊成一团像上了一层磨砂滤镜。那时候我才意识到图像去雨这活儿远不是找个网络、丢数据、train这么简单。后来我花了一整周时间把合成数据、网络结构、损失函数和评测方式全部返工才算摸清门道。网上关于图像去雨的PyTorch实现代码不少但大多分两类一类是直接搬公开仓库跑个实验就完事另一类是站在很高的角度讲论文里的数学推导真到自己写代码时还是会卡住。这篇博文不一样我把自己完整跑通的实现过程全部写出来从合成带雨训练集开始到设计一个参数少但可用的卷积网络再到训练、评测、推理部署每一段代码都带详细注释。整套流程全部使用PyTorch实现照着敲就能跑适合刚入门底层图像恢复任务、以及想做去雨方向课程设计或毕设的同学参考。1. 图像去雨到底在解决什么问题以及为什么用PyTorch重写实现1.1 雨纹干扰的本质它不是普通噪声而是有方向的结构信号图像去雨属于底层图像恢复任务和去噪、去模糊、去雾是同一个大类。很多人把雨纹当成一种密度较高的椒盐噪声来处理这是最大的误区。雨纹在成像上的表现是细长的、有一定方向和透明度的高亮线条它们在空间上稀疏但结构感极强且不同景深位置的雨纹模糊程度还不一样。正是这种结构化特征让通用去噪模型在雨图上表现很差。高斯去噪模型假设噪声在像素间是独立的而雨纹的像素点之间有强烈的空间相关性——相邻像素在方向和亮度上高度一致。换句话说通用去噪会把雨纹当成需要保留的边缘细节结果越去雨雨纹越清晰或者把真正的物体边缘给磨掉。这也是为什么需要专门设计面向雨纹结构的小网络而不是套一个大号的通用CNN。另外一个容易被忽略的点是雨纹往往叠加在背景纹理之上这给网络出了道难题网络要区分高亮的雨线和高亮的物体边缘两者在局部形态上可能很接近。解决这个问题的思路之一是扩大感受野让网络能在更大的上下文里判断一个像素是否属于雨线。具体怎么做会在第三章的网络设计里展开。1.2 开发环境与依赖版本说明我选PyTorch来实现这套代码原因很实际动态图调试太方便了网络中间任何一层的输出都能直接print出来看而且数据加载、GPU训练、模型导出这套生态非常成熟社区里找任何一个小模块都有现成参考。我的环境配置如下仅供对照参考Python 3.9PyTorch 1.12.1CUDA 11.3版本CPU也能跑只是慢torchvision 0.13.1opencv-python 4.6.0numpy 1.23.0tqdm 4.64.0import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader import cv2 import numpy as np import os import random from tqdm import tqdm # 固定随机种子保证每次实验结果可复现 def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42)注意OpenCV在读取图片时返回的是BGR通道顺序而PyTorch训练时一般用RGB。这个坑会导致训练时颜色错乱、推理时输出偏色所以我在代码里统一用cv2.cvtColor(img, cv2.COLOR_BGR2RGB)做转换。2. 训练数据从零合成模拟雨纹背后的成像逻辑2.1 为什么没直接下载公开数据集而是自己合成雨图公开的图像去雨数据集确实不少比如Rain100L、Rain100H、DID-MDN配套数据但实际用起来有几个烦人的问题文件格式不统一有的是.mat文件有的是多张图打包压缩每张图的尺寸差异大需要额外写脚本统一处理部分数据集下载源不稳定。更关键的是自己合成雨图能精确控制雨纹密度、方向、长度和透明度后面做消融实验时可以单独验证模型对斜向雨的鲁棒性或者雨纹太密集时效果如何下降。这种控制力是现成数据集给不了的。所以我采用了最直观的合成方案在干净图像上随机绘制多方向、多长度的半透明线条来模拟雨纹。这个方案不完全等同于真实雨图成像但作为训练数据已经足够让模型学到什么是雨纹、什么不是雨纹。2.2 雨纹生成算法与带注释的完整代码雨纹模拟的核心逻辑分三步生成随机方向与长度的线段、将线段绘制到掩码图上、对掩码做高斯模糊后与原图融合。高斯模糊这一步很关键直接画硬线条会得到边缘锐利的假雨纹网络学到的特征和真实雨纹不匹配。def add_rain(image, num_drops30, max_length30, angle_range(-30, 30)): 在干净图像上合成雨纹 image: HWC格式RGB图像取值范围0-255dtypeuint8 num_drops: 雨线数量 max_length: 雨线最大像素长度 angle_range: 雨线角度范围[-30, 30]模拟自然降雨的倾角 h, w, _ image.shape # 雨线掩码图先画线再做高斯模糊 rain_mask np.zeros((h, w, 3), dtypenp.float32) for _ in range(num_drops): # 随机选择雨线的起点 x0 np.random.randint(0, w) y0 np.random.randint(0, h) # 随机雨线长度控制在max_length的30%-100%之间 length np.random.randint(int(max_length * 0.3), max_length 1) # 随机雨线倾角在给定的范围内 angle np.deg2rad(np.random.uniform(angle_range[0], angle_range[1])) # 计算终点雨线通常向下倾斜所以dx和dy都取正值 dx int(length * np.sin(angle)) dy int(length * np.cos(angle)) x1 min(w - 1, max(0, x0 dx)) y1 min(h - 1, max(0, y0 dy)) # 在掩码图上画一条粗度为1的白色线段 # 后续通过高斯模糊把线变成有厚度的雨带 cv2.line(rain_mask, (x0, y0), (x1, y1), (1.0, 1.0, 1.0), 1) # 高斯模糊让雨线变得柔和模拟雨滴在镜头中的散射效果 rain_mask cv2.GaussianBlur(rain_mask, (5, 5), 1.0) # 雨纹是亮度增强所以是加法而非乘法 # alpha是雨纹强度系数取0.6-1.0之间的随机值 alpha np.random.uniform(0.6, 1.0) rainy_image image.astype(np.float32) alpha * rain_mask * 255.0 rainy_image np.clip(rainy_image, 0, 255).astype(np.uint8) return rainy_image合成雨图的几个参数需要说下num_drops控制雨的密度太少了模型学不到东西太多了模型会把背景细节当成雨纹一并抹掉我实测在30左右比较平衡max_length控制雨线长度这是和真实雨图差别最大的参数真实图中近处雨滴的轨迹长且清晰远处则短而模糊但作为基础版本固定长度范围已经够用了。2.3 Dataset与DataLoader实现数据增强的细节决定训练稳定性有了单张图合成雨纹的函数接下来把它们串成PyTorch标准的Dataset接口。这一步很多人会忽略数据增强的作用这里做了随机裁剪和随机水平翻转理由很简单如果让网络每轮都看到完整的原始尺寸图它很容易记住某一批图的位置特征导致验证集效果虚高泛化能力却很一般。class RainDataset(Dataset): 从干净图像目录读取图片在训练时动态生成雨图 def __init__(self, clean_img_dir, patch_size224, is_trainTrue): self.clean_img_dir clean_img_dir self.image_paths [ os.path.join(clean_img_dir, f) for f in os.listdir(clean_img_dir) if f.endswith((.png, .jpg, .jpeg, .bmp)) ] self.patch_size patch_size self.is_train is_train def __len__(self): return len(self.image_paths) def __getitem__(self, idx): # 读取干净图像并转为RGB img_path self.image_paths[idx] clean cv2.imread(img_path, cv2.IMREAD_COLOR) if clean is None: raise ValueError(f图片读取失败: {img_path}) clean cv2.cvtColor(clean, cv2.COLOR_BGR2RGB) if self.is_train: # 训练时随机裁剪到patch_size并随机水平翻转 h, w, _ clean.shape if h self.patch_size or w self.patch_size: clean cv2.resize(clean, (self.patch_size, self.patch_size), interpolationcv2.INTER_AREA) else: x np.random.randint(0, w - self.patch_size 1) y np.random.randint(0, h - self.patch_size 1) clean clean[y:y self.patch_size, x:x self.patch_size] if np.random.rand() 0.5: clean cv2.flip(clean, 1) # 合成雨图 rainy add_rain(clean) # BGR顺序转RGB这里已经转好了 # 转成Tensor并归一化到[0, 1]区间 clean_tensor torch.from_numpy(clean.transpose(2, 0, 1)).float() / 255.0 rainy_tensor torch.from_numpy(rainy.transpose(2, 0, 1)).float() / 255.0 return rainy_tensor, clean_tensor# 使用方式 train_dataset RainDataset(clean_img_dir./data/clean_images/train, patch_size224, is_trainTrue) val_dataset RainDataset(clean_img_dir./data/clean_images/val, patch_size224, is_trainFalse) train_loader DataLoader(train_dataset, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size8, shuffleFalse, num_workers2, pin_memoryTrue)提示验证时我没有做随机裁剪而是直接resize到224x224是为了保证每次验证的输出尺寸一致从而计算稳定的PSNR和SSIM。如果你希望验证结果更严谨也可以使用滑窗裁剪再把所有结果拼回来。3. 网络模型设计让网络学习“雨纹残差”而不是直接生成背景3.1 残差学习的出发点擦黑板思维网络结构上我选择残差学习的思路。直接让网络输出干净背景图当然可以但这对网络要求太高它需要同时完成识别雨纹区域和生成被遮挡的背景纹理两件事尤其是背景生成部分本质上是一个图像重建问题需要网络有很强的生成能力计算开销和过拟合风险都会上升。残差学习的思路完全不同。雨图可以近似看成干净背景 雨纹叠加那么从雨图中减去雨纹就得到干净背景。网络只需要预测雨纹部分最后用输入减去预测结果即可。这就像擦黑板我们不需要完整重画黑板上被粉笔字覆盖的内容只需要把粉笔字擦掉底下的内容自然就露出来了。好处是显而易见的网络要预测的目标从整张真实图像变成一个稀疏的、结构化的雨纹图训练难度大幅下降。这个思路在图像去噪中同样有效——预测噪声残差比预测干净图像更容易收敛。3.2 网络结构与带注释的PyTorch代码网络主体是一个4层卷积堆叠的小网络。我没有选择很深的结构原因有两个去雨任务中雨纹是局部特征感受野不需要覆盖整幅图像层数太多容易丢失高频细节而高频细节里有大量的背景纹理。class DerainNet(nn.Module): 一个用于图像去雨的简单卷积基线网络 输入: 3通道雨图 (B, 3, H, W) 输出: 3通道去雨图 (B, 3, H, W) def __init__(self): super(DerainNet, self).__init__() # 四层卷积通道数从3 - 32 - 32 - 32 - 3 # 卷积核大小统一用3x3padding1保证特征图尺寸不变 self.net nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(32, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(32, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(32, 3, kernel_size3, padding1), ) # kaiming初始化能保证前向传播时激活值不会快速消失或膨胀 for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.constant_(m.bias, 0) def forward(self, x): # 网络只学习雨纹残差 residual self.net(x) # 输入减去残差得到去雨结果 return x - residual感受野计算一下三层3x3卷积堆叠后最远能看到7x7的区域这个大小对捕获一条宽1-2像素、长十几像素的雨线来说刚刚够用。如果你发现模型对长雨线效果不好可以再加一层卷积感受野会扩大到9x9。注意网络最后没有任何激活函数。因为残差值既可能是正数雨纹增亮区域也可能是负数如果雨纹带轻微阴影效果线性输出比ReLU更合理。4. 训练循环与损失设计精度提升的关键细节4.1 损失函数从MSE到MSE加SSIM的组合策略训练图像去雨模型最常用的损失函数是MSE均方误差。原因很直接MSE和PSNR在数学上是对应关系PSNR越高代表像素级误差越小而MSE越小PSNR越高。但只用MSE有个明显问题它假设每个像素是独立的不考虑像素间的结构关系。训练后期你会发现PSNR在涨但视觉上图像变油了——细节被抹平边缘不够锐利。这是因为MSE对局部结构的轻微错位惩罚不够大。我的建议是加一个简单的SSIM损失作为辅助SSIM衡量的是两个图像块在亮度、对比度和结构三个维度的相似度能迫使网络保留更多结构信息。不过SSIM计算代价略高而且在不同实现里公式细节有差异初学者可以先只用MSE跑通基线再逐步加入结构损失。4.2 训练主循环与Checkpoint保存每个环节都有代码注释# 损失函数MSE这里不直接定义而是在训练循环里调用 criterion nn.MSELoss() model DerainNet().cuda() optimizer optim.Adam(model.parameters(), lr1e-3) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50, eta_min1e-6) def train_one_epoch(model, loader, criterion, optimizer, epoch): model.train() total_loss 0.0 pbar tqdm(loader, descfEpoch {epoch1}) for rainy, clean in pbar: rainy rainy.cuda() clean clean.cuda() # 前向传播 output model(rainy) # 计算MSE损失 loss criterion(output, clean) # 梯度清零反向传播更新权重 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() pbar.set_postfix(lossloss.item()) return total_loss / len(loader) def validate(model, loader): 验证函数计算模型在验证集上的PSNR 这里不开梯度节省显存并加速 model.eval() total_psnr 0.0 with torch.no_grad(): for rainy, clean in loader: rainy rainy.cuda() clean clean.cuda() output model(rainy) # PSNR计算是基于0-1范围的均方误差 mse torch.mean((output - clean) ** 2, dim[1, 2, 3]) psnr 10 * torch.log10(1.0 / (mse 1e-10)) total_psnr psnr.sum().item() return total_psnr / len(loader.dataset)然后是完整的训练脚本我把模型保存、学习率调整都写进去epochs 60 best_psnr 0.0 for epoch in range(epochs): train_loss train_one_epoch(model, train_loader, criterion, optimizer, epoch) val_psnr validate(model, val_loader) # 学习率按照余弦退火调度调整 scheduler.step() print(fEpoch {epoch1}, Loss: {train_loss:.6f}, Val PSNR: {val_psnr:.2f}) # 保存验证集上PSNR最高的模型 if val_psnr best_psnr: best_psnr val_psnr torch.save(model.state_dict(), ./best_model.pth) print(f保存新模型PSNR: {val_psnr:.2f})4.3 调参实战经验几个影响训练质量的隐性因素第一batch size的选择。我用8再大当然可以但8GB显存的显卡跑224x224输入时batch size建议不超过16。如果显存不够优先降低batch size而不是降低图像分辨率因为裁切尺寸变小会让网络看到的信息变少泛化性下降。第二学习率。用Adam时初始学习率1e-3是比较稳的但到了40个epoch之后如果发现loss在缓慢震荡而PSNR不再上升通常不是模型问题而是学习率太大了。我用的余弦退火调度器会平滑地把学习率降到1e-6整体训练曲线很稳定。第三如果训练集只有几十张图网络很容易过拟合。我的经验是每个epoch做随机裁剪和翻转等效于把训练集扩大了很多倍但如果源图太少还是需要换用ImageNet预训练权重做迁移学习单纯靠加大增强力度效果有限。提示检查过拟合的方法很简单打印训练集和验证集的PSNR。如果训练集PSNR持续涨、验证集不涨甚至下跌基本就是过拟合优先考虑增加数据增强力度或者缩小网络容量。5. 评测指标与结果可视化PSNR高不等于效果好5.1 PSNR和SSIM的纯PyTorch实现避开sklearn依赖很多评测代码依赖skimage.metrics.peak_signal_noise_ratio和structural_similarity但我更推荐自己写一个轻量版本一来不依赖额外库二来能彻底理解指标的计算原理。PSNR的定义是峰值信噪比对于0-1范围的数据公式是PSNR 10 * log10(1 / MSE)def compute_psnr(img1, img2): 计算两张0-1范围内图像的PSNR img1, img2: PyTorch Tensor, 形状为 (B, 3, H, W) mse torch.mean((img1 - img2) ** 2) if mse.item() 1e-10: return 100.0 # 完全一致时给出一个很大的值 psnr 10.0 * torch.log10(1.0 / mse) return psnr.item()SSIM稍微复杂一点我不想贴一个过于臃肿的实现这里给出一个简化版本它用高斯窗口计算局部均值和方差再合成结构相似度def gaussian_kernel(size11, sigma1.5): 生成一维高斯核并外积成二维核用于SSIM计算 coords torch.arange(size, dtypetorch.float32) - size // 2 g torch.exp(-(coords ** 2) / (2 * sigma ** 2)) g / g.sum() return g.outer(g).view(1, 1, size, size) def compute_ssim(img1, img2, window_size11, sigma1.5): 简化版SSIM忽略通道合并逐通道计算后取平均 img1, img2: 0-1范围的Tensor形状(B, 3, H, W) device img1.device window gaussian_kernel(window_size, sigma).to(device) # 常数C1, C2防止除零 C1 0.01 ** 2 C2 0.03 ** 2 # 计算均值 mu1 torch.nn.functional.conv2d(img1, window, paddingwindow_size // 2) mu2 torch.nn.functional.conv2d(img2, window, paddingwindow_size // 2) mu1_sq, mu2_sq mu1 ** 2, mu2 ** 2 mu1_mu2 mu1 * mu2 # 计算方差和协方差 sigma1_sq torch.nn.functional.conv2d(img1 * img1, window, paddingwindow_size // 2) - mu1_sq sigma2_sq torch.nn.functional.conv2d(img2 * img2, window, paddingwindow_size // 2) - mu2_sq sigma12 torch.nn.functional.conv2d(img1 * img2, window, paddingwindow_size // 2) - mu1_mu2 ssim_map ((2 * mu1_mu2 C1) * (2 * sigma12 C2)) / \ ((mu1_sq mu2_sq C1) * (sigma1_sq sigma2_sq C2)) return ssim_map.mean().item()5.2 结果可视化把雨图、去雨图、干净图画一起比只看数值更靠谱评估一个去雨模型不能只盯PSNR。我自己吃过这个亏模型PSNR到了32dB但肉眼一看背景细节已经糊成水彩画了。PSNR只衡量像素差异的全局平均值它不能反映局部结构是否被抹平。所以我每次验证都做一张三连图左边是雨图输入中间是模型输出右边是干净原图。代码很简单def save_comparison(rainy, output, clean, save_path): 将输入雨图、模型输出、干净原图横向拼接保存 rainy/output/clean: Tensor (3, H, W)0-1范围 # 转成numpy并还原到0-255范围 rainy_np rainy.cpu().numpy().transpose(1, 2, 0) * 255.0 output_np output.cpu().numpy().transpose(1, 2, 0) * 255.0 clean_np clean.cpu().numpy().transpose(1, 2, 0) * 255.0 combo np.hstack([rainy_np, output_np, clean_np]).astype(np.uint8) # 保存为RGB图注意要转回BGR给cv2 combo_bgr cv2.cvtColor(combo, cv2.COLOR_RGB2BGR) cv2.imwrite(save_path, combo_bgr)然后每个epoch结束时挑几张验证图跑一次保存我就能直观看到模型有没有出现过度平滑的现象。如果输出图像像被水洗过一样那说明要调整损失函数比如加大SSIM的权重。6. 推理部署与常见问题排查从训练到真正能用的一段路6.1 单张图片推理全流程注意归一化必须和训练一致模型训练好了最终要拿去处理任意一张图。推理流程和训练时有几个细节不同我在这里完整写一遍同时也解释为什么有些细节容易出错。def inference_single_image(model, image_path, output_path, use_gpuTrue): 对单张雨图做去雨推理 device torch.device(cuda if use_gpu and torch.cuda.is_available() else cpu) model model.to(device) model.eval() # 1. 读取图像并转为RGB img cv2.imread(image_path) if img is None: raise ValueError(f无法读取图片: {image_path}) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 2. 保存原图尺寸推理后要把输出resize回原尺寸 h, w, _ img.shape img_resized cv2.resize(img_rgb, (224, 224), interpolationcv2.INTER_AREA) # 3. 转Tensor并归一化到0-1范围增加batch维度 input_tensor torch.from_numpy(img_resized.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0 input_tensor input_tensor.to(device) # 4. 推理并禁止计算梯度 with torch.no_grad(): output_tensor model(input_tensor) # 5. 处理后保存先转回numpyclamp到0-1区间再resize回原始尺寸 output_np output_tensor.squeeze(0).cpu().numpy().transpose(1, 2, 0) output_np np.clip(output_np, 0.0, 1.0) * 255.0 output_np output_np.astype(np.uint8) output_rgb cv2.resize(output_np, (w, h), interpolationcv2.INTER_LINEAR) # 6. RGB转BGR存盘 output_bgr cv2.cvtColor(output_rgb, cv2.COLOR_RGB2BGR) cv2.imwrite(output_path, output_bgr)这里有个容易踩的坑如果训练时用了归一化中常见的mean/std操作推理时也必须用完全相同的mean/std。我这个代码训练时只除以了255没有做mean/std归一化所以推理时也保持一致的归一化方式。如果两边不一致模型输出的颜色会偏移而且很难发现原因。6.2 推理阶段常踩的几个坑输入尺寸不一致的问题。我在上面用了resize到224x224再推理的方法简单但会丢失细节尤其在大图上效果不好。更好的做法是滑窗推理把大图切成若干224x224的patch分别推理最后拼回原尺寸。代价是推理时间变长但对细节保留非常明显。另一个坑是输出图像的暗部偏色。如果训练时合成雨纹的alpha值总是偏大模型会在真实无雨图片上产生去雨过猛的问题把暗部细节也当成雨纹去掉。这就需要在训练数据里混入一些不含雨纹的干净图片让模型学会没有雨纹时保持原样。6.3 从基线到更好效果的三条改进路径这套基线模型的作用是跑通全流程如果要做更深入的课题我有三条建议第一把普通卷积块换成残差块或者密集连接块训练稳定性会有明显提升。第二在卷积之后加一个简单的通道注意力模块让网络学会重点处理雨纹出现概率高的通道。第三把MSE损失换成Charbonnier损失配合SSIM损失对离群点的抗干扰性更强。不过我必须强调一点在追求更先进结构之前先把数据合成质量和评测流程做扎实。很多人在小改动上纠结半天结果发现是合成雨纹的程序写错了这亏我吃过。写在最后关于这套代码的一些个人心得整套流程跑下来我最想分享的体会是图像去雨的提升瓶颈往往不在网络结构而在看不见的细节里。数据合成时雨纹密度和长度的分布决定了模型的上限训练时损失函数的结构约束决定了输出是不是耐看评测时如果不看主观效果很容易被PSNR欺骗。我建议每一个做这个方向的同学都从合成数据开始亲手搭一遍这整套代码训练过程中每10个epoch保存一次验证输出图盯着图像看趋势比只看曲线靠谱得多。跑通之后你会对残差学习、感受野、损失函数设计这些概念有完全不同的理解。最后分享一个小技巧训练接近收敛时把验证集中效果最差的几张图挑出来看。分析这些图面的共同点往往能准确指出当前模型的短板——是长雨纹处理不了还是密集雨纹区域的背景纹理被抹掉。针对短板去调整数据或网络比盲目堆参数高效得多。
