简介本资源为基于波形扩散模型的低光照图像增强算法Python实现源码包面向计算机、人工智能、电子信息等相关专业的在校学生、教师及企业研发人员可用于课程设计、毕业设计、大作业或初期项目立项演示。算法借助波形扩散模型有效提升低光照图像质量并配套预训练模型与多个公开数据集涵盖训练、评估脚本及配置流程便于快速复现与二次开发。压缩包共34个文件以15个py源码文件为核心辅以14个pyc编译文件、2张png效果图、1个yml配置、1个md说明及1个txt依赖清单整体约21.1MB目录结构清晰模型、数据集、工具模块划分明确。目前已有491人学习关注。读者可从中获得完整的算法实现思路、扩散模型与波形变换结合的关键代码、数据增强与采样优化细节以及可运行的训练评估流程适合入门进阶学习或在此基础上修改扩展新功能。1. 波形扩散模型做低光照增强为什么它比直接提亮更值得折腾夜里十一点在停车场拍的那张照片你把它丢进任何一款修图软件里拉高亮度得到的往往是一层灰蒙蒙的雾暗部噪点被一起放大颜色发白发绿人脸像糊了一层蜡。低光照图像增强要解决的从来不是把像素值乘个系数这么简单它本质是一个病态逆问题传感器在极暗环境下采集到的信号信噪比极低光子噪声、读出噪声、量化误差混在一起亮度和色彩信息已经被破坏你看到的暗只是表象真正丢掉的是结构和色度。波形扩散模型Wavelet Diffusion Model这条路子是把扩散模型的生成能力和小波变换的多尺度分解结合起来。普通扩散模型直接在像素空间做去噪计算量大、细节容易糊换到小波域之后低频子带承载整体亮度和色调高频子带承载边缘和纹理模型可以针对不同频带用不同的去噪强度暗部细节恢复得更干净。这份 python 实现源码.zip 对应的就是这套思路的工程落地训练一个条件扩散模型输入是低光图的小波系数输出是正常光照下的小波系数再逆变换回像素空间。适合谁看已经会写 python、装过 pytorch、跑过至少一个图像复原或生成模型的人想从 Retinex、直方图均衡、GAN 那套老方法里跳出来试试扩散模型在 low-light enhancement 上到底能不能打的人以及手里有配对数据低光/正常光想快速复现一版 baseline 的人。如果你连 python 环境都没配过建议先把 python 安装、python 安装 numpy 库的方法、vscode python 环境配置这几件事做完再回来否则后面每一步都会卡在环境上。这篇不吹模型多强只讲清楚三件事波形扩散模型为什么适合低光照、源码里各个模块怎么串起来、以及我在复现时踩过的那些坑。读完你应该能自己把训练跑起来并且知道 loss 不降的时候该去动哪个参数。2. 波形扩散模型的核心机制与源码模块拆解2.1 为什么把扩散过程搬到小波域先讲清楚动机不然你读源码会一头雾水。标准 DDPM 的前向过程是对整张图逐步加高斯噪声反向过程学一个网络去预测噪声。放到低光照增强里问题在于暗部区域的信号本身就很弱你在像素空间加噪等于把本来就不多的有效信息彻底淹没网络要同时学会提亮和去噪两件事训练极不稳定。小波变换这里一般用 Haar 或 db2把图像拆成 LL、LH、HL、HH 四个子带。LL 是低频近似基本就是图像的亮度骨架和色调分布LH/HL/HH 是水平、垂直、对角方向的高频细节。低光图和正常光图的差异在 LL 子带上表现为整体亮度偏移在高频子带上表现为细节丢失和噪声放大。分开处理的好处是LL 子带用较强的条件约束去对齐亮度高频子带用较弱的噪声调度保留纹理网络不用在一个空间里同时兼顾两种量级差异巨大的信号。源码里对应的模块通常叫DWT和IWT离散小波变换与逆变换一般用 pytorch 手写卷积实现不依赖 pywt因为要放进计算图里端到端训练。前向扩散在四个子带上分别进行时间步 t 是共享的但每个子带的噪声方差可以按频带缩放这是这套方法区别于普通 DDPM 的关键设计。2.2 源码目录结构与数据流拿到源码包之后先别急着python train.py。我一般会先花十分钟把目录结构和数据流捋一遍不然报错的时候你连改哪个文件都不知道。典型的结构大致是这样lowlight_wavelet_diffusion/ ├── configs/ │ └── default.yaml # 训练超参、数据路径、频带权重 ├── data/ │ └── dataset.py # 配对数据加载返回 low/normal 两张图 ├── models/ │ ├── dwt.py # 小波变换与逆变换 │ ├── unet.py # 去噪主干网络 │ └── diffusion.py # 前向加噪、反向采样调度 ├── utils/ │ ├── metrics.py # PSNR / SSIM / LPIPS │ └── visualize.py # 训练过程可视化 ├── train.py └── infer.py数据流是这样的dataset.py读入一对图低光、正常光归一化到 [-1, 1]dwt.py把两张图都做小波分解得到各自的四个子带diffusion.py对正常光子带做前向加噪得到 t 时刻的噪声子带unet.py接收低光子带 噪声子带 时间步嵌入预测噪声loss 是预测噪声和真实噪声的 MSE可以按频带加权。推理时从纯噪声出发用低光子带作为条件逐步去噪最后IWT还原成图像。2.3 去噪主干网络的关键参数主干一般用 U-Net 变体源码里能调的参数集中在configs/default.yaml。下面这张表是我复现时实际会去动的几个默认值只是起点不是最优参数含义常见取值调整建议base_channelsU-Net 第一层通道数64显存不够降到 32细节会掉channel_mult各层通道倍率[1,2,4,8]别超过 4 层暗部容易过平滑num_res_blocks每个分辨率残差块数2加到 3 提升细节训练变慢timesteps扩散总步数1000推理想快可降到 250需重训band_weight高频子带 loss 权重[1,1,2,2]高频给高权重纹理更锐lr学习率1e-4发散就减半收敛慢可加 warmupband_weight这个参数是这套方法的核心调优点。LL 子带权重给太高结果就是整体亮度对了但细节糊高频权重给太高噪点会被当成纹理保留下来。我一般从 [1,1,2,2] 起步看验证集 SSIM 再微调。2.4 最小可跑通的训练命令环境准备好之后最小训练流程是这样# 安装依赖注意 torch 版本要和 CUDA 匹配 pip install torch torchvision numpy opencv-python pyyaml tqdm # 单卡训练先用小数据集验证流程能跑通 python train.py --config configs/default.yaml \ --data_root ./data/LOLdataset \ --batch_size 4 \ --epochs 100 \ --gpu 0# train.py 里核心训练循环的简化逻辑 for epoch in range(cfg.epochs): for low, normal in dataloader: low low.cuda(); normal normal.cuda() # 小波分解得到四个子带 low_bands dwt(low) normal_bands dwt(normal) # 随机采样时间步对正常光子带加噪 t torch.randint(0, cfg.timesteps, (low.size(0),)).cuda() noise torch.randn_like(normal_bands) noisy_bands q_sample(normal_bands, t, noise) # 网络以低光子带为条件预测噪声 pred unet(noisy_bands, low_bands, t) # 按频带加权计算 loss loss weighted_mse(pred, noise, cfg.band_weight) optimizer.zero_grad(); loss.backward(); optimizer.step()逻辑说明dwt把两张图都拆成子带注意低光图只做分解不做加噪它是条件q_sample是前向扩散的闭式解直接按累积噪声系数采样不用循环weighted_mse按频带给不同权重这是控制细节和亮度平衡的旋钮。参数上batch_size先给 4 是为了确认显存够用跑通后再往上加epochs给 100 只是看 loss 趋势真正收敛通常要 300 以上。提示第一次跑务必把batch_size设小、epochs设小先确认 loss 能稳定下降、显存不爆再放大规模。直接上大配置报错信息会把你淹没。3. 数据准备与训练配置从配对数据集到可复现的 baseline3.1 配对数据的组织方式低光照增强是有监督任务需要低光/正常光配对图。常见数据集是 LOL、SID、MIT-Adobe FiveK 这几类源码默认按文件夹配对读取。目录结构一般长这样LOLdataset/ ├── train/ │ ├── low/ # 低光图 │ └── high/ # 对应正常光图文件名必须一一对应 └── test/ ├── low/ └── high/dataset.py里读取逻辑的关键是文件名对齐。我见过太多人卡在这里low 和 high 文件夹里图片数量一样但名字对不上训练时随机配对loss 看着在降实际学的是乱配的映射推理出来颜色全错。稳妥做法是在__init__里做一次断言import os class PairedDataset(Dataset): def __init__(self, root, splittrain): self.low_dir os.path.join(root, split, low) self.high_dir os.path.join(root, split, high) low_names sorted(os.listdir(self.low_dir)) high_names sorted(os.listdir(self.high_dir)) # 文件名必须严格一一对应否则直接报错别让它悄悄跑 assert low_names high_names, low/high 文件名不匹配检查数据集 self.names low_names def __getitem__(self, idx): name self.names[idx] low cv2.imread(os.path.join(self.low_dir, name)) high cv2.imread(os.path.join(self.high_dir, name)) # 统一尺寸扩散模型对尺寸敏感最好裁成 2 的幂 low cv2.resize(low, (256, 256)) high cv2.resize(high, (256, 256)) return to_tensor(low), to_tensor(high)逻辑说明assert那行是后悔药宁可启动就崩也别让模型在错配数据上白跑一天。cv2.resize到 256 是因为小波分解要做多层尺寸最好是 2 的整数次幂否则边界处理会引入伪影。参数上如果你显存够可以裁到 512细节恢复明显更好但 batch_size 要相应减半。3.2 训练配置怎么改才不翻车configs/default.yaml是唯一入口改配置比改代码安全。我复现时的习惯是先锁定三件事学习率、频带权重、时间步。学习率 1e-4 是扩散模型的常见起点配合 cosine 调度频带权重按前面说的从 [1,1,2,2] 起步时间步 1000 是标准值别乱动除非你清楚重训的代价。# configs/default.yaml 关键片段 model: base_channels: 64 channel_mult: [1, 2, 4, 8] num_res_blocks: 2 diffusion: timesteps: 1000 beta_schedule: cosine # 比 linear 更稳暗部收敛快 loss: band_weight: [1.0, 1.0, 2.0, 2.0] train: lr: 1.0e-4 batch_size: 4 epochs: 300 warmup_steps: 1000beta_schedule选 cosine 而不是 linear是我踩坑之后改的。linear 调度在低光照任务里前期噪声加得太猛LL 子带的亮度信息很快被破坏网络学不到有效的亮度映射表现为 loss 前期震荡。cosine 调度在两端更平缓暗部对齐更稳。3.3 训练过程该盯哪些指标别只盯 loss。扩散模型的 loss 是噪声预测的 MSE它降不代表图像质量好。我一般同时看三个训练 loss、验证集 PSNR、以及每隔若干 epoch 存一张推理图肉眼比对。# utils/metrics.py 里 PSNR 和 SSIM 的调用 from skimage.metrics import peak_signal_noise_ratio as psnr from skimage.metrics import structural_similarity as ssim def evaluate(pred, gt): # pred/gt 都是 [0,1] 范围的 numpy 数组 p psnr(gt, pred, data_range1.0) s ssim(gt, pred, data_range1.0, channel_axis2) return p, s逻辑说明PSNR 反映整体亮度对齐程度SSIM 反映结构保真度。低光照增强里经常出现 PSNR 高但 SSIM 低的情况说明亮度对了但结构糊了这时候要去调高频子带的权重。参数上data_range必须和你的归一化范围一致用 [0,1] 就写 1.0用 [0,255] 就写 255写错指标会离谱。注意验证集一定要和训练集来自同一分布。我见过有人拿 LOL 训练、拿手机夜景测试PSNR 掉十几个点然后怀疑模型有问题其实是域差异不是代码 bug。4. 推理、可视化与效果验证的实操细节4.1 从纯噪声采样出增强图推理流程和训练相反从标准高斯噪声出发以低光图的小波子带为条件逐步去噪最后逆变换。源码里infer.py封装了这套逻辑核心是采样循环。torch.no_grad() def sample(low_img, model, diffusion, cfg): low_bands dwt(low_img) # 条件不做加噪 # 从纯噪声开始形状和子带一致 x torch.randn_like(low_bands) for t in reversed(range(cfg.timesteps)): t_batch torch.full((x.size(0),), t, devicex.device) pred_noise model(x, low_bands, t_batch) x diffusion.p_sample(x, pred_noise, t_batch) return iwt(x) # 逆小波变换回像素空间逻辑说明dwt(low_img)只算一次整个采样过程复用别在循环里重复算浪费时间。p_sample是反向扩散的单步更新内部按 beta 调度算均值和方差。参数上timesteps必须和训练时一致训练用 1000 推理用 250 会直接崩除非你用了 DDIM 这类可跳步的采样器那需要单独配置。4.2 可视化对比怎么做才有说服力单看一张增强图容易被骗因为人眼对亮度很敏感稍微亮一点就觉得好。我一般做三栏对比低光原图、增强结果、正常光参考图再叠一张误差热力图。import matplotlib.pyplot as plt def visualize(low, pred, gt, save_path): fig, axes plt.subplots(1, 4, figsize(16, 4)) axes[0].imshow(low); axes[0].set_title(Low) axes[1].imshow(pred); axes[1].set_title(Enhanced) axes[2].imshow(gt); axes[2].set_title(Reference) # 误差热力图看增强结果和参考图的差异分布 err np.abs(pred - gt).mean(axis2) axes[3].imshow(err, cmaphot); axes[3].set_title(Error) for ax in axes: ax.axis(off) plt.savefig(save_path, dpi150, bbox_inchestight)逻辑说明误差热力图是关键它能告诉你模型在哪里翻车。如果误差集中在暗部边缘说明高频子带恢复不足如果误差均匀分布在亮区说明 LL 子带对齐有问题。参数上dpi给 150 够看bbox_inchestight去掉白边方便直接放进报告。4.3 定量指标怎么读才不被误导PSNR、SSIM、LPIPS 三个指标各有盲区。PSNR 对亮度偏移敏感但对结构不敏感SSIM 看结构但会忽略颜色LPIPS 用感知特征更接近人眼但数值不直观。我一般三个一起看再配合误差热力图。指标反映什么盲区低光照场景的参考值PSNR整体像素误差结构、颜色20-28 dB 算正常SSIM结构相似度颜色偏移0.75-0.90 算不错LPIPS感知差异数值不直观越低越好0.2 以下算好如果 PSNR 高但 LPIPS 也高说明像素对上了但感知上不自然常见于过度平滑的结果这时候要回头调高频权重。如果 SSIM 高但颜色偏检查 LL 子带的色度通道是不是被单独处理了。提示指标只在同一数据集内比较有意义。跨数据集比 PSNR 是耍流氓域差异带来的波动比模型改进还大。5. 避坑与排查复现波形扩散模型时最容易翻车的五件事5.1 现象loss 前期剧烈震荡几百步后直接 NaN原因学习率太大或者beta_schedule用了 linear 导致前期噪声过强梯度爆炸。低光照任务里 LL 子带的数值范围比高频子带大很多统一学习率容易让高频部分梯度失控。解决先把学习率降到 5e-5加 1000 步 warmupbeta_schedule换成 cosine如果还炸在 loss 里对高频子带做梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)。我一般还会打印每个子带的 loss 分量看是哪个频带先炸。5.2 现象增强结果整体偏灰暗部提亮了但颜色发白原因LL 子带只对齐了亮度没管色度。小波分解如果用的是 RGB 三通道分别做色度信息散落在各通道的 LL 里网络容易只学到亮度映射。解决把图像转到 YCbCr 或 Lab只对 Y亮度通道做扩散增强CbCr 用简单的增益补偿或者直接保留。源码里如果没做颜色空间转换自己加一个效果立竿见影。参数上色度增益别超过 1.5否则会偏色。5.3 现象推理时显存爆掉训练时却没事原因推理没有torch.no_grad()或者采样循环里每一步都保留了计算图。扩散模型采样要跑几百步每步都建图显存直接翻几百倍。解决采样函数必须加torch.no_grad()如果还要更省用torch.inference_mode()batch 设成 1逐张推理。另外检查dwt和iwt里有没有不必要的中间变量缓存。5.4 现象高频细节糊成一片边缘像被抹过原因高频子带 loss 权重太低或者 U-Net 的channel_mult层数太多深层感受野过大把高频当噪声平滑掉了。解决把band_weight里高频两项从 2 提到 3 甚至 4channel_mult砍掉最后一层从 [1,2,4,8] 改成 [1,2,4]如果还糊检查num_res_blocks是不是太小加到 3。我踩过这个坑调完高频权重SSIM 直接涨了 0.04。5.5 现象训练集指标很好测试集一塌糊涂原因过拟合或者训练测试数据分布不一致。低光照数据集通常很小LOL 才几百对模型参数量一大就记住训练集了。解决加数据增强随机裁剪、水平翻转、亮度微扰加 dropout 或者 weight decay如果测试集是不同设备拍的考虑做域适应或者至少做一次颜色归一化。别指望在小数据集上堆参数能泛化这是血泪经验。6. 进阶技巧用频带自适应权重把暗部细节再抠出来前面讲的band_weight是固定值训练全程不变。但低光照图像有个特点不同区域的暗部程度差异很大一张图里可能既有接近全黑的角落也有相对亮的区域。固定权重要么对极暗区不够要么对较亮区过度增强。我后来改成按 LL 子带的局部均值动态算权重效果比固定值稳。思路是这样LL 子带反映局部亮度均值越低说明该区域越暗需要越强的高频恢复力度。给每个空间位置算一个权重图再和高频子带的 loss 逐像素相乘。def adaptive_band_weight(low_bands, base_weight): # low_bands[0] 是 LL 子带取局部均值衡量暗度 ll low_bands[:, 0:1] # [B,1,H,W] # 局部均值用平均池化近似 local_mean F.avg_pool2d(ll, kernel_size8, stride8) local_mean F.interpolate(local_mean, sizell.shape[-2:], modenearest) # 越暗权重越大加 0.5 防止全黑区域权重爆炸 darkness 1.0 - local_mean weight_map base_weight * (1.0 darkness) return weight_map逻辑说明avg_pool2d加interpolate是为了得到平滑的局部亮度估计避免逐像素算导致权重图噪声太大。darkness越大表示越暗权重相应放大。参数上kernel_size给 8 是经验值太小权重图会抖太大就失去局部性base_weight还是从 [1,1,2,2] 起步自适应部分只做微调。这套改法我在两个数据集上试过PSNR 提升不明显但 LPIPS 和肉眼观感改善明显尤其是夜景里暗部的人脸和招牌文字边缘干净很多。代价是训练时多算一次池化和插值显存和时间增加大概 5%可以接受。验证这套改动有没有用别只看指标。我的习惯是固定几张典型图——一张极暗、一张有强光源、一张色彩丰富——每次改完都跑这三张拼成对比图存下来。时间长了你会有一组自己的回归测试图比任何指标都直观。这个习惯帮我省了无数次指标涨了但图更丑的返工。最后说个心态上的事扩散模型调参是个耐心活一次只动一个变量动完记录指标和观感别同时改学习率、权重、网络结构不然出了问题你根本不知道是谁的锅。我早期就是贪快一次改三处结果 loss 崩了排查了一整晚。希望帮到你。本文还有配套的精品资源点击获取