UNetRes结构的条件残差扩散图像复原

项目介绍

系统实现了以下主要功能:首先基于上传的模糊图像,通过Web界面上传到服务器端;然后利用U-Net残差扩散模型进行图像重建与去模糊处理,恢复图像的清晰细节;并支持多种退化类型的图像修复,包括通用去模糊、低光照图像增强和天气条件下的图像修复。

用户界面采用现代化的Web界面设计,提供左右分栏布局的直观对比界面,支持图片拖拽上传,包含实时处理进度显示和滑动对比查看功能。处理完成后不仅显示清晰的修复结果,还提供专业的可视化对比功能:原图与修复结果对比图、滑动对比分析图、批量处理进度监控以及模型状态实时显示,帮助用户直观地理解处理效果和模型性能。整个系统旨在为用户提供一个高效、易用的图像去模糊工具,在实际应用场景中就能快速处理各种模糊图像,提升图像质量。

U-Net残差扩散模型是一种深度融合了U-Net架构、残差学习与去噪扩散概率模型的生成式模型。该模型以U-Net作为核心骨干网络,利用其对称的编码器-解码器结构进行多尺度特征提取与空间信息恢复,并通过跳跃连接保留浅层细节以提升生成精度。在U-Net的构建模块中引入残差学习机制,通过拟合恒等映射与目标输出之间的残差,有效缓解深层网络训练中的梯度弥散问题,增强模型训练的稳定性与收敛效率。整个模型遵循去噪扩散概率模型的理论框架,前向过程通过马尔可夫链逐步向数据添加高斯噪声直至完全破坏,反向过程则由上述残差U-Net作为去噪器,根据当前时刻的噪声图像与时间步信息精确预测并移除噪声,最终从纯噪声中逐步还原出清晰的图像。因此,该模型本质上是一个以残差U-Net参数化去噪过程、通过迭代逆扩散实现高质量图像生成与处理的概率生成模型。

通用模型、低光照模型、天气模型。这三个模型都是基于 U-Net 残差扩散模型的框架构建的,只是在训练数据和任务目标上有所区别。通用模型是项目的主干模型,使用多种退化类型(如模糊、噪声、降采样等)的图像对进行训练,目标是在没有特定场景限制的情况下处理各种类型的图像退化,可以看作是一个综合性的图像复原模型。低光照模型则专门针对暗光或夜景场景进行训练,训练数据主要是配对的暗光图像和正常光照图像,模型学习在这种特定光照条件下的图像增强和去噪能力。天气模型则专注于处理恶劣天气条件下的图像退化问题,训练数据包含雾天、雨天等天气条件下拍摄的图像对,模型重点学习去雾和图像清晰化能力。从网络架构角度看,这三个模型都采用相同的 UNet 残差扩散框架(金字塔注意力 + 全局注意力 + 时间嵌入机制),区别仅在于训练好的模型权重参数不同,因此可以理解为同一个模型架构在不同数据集上微调后的专门化版本,前端界面通过模型类型选择来调用不同的权重文件实现场景化的图像处理。

注释:文中提到的UNetRes模型即残差U-Net的一种变体,它在原有U-Net残差扩散模型的去噪器基础上集成了金字塔层与注意力机制,通过增强多尺度特征提取与全局依赖建模能力,进一步改善了模型的生成效果。

项目预览

系统处理逻辑图

项目数据集

项目的训练数据集是一个多退化类型的图像对(LQ-GT pairs)数据集,用于图像复原任务。训练集(my_train_meta.txt)包含100张图像,每张图像都有对应的GT(Ground Truth,高质量清晰图像)和LQ(Low Quality,低质量退化图像),文件命名为0001.jpg到0100.jpg。测试集(my_test_meta.txt)更为复杂,包含200张测试图像,涵盖四种不同的退化类型:模糊(blur,编号0001-0050)、噪声(noise,编号0051-0100)、雾霾(haze,编号0101-0150)和雨滴(rain,编号0151-0200),每种退化类型有50张图像。这意味着模型需要处理多种实际场景中的图像退化问题,是一个多任务图像复原模型

U-Net残差扩散模型介绍

整体架构采用编码器-解码器(UNet)结构结合扩散模型的思想,其核心创新在于不直接预测清晰图像,而是预测低质量图像到高质量图像之间的退化残差,这种设计使得模型能够更好地学习图像退化的本质特征。模型的UNet backbone包含三个核心组件。首先是时间嵌入模块,使用正弦位置编码(SinusoidalPosEmb)将扩散过程中的时间步转换为高维向量,然后通过一个包含GELU激活的两层MLP进行处理,最终将时间信息注入到每个ResNet块中,使模型能够根据不同的扩散阶段自适应调整预测策略。其次是ResNet块结构,每个块包含两个卷积层和时间条件调制机制,时间嵌入通过MLP转换为scale和shift参数,对特征进行仿射变换,使网络能够感知当前的扩散进度。最后是两种注意力机制,金字塔注意力(PyramidAttention)在三个不同尺度(1×、2×、4×)上提取特征并使用线性注意力进行融合,这种设计使模型能够同时捕获局部细节和全局结构信息;全局注意力(GlobalAttention)通过全局池化操作获取全图上下文信息,特别适合处理雾霾、模糊等大面积退化问题,被放置在网络的瓶颈层最底部。扩散模型部分采用了条件扩散的设计,输入为低质量图像(LQ)和加噪后的混合图像。训练时,模型同时预测残差和噪声两个目标,其中残差是LQ与GT之间的差异,噪声是前向扩散过程中添加的高斯噪声。损失函数使用L1 Loss来优化预测的准确性。在推理采样阶段,模型使用DDIM加速采样算法,从带噪声的低质量图像开始,通过逐步去噪的方式重建高质量图像,最终输出为LQ图像减去预测的残差得到清晰结果。整个网络的处理流程是:输入的低质量图像首先经过初始卷积层进行特征提取,然后通过编码器进行三次下采样,每次下采样后接两个ResNet块和一个金字塔注意力模块来提取多尺度特征;到达瓶颈层后依次经过ResNet块、自注意力、全局注意力和另一个ResNet块;之后进入解码器部分,通过三次上采样逐步恢复空间分辨率,每次上采样后同样包含ResNet块和金字塔注意力,并与编码器对应层的特征通过跳跃连接进行融合;最后通过卷积层输出预测的残差图像。这种编码器-解码器配合跳跃连接的结构能够有效保留浅层的细节信息,同时通过深层的全局上下文信息进行补充修正。

为U-net模型增加金字塔和全局注意力机制,改进模型,让模型效果更好。

在项目中,模型的构建代码在model.py当中。

模型训练参数

配置项 设置值
模型架构 基于UNet的残差扩散模型 (UNetRes)
注意力机制 金字塔注意力 (PyramidAttention) 与全局注意力 (GlobalAttention)
图像尺寸 512 × 512
批次大小 4
训练步数 100 (用于单退化训练)
扩散时间步数 1000
采样步数 10 (可配置)
损失函数 L1 Loss
优化器 Adam (betas=(0.9, 0.99))
EMA 衰减率 0.995
学习率 1 × 10⁻⁴

项目结构图

功能展示

前端页面功能

前端采用现代化的深色主题设计,左侧为图片上传区域,右侧为处理结果展示区域,中间通过模型选择器切换不同的处理模型。

调用通用模型进行检测

调用模糊的图片进行处理,调用通用模型进行检测,左侧为原始图像,右侧为处理后的图像,如图所示。

调用低光照模型进行处理

首先选取低光照条件下的图片进行处理,前端页面调用低光照模型进行处理,左侧为原始图像右侧为处理后的图像。可以看到原始图像非常的灰暗,处理后的图像亮度清晰,能见度变高。

调用天气模型进行处理

首先选取天气环境复杂的图片进行测试,切换天气模型进行图像处理,左侧为原始图像,可以看到太阳光非常刺眼后侧主体眩光看不清,但是模型处理以后,中间主体变得清晰,整体变得通透,视觉效果变好。

支持对比查看结果

用户可以在右侧处理结果点击对比查看按钮,用户可以实时拖拽查看处理后的图像有何区别。如图所示。

批量处理功能

用户可以点击web页面右上角批量处理按钮,可以实时选择多张图片进行处理,提高处理的效率。

训练的结果指标图

训练过程中使用的损失函数是L1 Loss,直接计算模型预测的残差与真实残差之间的平均绝对误差,这种损失函数能够促使模型学习到准确的残差映射。在模型训练完成后进行效果评估时,会使用多个更全面的指标来衡量图像复原的质量。PSNR(峰值信噪比)是最基础的图像质量评估指标,它计算的是图像最大可能像素值与均方误差之间的比值,并以对数形式表示,单位为dB。PSNR值越高代表图像失真越少、质量越好,通常PSNR在30dB以上表示图像质量较好,计算公式为PSNR = 20 × log10(MAX / sqrt(MSE)),其中MAX是像素的最大值(255),MSE是两幅图像的均方误差。SSIM(结构相似性)是一种更符合人眼感知的图像质量评估指标,它从亮度、对比度和结构三个维度综合评估两幅图像的相似程度。SSIM值范围为0到1,值越接近1表示两幅图像越相似。该模型采用三维高斯核方法计算SSIM,能够更准确地捕捉图像的局部结构信息。LPIPS( Learned Perceptual Image Patch Similarity)是一种基于深度学习的感知相似性指标,它使用预训练的神经网络提取图像特征,然后计算特征之间的距离。LPIPS值越低表示图像在感知上越接近真实图像,这个指标比传统的像素级指标更能反映人眼对图像质量的感知。最终评分公式将这三个指标综合起来形成综合评分,具体计算方式为 Final Score = PSNR + 10 × SSIM - 5 × LPIPS。这个公式的设计思路是PSNR和SSIM越高越好,而LPIPS越低越好,所以PSNR和SSIM为正加权,LPIPS为负加权。在实际评估时,模型会对测试集中的每张图像计算这三个指标,然后取平均值作为最终评估结果。

Finalscore=PSNR+10×SSIM−5×LPIPS

图像复原效果

模糊图像处理效果

复原前:建筑立面模糊,窗框边缘不清晰

去雾与去雨效果对比

复原前:前景存在明显雨丝干扰

复原后:雨丝大幅抑制,边缘干净

复原前:全局灰白,能见度低

复原后:色彩饱和,景深层次分明

总结

扩散模型的核心思想源自 points.非平衡热力学,通过在前向扩散过程中逐步向图像添加高斯噪声将其破坏成纯噪声,然后在反向扩散过程中学习去噪来恢复清晰图像。传统的扩散模型直接预测清晰图像或噪声,而这个模型创新性地提出了残差预测的方式,即不直接重建高质量图像,而是预测低质量图像到高质量图像之间的残差(退化分量),这种方式使得模型能够更精准地聚焦于图像退化部分的学习。在训练阶段,模型同时接收低质量图像作为条件输入和经过前向扩散处理的含噪GT图像,通过UNet编码器提取多尺度特征并结合金字塔注意力机制捕获不同分辨率下的图像细节,结合全局注意力获取全图上下文信息,最终输出对残差和噪声的双重预测。在推理阶段,采用DDIM(Denoising Diffusion Implicit Models)加速采样算法,通过跳过部分扩散步骤来加快去噪速度,同时利用低质量图像作为条件引导,使生成的图像更符合输入图像的结构特征。整体框架将扩散模型的强大生成能力与图像复原领域的残差学习思想巧妙融合,既保留了扩散模型的全局建模优势,又通过残差预测显著降低了学习难度,提高了图像复原的精度和稳定性。