IRN论文阅读笔记
《Invertible Rescaling Network and Its Extensions》
可逆重缩放网络及其扩展
Zotero 插件商店:https://zotero-chinese.com/plugins/
补充概念
上采样与降采样。
上采样是指在图像或信号处理中,将图像的分辨率从**低分辨率(LR)提高到高分辨率(HR)**的过程。它的核心目标是增加像素数量,同时尽可能保留或恢复图像细节。
**下采样(降采样)是指将图像或信号从高分辨率(HR)降低到低分辨率(LR)**的过程。核心目标是减少像素数量,同时尽量保留重要的图像信息。
简单来说:
- 下采样用于减小分辨率或采样率,典型用途包括数据压缩、多尺度分析等。实现时应注意混叠问题,常需配合滤波。
- 上采样用于扩大分辨率或采样率,典型用途包括超分辨率、编码器-解码器网络等场景。不同插值方法在质量和速度上各有取舍。
下采样: 在退化阶段进行,用于模拟图像退化过程(如降分辨率、模糊化)。
上采样: 在恢复阶段进行,用于从退化图像中重建高分辨率图像或恢复细节信息。
降采样和上采样是本文方法中的核心环节,分别对应退化和恢复的两个方向,并在损失函数的引导下相互协作,最终实现高质量的图像重建。
Haar小波变换是一种可逆的小波变换方法,能够将图像分解为低频分量(代表主要内容)和高频分量(代表细节信息)。
一篇博客的见解:
这篇论文主要是模拟图像在传输过程中,降尺度和升尺度的问题:高分辨率数字图像(HR图像)通常会按比例缩小(LR图像)以适应各种显示屏或者节省存储成本和带宽,同时终端设备采用后放大方法恢复原始的分辨率或放大图像中的细节。形成一个HR →LR → HR 可逆过程,因此这里,高分辨率图像HR一开始就是可以获得的。对于超分辨率任务,真实情况下只有一个LR图像的输入,HR图像是不可获得的。
但作者给出的超分结果,是将HR作为模型的输入,通过正向网络生成LR图像,然后将其输入到反向网络重建出最终的高分辨率结果(SR图像),这里的HR是必要的。当本人模拟真实情况下仅有LR输入时,测试的结果性能非常低。
作者给出的解释:
**图像缩放(image rescaling)和超分辨率(super-resolution)**是不同的任务。IRN是下采样HR得到LR,然后从LR重建出HR;SR的目标是放大任意的LR图像。
如果仅仅使用IRN的架构对由bicubic下采样得到LR图像和HR图像进行配对训练(潜变量z作为0padding),这也是许多SR方法的设置,其性能不如它们。原因包括IRN不是主要为先验学习设计的,而且参数很少。IRN的提升来自于对降尺度和升尺度的可逆建模。
摘要
图像重缩放是一种常用的双向操作,首先将高分辨率图像降采样以适配各种显示屏或减少存储和带宽需求,其后再将对应的低分辨率图像上采样以恢复原始分辨率或细节。然而,非单射的降采样映射会丢弃高频内容,从而导致逆向恢复任务中的不适定问题。这可以抽象为一个由于信息丢失而引起的通用图像退化-恢复问题。在本研究中,我们提出了一种新颖的可逆框架以解决这一问题,该框架从新的视角,即可逆双射变换,建模双向的退化和恢复操作。可逆性使得该框架能够通过分布的形式建模预退化过程中的信息丢失,从而缓解后续恢复过程中的不适定性。具体而言,我们开发了可逆模型,在生成有效退化图像的同时,将丢失内容的分布转化为潜在变量的固定分布。在恢复阶段,通过对生成的退化图像和随机生成的潜在变量应用逆变换,使得恢复变得可行。我们从图像重缩放任务出发,将模型具体化为可逆重缩放网络,并将其扩展到类似的去色–上色任务中。此外,我们还提出将可逆框架与现有退化方法(如图像压缩)相结合,以实现更广泛的应用。实验结果表明,在从降采样和去色图像的上采样及上色重建、以及图像压缩的率失真性能方面,我们的模型在定量和定性评估中都显著优于现有方法。
引言
在互联网高分辨率(HR)图像/视频迅速增长的时代,图像重缩放变得越来越重要。为了高效存储、传输和共享,这些大尺寸数据通常被降采样,以显著减小体积,从而更加适合带宽,同时维持可视化的内容,用于预览或适配不同分辨率的屏幕。另一方面,由用户需求驱动的逆向恢复任务旨在将降采样的低分辨率(LR)图像放大为更高分辨率或恢复到原始尺寸,以便呈现生动的细节。然而,由于降采样的非单射特性,根据奈奎斯特-香农采样定理,高频内容在降采样过程中丢失。这种信息丢失导致了一个难以处理的逆向任务的不适定问题,因为同一个降采样的LR图像可能对应于多个可能的HR图像,因此给恢复过程带来了巨大挑战。
这个问题可以被抽象为一个由于维度缩减导致信息丢失的通用图像退化–恢复问题。类似的例子还包括图像去色–上色和图像压缩问题。在以下部分中,我们首先聚焦于这个通用问题,然后考虑其具体实例化。
已有许多努力尝试通过机器学习算法来缓解这一不适定性问题。例如,许多研究关注单向恢复任务,例如对于图像重缩放任务,他们选择使用超分辨率(SR)方法,通过施加或学习先验(即针对该逆向任务中给定低分辨率图像所对应的所有可能的高分辨率图像的偏好)来实现低分辨率图像的上采样。然而,主流的超分辨率算法(如 Dai 等人,2019;Dong 等人,2015;Lim 等人,2017;Wang 等人,2018;Zhang 等人,2018a,2018b)使用预定义且不可调节的降采样方法(例如双三次插值)来指导上采样的学习,这忽略了这两个互为逆操作之间的兼容性。因此,仅仅应用单向恢复方法(例如 SR)无法充分利用该任务的双向特性,导致恢复效果不理想。
一些最近的工作尝试通过编码器-解码器框架统一这些双向操作,而不是将它们分为两个独立的任务。在这些图像重缩放的方法中,编码器作为一个基于学习的上采样优化降采样模块,与上采样解码器(Kim et al., 2018)或现有的超分辨率模块(Li et al., 2018; Sun and Chen, 2020)联合训练。这种编码器-解码器框架也被应用于类似的退化–恢复任务(Xia et al., 2018; Ye et al., 2020)。通过考虑任务的双向特性,这种集成训练方法可以大幅提高图像重建的质量。然而,这些研究只是通过训练目标将两种操作简单地联系起来,而没有尝试充分利用任务的互补特性或捕捉丢失内容的特征。因此,结果也无法达到预期。
在本文中,我们提出了一种新颖的可逆框架,通过可逆双射变换从根本上缓解这一不适定性问题。受这对任务的互补性质启发,在正向过程中保留丢失信息的知识(例如,在图像重缩放任务中保留的高频内容)将极大地帮助逆向恢复。然而,存储或传输所有丢失的内容以实现精确恢复是不可行的。
为了很好地解决这一挑战,我们以分布的形式处理这些内容,并假设合理的丢失内容遵循某种分布。我们开发了一种新颖的可逆模型,以分布变换函数的形式捕捉分布的知识。
具体来说,在正向过程中,我们的可逆模型通过可逆变换将原始图像 x转化为一个退化图像 y和一个辅助潜在变量 z。y属于一个有效退化图像的目标集合,例如,在图像重缩放任务中,给定高分辨率图像 x 的视觉上令人满意的低分辨率图像集合;而 z 是一个遵循固定预设分布 p(z)(例如各向同性高斯分布)的随机变量。
y 和 z 的联合分布通过双射变换从 x 的分布中得到,因此从统计建模的角度来看,随机变量 z包含了 y所丢失的“信息”。
学习这一双射变换使我们的模型能够捕捉丢失内容的知识。然后,在逆向恢复过程中,通过从预设分布中随机采样的 z 和退化图像 y 一起,通过模型的逆函数可以恢复原始图像的大部分内容。
我们考虑了这一双向问题的两个具体实例,即图像重缩放和图像去色–上色。在具体架构方面,我们从图像重缩放任务出发,开发了可逆重缩放网络(IRN),它可以轻松扩展并适配于去色–上色任务。
为了实现这一可逆框架,在训练过程中需要克服一些挑战。我们的基本目标包括以高质量重建原始图像,并生成属于目标集合的退化图像,例如视觉上令人满意的低分辨率(LR)图像集合。进一步的目标是使用与图像无关的 z(即 z∼p(z) 而不是与图像相关的 z ∼ p(z|y)来完成恢复任务,因为独立分布 (p(z) 的统计建模和采样更加简单,无需处理条件 y。这是可行的,因为对于任何具有密度的随机向量(即 z∼p(z∣y),都存在一个双射 fy,使得 fy(z)∼N(0,I)(Hyvärinen 和 Pajunen,1999)。为了实现这些目标,我们结合了重建损失、指导损失和分布匹配损失,设计了一种新颖、紧凑且高效的目标函数。注意,最后一部分旨在将恢复的图像对齐到真实的原始图像流形上,同时强制 z遵循与图像无关的分布 p(z)。这一点无法通过传统生成对抗网络(GANs)或最大似然估计(MLE)方法简单实现。我们将 y 的分布表述为 x的推前经验分布,结合独立分布 p(z) 反向通过我们的可逆模型,以恢复 x的分布。因此,我们的分布匹配关注于这个恢复的分布与 x 的数据分布之间,并在实践中最小化它们之间的 JS 散度。
此外,我们表明,根据我们模型的可逆性,一旦在 x 上实现了分布匹配,那么这种匹配也会在 (y,z)空间上成立,其中 z 是与图像无关的。此外,我们提出将我们的可逆框架与现有退化方法相结合,并通过图像重缩放和图像压缩的结合对其进行了实例化。我们展示了将我们的可逆框架与从这些退化中恢复的结合效果。我们注意到,可逆框架和模型还可以有许多其他的通用应用,例如图像隐写术、视频重缩放、图像去噪等。有关将可逆框架和模型应用于各种任务的更多详细信息,请参阅自我们初版工作发表以来的相关研究。我们的贡献的结论如下:
- 据我们所知,我们是首个使用可逆双射变换建模图像退化与恢复互为逆过程的研究。精心设计的可逆性使该框架能够对信息丢失进行建模,从而缓解这一双向问题中的不适定性。
- 我们提出了一种新颖的模型设计和高效的训练目标来实现这一框架。它强制潜在变量 z 遵循一个简单的与图像无关的分布,从而能够基于该分布的样本实现高效的逆向上采样。我们针对图像重缩放任务设计了具有精心构建架构的 IRN,并证明其可以轻松适应类似的图像去色–上色任务。
- 提出的 IRN 及其具有尺度灵活性和高效性的变体在从降采样的 LR 图像重建 HR 图像方面,显著优于最先进的降采样-超分辨率方法和编码器-解码器方法。同时,与这些方法相比,IRN 的参数大幅减少,表明该模型具有轻量化和高效性的特点。
- 我们进一步提出将可逆框架与现有退化方法的恢复相结合,例如图像重缩放与压缩的结合,以实现更广泛的应用。实验表明,在这些场景中也取得了改进。
相关工作
降采样后的图像上采样:
- 传统方法和编码器-解码器框架均存在建模能力不足的问题。
- 需要进一步研究如何统一降采样和上采样任务。
图像去色-上色:
- 当前方法难以平衡真实性与多样性,需解决任务的双向建模问题。
图像压缩:
- 图像压缩与恢复多为独立处理,缺乏协同优化。
- 实现压缩与恢复任务的一体化建模是未来的重要方向。
可逆神经网络:
- 可逆神经网络为双向任务建模提供了理论基础。
- 其潜力尚未被完全挖掘,需探索更多的应用场景
方法
作者oral:ECCV 2020 Oral | 可逆图像缩放:完美恢复降采样后的高清图片 - 知乎
作者对本文方法的介绍,上面是其介绍的链接。
非常精彩的工作,我的理解是本文通过类似小波变换将图像分为了高频和低频的部分,低频部分就是下采样的图片,而高频信息则映射到一个高斯分布Z’~N(0,I)中,这样相当于对于超分网络有了高频信息的先验,有了先验,自然性能会很猛。这一点也是和不同数据集自身分布是不一致的统一。
看了其他评论,我个人觉得其实并不是压缩,如果是压缩的话,映射的高频域也是需要进行量化编码的,而文中直接丢弃了,换句话说没有后验信息,所以感觉不是压缩。
总的来说,论文的核心方法就是利用可逆神经网络建立一个严格的双向映射 x ↔ (y, z),在降采样阶段将图像信息分解为视觉有效的低分辨率部分 y 和服从预设分布的潜变量 z,在重构阶段利用 y 和随机采样的 z 重构出 x,从而在理论上克服传统下采样中信息丢失带来的难题。此方法不仅能实现较好的 PSNR/SSIM,还能在感知质量和模型轻量化上取得优势。
理论上,由于模型是基于严格可逆的双射设计,即将原始图像 x 映射为 (y, z) 后,再通过逆映射 f⁻¹(y, z) 恢复 x,因此如果模型完全理想,并且 z 是按照预先指定的分布采样,那么任取 z(例如随机采样)都应能重构出与原始图像 x 等价的图像。
但在实际训练中,由于模型的能力和训练过程中的数值近似,不同的 z 采样往往会导致重构图像在细微的高频细节上存在轻微差异。这种差异通常体现在那些不易察觉的纹理或噪声上,而整体结构和低频内容保持不变。论文中的实验表明,不同 z 采样得到的重构图像在 PSNR 上的差异极小(通常小于 0.02 dB),这说明随机采样的 z 主要影响的是高频的、细微的细节,而不会导致视觉上明显的失真。
因此,尽管重构过程中 z 是随机采样的,但在训练充分且分布匹配达到要求的情况下,模型可以稳定地恢复出与原始图像非常接近的重构图像。
在本节中,我们首先正式地介绍了图像退化-恢复问题的一般数学公式。3.1.然后,我们在章节中描述了这个双向问题的可逆建模框架。3.2.对于具体的模型,我们从图像缩放开始,阐述了节中IRN的具体可逆结构和训练方法。3.3.然后,我们在章节中展示了IRN对类似的脱色-着色任务的适应。3.4.最后,我们提出将可逆框架与现有的退化方法结合起来。
论文问题的数字表示:
该方法将图像退化(如降采样、去色等)和恢复建模为一个双向问题,公式如下:
可逆框架的设计:
本文的可逆框架采用了一种双向可逆变换,在前向过程中将原始图像分解为:
退化图像y:符合目标退化集合(例如降采样得到的低分辨率图像)。
辅助变量z:捕获退化过程中的丢失信息,以固定的分布(如高斯分布)表示。
损失函数的设计:
重构损失
引导损失
分布匹配损失
Haar小波变换(Haar Transformation):
- Haar变换将输入图像分解为低频成分和三个方向的高频成分。低频部分直接用于生成LR图像,高频部分用来捕捉更多细节信息。
- 降低了图像的空间分辨率,同时保留尽可能多的信息。
IRN 的框架主要分为三个部分:
- Haar 小波变换(Haar Transformation):
- 首先对输入高分辨率图像 x 进行 Haar 小波变换,将图像分解为低频部分 h1lh_1^lh1l 和高频部分 h2lh_2^lh2l。
- 低频部分$ h_1^l$: 代表图像的主要结构和内容(类似于降采样结果)。
- 高频部分h2lh_2^lh2l: 捕获图像的纹理细节和边缘特征。
- 这种分解为后续信息处理提供了基础。
- 可逆模块(InvBlocks):
- InvBlocks 用于对 Haar 小波的低频和高频信息进行进一步的处理和变换。
- 结构解析:
- 对 低频部分呢和 高频部分分别进行不同的变换:
- ϕ:用于处理低频部分 h1l。
- η:用于处理高频部分 h2l,捕获更多细节。
- ρ:将 h1l和 h2l的信息结合,增加网络的表达能力。
- 这些操作的结果通过可逆耦合层组合后,形成新的高频和低频特征 h1l+1和 h2l+1。
- 对 低频部分呢和 高频部分分别进行不同的变换:
- 通过多层 InvBlocks 的堆叠,进一步提升了特征表达的多样性。
- 降采样模块(Downscaling Module):
- 降采样模块用于逐步降低输入图像的空间分辨率,每次分辨率降低 2 倍。
- 输出:
- 最终得到低分辨率图像 y。
- 通过学习得到潜在变量 z,用于保存高频信息中丢失的细节。
框架前向过程:下采样
框架的反向过程:上采样
实验(见IRN_exp)
我们的实验包括三个部分:图像重新缩放、图像去色-着色,以及图像重新缩放和压缩的组合。对于所有任务的训练,我们采用广泛使用的DIV2K(Agustsson & Timofte,2017)图像恢复数据集来训练我们的模型。它包含 800 张高质量的 2K 分辨率训练图像和 100 张验证图像。此外,对于前两项任务,我们在4个额外的标准数据集上评估我们的模型,即Set5 , Set14 , BSD100 和 Urban100 ;对于第三项任务,我们还在广泛使用的柯达数据集上评估了我们的模型 (Franzen, 1999)。对于图像缩放,按照 (Lim, et al., 2017) 中的设置,我们定量评估峰值噪声信号比 (PSNR) 和 SSIM 在 YCbCr (Y, Cb, Cr) 颜色空间中表示的图像的 Y 通道上。我们还评估了LPIPS,PI和FID作为感知评估的定量指标。对于其他两项任务,我们在 RGB 三通道色彩空间上评估 PSNR 和 SSIM。
项目代码以及测试
作者oral:ECCV 2020 Oral | 可逆图像缩放:完美恢复降采样后的高清图片 - 知乎
github:https://github.com/pkuxmq/Invertible-Image-Rescaling?tab=readme-ov-file
数据集下载:datasets - Google 云端硬盘
预训练模型下载:pretrained_models_IRN_all - Google 云端硬盘
invertible image rescaling可逆图片缩放网络配置(仅用cpu) - 简书
DIV2K 数据集下载:DIV2K Dataset
半调数据集:
CIFAR-10 and CIFAR-100 datasets
zalandoresearch/fashion-mnist: A MNIST-like fashion product database. Benchmark
CIFAR-10 and CIFAR-100 datasets
环境以及安装
- Python 3 (Recommend to use Anaconda)
- PyTorch >= 1.0
- NVIDIA :GPU + CUDA
- Python packages:
pip install numpy opencv-python lmdb pyyaml
- TensorBoard:
- PyTorch >= 1.1:
pip install tb-nightly future
- PyTorch == 1.0:
pip install tensorboardX
- PyTorch >= 1.1:
PyTorch官网上直接根据服务器GPU CUDA版本下载对应的PyTorch版本就行。
数据集下载
Commonly used training and testing datasets can be downloaded here.
从这里下载数据集,进行实验。
修改配置文件
载入需要测试的图片路径,以及预训练的IRN 4×模型,输出HR图像。
test_IRN_x4.yml
name: IRN_x4
suffix: ~ # add suffix to saved images
model: IRN
distortion: sr
scale: 4
crop_border: ~ # crop border when evaluation. If None(~), crop the scale pixels
gpu_ids: [0]datasets:test_1: # the 1st test datasetname: set5mode: LQGTdataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/Set5/GTmod12 # path to test HR imagesdataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloadertest_2: # the 2st test datasetname: set14mode: LQGTdataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/Set14/GTmod12 # path to test HR imagesdataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloadertest_3: # the 3st test datasetname: B100mode: LQGTdataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/BSDS100 # path to test HR imagesdataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloadertest_4: # the 3st test datasetname: Urban100mode: LQGTdataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/urban100 # path to test HR imagesdataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloadertest_5:name: val_DIV2Kmode: LQGTdataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/DIV2K_valid_HR # path to test HR imagesdataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader#### network
network_G:which_model_G:subnet_type: DBNetin_nc: 3out_nc: 3block_num: [8, 8]scale: 4init: xavier#### path
path:pretrain_model_G: /home/hxr/Invertible-Image-Rescaling-master/IRN_Pretrain_model/IRN_x4.pth
test3
name: IRN_x4
suffix: ~ # add suffix to saved images
model: IRN
distortion: sr
scale: 4
crop_border: ~ # crop border when evaluation. If None(~), crop the scale pixels
gpu_ids: [0]datasets:# test_1:# name: test1# mode: LQ# dataroot_GT: ~ # path to test HR images# dataroot_LQ: /home/hxr/Invertible-Image-Rescaling-master/Classical/CelebA_64 # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader# test_1:# name: test2# mode: LQGT# dataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/AHCQ-64 # path to test HR images# dataroot_LQ: ~test_1:name: test3mode: LQGTdataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/CelebA_256 # path to test HR imagesdataroot_LQ: ~# test_1: # the 1st test dataset# name: set5# mode: LQGT# dataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/Set5/GTmod12 # path to test HR images# dataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader# test_2: # the 2st test dataset# name: set14# mode: LQGT# dataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/Set14/GTmod12 # path to test HR images# dataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader# test_3: # the 3st test dataset# name: B100# mode: LQGT# dataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/BSDS100 # path to test HR images# dataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader# test_4: # the 3st test dataset# name: Urban100# mode: LQGT# dataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/urban100 # path to test HR images# dataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader# test_5:# name: val_DIV2K# mode: LQGT# dataroot_GT: /home/hxr/Invertible-Image-Rescaling-master/Classical/DIV2K_valid_HR # path to test HR images# dataroot_LQ: ~ # path to test reference LR images, not necessary, if not provided, LR images will be generated in dataloader#### network
network_G:which_model_G:subnet_type: DBNetin_nc: 3out_nc: 3block_num: [8, 8]scale: 4init: xavier#### path
path:pretrain_model_G: /home/hxr/Invertible-Image-Rescaling-master/IRN_Pretrain_model/IRN_x4.pth
test.py
import os.path as osp
import logging
import time
import argparse
from collections import OrderedDictimport numpy as np
import options.options as option
import utils.util as util
from data.util import bgr2ycbcr
from data import create_dataset, create_dataloader
from models import create_model#### options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='Path to options YMAL file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
opt = option.dict_to_nonedict(opt)util.mkdirs((path for key, path in opt['path'].items()if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key))
util.setup_logger('base', opt['path']['log'], 'test_' + opt['name'], level=logging.INFO,screen=True, tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))#### Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):test_set = create_dataset(dataset_opt)test_loader = create_dataloader(test_set, dataset_opt)logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set)))test_loaders.append(test_loader)model = create_model(opt)
for test_loader in test_loaders:test_set_name = test_loader.dataset.opt['name']logger.info('\nTesting [{:s}]...'.format(test_set_name))test_start_time = time.time()dataset_dir = osp.join(opt['path']['results_root'], test_set_name)util.mkdir(dataset_dir)test_results = OrderedDict()test_results['psnr'] = []test_results['ssim'] = []test_results['psnr_y'] = []test_results['ssim_y'] = []test_results['psnr_lr'] = []test_results['ssim_lr'] = []test_results['psnr_y_lr'] = []test_results['ssim_y_lr'] = []for data in test_loader:# img_path = data['LQ_path'][0]# img_name = osp.splitext(osp.basename(img_path))[0]# output_img = model.upscale(data['LQ'].cuda(),4).detach()[0].float().cpu()# sr_img = util.tensor2img(output_img)# # save images# suffix = opt['suffix']# if suffix:# save_img_path = osp.join(dataset_dir, img_name + suffix + '.png')# else:# save_img_path = osp.join(dataset_dir, img_name + '.png')# util.save_img(sr_img, save_img_path)model.feed_data(data)img_path = data['GT_path'][0]img_name = osp.splitext(osp.basename(img_path))[0]model.test()visuals = model.get_current_visuals()sr_img = util.tensor2img(visuals['SR']) # uint8srgt_img = util.tensor2img(visuals['GT']) # uint8lr_img = util.tensor2img(visuals['LR']) # uint8lrgt_img = util.tensor2img(visuals['LR_ref']) # uint8# save imagessuffix = opt['suffix']if suffix:save_img_path = osp.join(dataset_dir, img_name + suffix + '.png')else:save_img_path = osp.join(dataset_dir, img_name + '.png')util.save_img(sr_img, save_img_path)if suffix:save_img_path = osp.join(dataset_dir, img_name + suffix + '_GT.png')else:save_img_path = osp.join(dataset_dir, img_name + '_GT.png')util.save_img(srgt_img, save_img_path)if suffix:save_img_path = osp.join(dataset_dir, img_name + suffix + '_LR.png')else:save_img_path = osp.join(dataset_dir, img_name + '_LR.png')util.save_img(lr_img, save_img_path)if suffix:save_img_path = osp.join(dataset_dir, img_name + suffix + '_LR_ref.png')else:save_img_path = osp.join(dataset_dir, img_name + '_LR_ref.png')util.save_img(lrgt_img, save_img_path)# calculate PSNR and SSIMgt_img = util.tensor2img(visuals['GT'])gt_img = gt_img / 255.sr_img = sr_img / 255.lr_img = lr_img / 255.lrgt_img = lrgt_img / 255.crop_border = opt['crop_border'] if opt['crop_border'] else opt['scale']if crop_border == 0:cropped_sr_img = sr_imgcropped_gt_img = gt_imgelse:cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :]cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :]psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255)ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255)test_results['psnr'].append(psnr)test_results['ssim'].append(ssim)# PSNR and SSIM for LRpsnr_lr = util.calculate_psnr(lr_img * 255, lrgt_img * 255)ssim_lr = util.calculate_ssim(lr_img * 255, lrgt_img * 255)test_results['psnr_lr'].append(psnr_lr)test_results['ssim_lr'].append(ssim_lr)if gt_img.shape[2] == 3: # RGB imagesr_img_y = bgr2ycbcr(sr_img, only_y=True)gt_img_y = bgr2ycbcr(gt_img, only_y=True)if crop_border == 0:cropped_sr_img_y = sr_img_ycropped_gt_img_y = gt_img_yelse:cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border]cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border]psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255)ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255)test_results['psnr_y'].append(psnr_y)test_results['ssim_y'].append(ssim_y)lr_img_y = bgr2ycbcr(lr_img, only_y=True)lrgt_img_y = bgr2ycbcr(lrgt_img, only_y=True)psnr_y_lr = util.calculate_psnr(lr_img_y * 255, lrgt_img_y * 255)ssim_y_lr = util.calculate_ssim(lr_img_y * 255, lrgt_img_y * 255)test_results['psnr_y_lr'].append(psnr_y_lr)test_results['ssim_y_lr'].append(ssim_y_lr)logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}. LR PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'.format(img_name, psnr, ssim, psnr_y, ssim_y, psnr_lr, ssim_lr, psnr_y_lr, ssim_y_lr))else:logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}. LR PSNR: {:.6f} dB; SSIM: {:.6f}.'.format(img_name, psnr, ssim, psnr_lr, ssim_lr))# Average PSNR/SSIM resultsave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])ave_psnr_lr = sum(test_results['psnr_lr']) / len(test_results['psnr_lr'])ave_ssim_lr = sum(test_results['ssim_lr']) / len(test_results['ssim_lr'])logger.info('----Average PSNR/SSIM results for {}----\n\tpsnr: {:.6f} db; ssim: {:.6f}. LR psnr: {:.6f} db; ssim: {:.6f}.\n'.format(test_set_name, ave_psnr, ave_ssim, ave_psnr_lr, ave_ssim_lr))if test_results['psnr_y'] and test_results['ssim_y']:ave_psnr_y = sum(test_results['psnr_y']) / len(test_results['psnr_y'])ave_ssim_y = sum(test_results['ssim_y']) / len(test_results['ssim_y'])ave_psnr_y_lr = sum(test_results['psnr_y_lr']) / len(test_results['psnr_y_lr'])ave_ssim_y_lr = sum(test_results['ssim_y_lr']) / len(test_results['ssim_y_lr'])logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}. LR PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.\n'.format(ave_psnr_y, ave_ssim_y, ave_psnr_y_lr, ave_ssim_y_lr))
该论文主要在 DIV2K 数据集上进行模型训练 (800 张图像)。
评测时,会在 Set5、Set14、BSD100、Urban100、DIV2K 的验证集 以及 Kodak 这几个常见测试/验证数据集上做定量对比 (PSNR、SSIM 等)。
Training for image rescaling
First set a config file in options/train/, then run as following:
python train.py -opt options/train/train_IRN_x4.yml
Testing for image rescaling
First set a config file in options/test/, then run as following:
python test.py -opt options/test/test_IRN_x4.yml
Training for image decolorization-colorization
First set a config file in options/train/, then run as following:
python train_IRN-Color.py -opt options/train/train_IRN_color.yml
Testing for image decolorization-colorization
First set a config file in options/test/, then run as following:
python test_IRN-Color.py -opt options/test/test_IRN_color.yml
Training for combination with image compression
First set a config file in options/train/, then run as following:
python train_IRN-Compression.py -opt options/train/train_IRN-Compression_x2_q90.yml
Testing for combination with image compression
First set a config file in options/test/, then run as following:
python test_IRN-Compression.py -opt options/test/test_IRN-Compression_x2_q90.yml
Pretrained models can be downloaded from Google Drive or Baidu Drive (extraction code: rx0z).
加密隐写图像传输——接收方接收超分辨率重建后的加密隐写图像
from fastai.vision.all import *
from diffusers import DDIMPipeline,DiffusionPipeline# def set_random_seed(seed=0):
# torch.manual_seed(seed + 0)
# torch.cuda.manual_seed(seed + 1)
# torch.cuda.manual_seed_all(seed + 2)
# np.random.seed(seed + 3)
# torch.cuda.manual_seed_all(seed + 4)
# random.seed(seed + 5)# set_random_seed()
latent_channel=3
args={"gen_seed": 0,"guidance_scale":7.5,"num_inference_steps":200,"num_inversion_steps":200,"channel_copy":latent_channel,"hw_copy":1,"output_path":"output","model_path": "google/ddpm-celebahq-256",
}args=dict2obj(args)
image_length=256
secret_size=(image_length,image_length)
batch_size=1
latent_shape = (batch_size, latent_channel, image_length, image_length)sd_pipe = DiffusionPipeline.from_pretrained(args.model_path)
sd_pipe.to('cuda')
image=sd_pipe(num_inference_steps=200).images[0]
print(image.size)
image
基于 IRN(Invertible Rescaling Network)的超分辨率功能,我们可以实现一个加密隐写小图到加密隐写大图的转换。IRN 的主要作用是提供可逆的图像缩放,即能够在低分辨率图像(LR)与高分辨率图像(HR)之间进行无损转换,同时保持信息完整性。这使得我们可以将隐写信息嵌入小图,并在放大后保持隐写信息的完整性。
IRN将小图转化成大图
250221
250225
SR5-7.py
import os
import os.path as osp
import logging
import time
import torch
from data import create_dataset, create_dataloader
from models import create_model
import utils.util as utilclass NoneDict(dict):"""A dict class that returns None for missing keys instead of raising a KeyError."""def __missing__(self, key):return Nonedef dict_to_nonedict(opt):"""Recursively convert all nested dicts into NoneDict,so that missing keys return None instead of raising a KeyError."""if isinstance(opt, dict):new_opt = {}for k, v in opt.items():new_opt[k] = dict_to_nonedict(v)return NoneDict(**new_opt)elif isinstance(opt, list):return [dict_to_nonedict(e) for e in opt]else:return optdef parse_opt(opt, is_train=False):"""Process the input opt dict:1. Set GPU environment.2. Mark train/test mode.3. If SR task, override dataset scale.4. Prepare and create necessary directories under `opt['path']`."""# 1. 设置 GPU 设备gpu_list = ','.join(str(x) for x in opt.get('gpu_ids', []))os.environ['CUDA_VISIBLE_DEVICES'] = gpu_listprint('export CUDA_VISIBLE_DEVICES=' + gpu_list)# 2. 标记是否是训练模式opt['is_train'] = is_train# 3. 如果是 SR 任务,对 dataset 的 scale 做覆盖if opt['distortion'] == 'sr' and 'scale' in opt:scale = opt['scale']for phase, dataset_opt in opt['datasets'].items():# "train_1" => "train", "test_1" => "test"dataset_opt['phase'] = phase.split('_')[0]dataset_opt['scale'] = scale# 4. 准备运行目录结构current_root = os.getcwd()if 'path' not in opt:opt['path'] = {}opt['path']['root'] = current_rootif is_train:# 训练模式:把结果写到 experiments/<name>/ ...experiments_root = osp.join(current_root, 'experiments', opt['name'])opt['path']['experiments_root'] = experiments_rootopt['path']['models'] = osp.join(experiments_root, 'models')opt['path']['training_state'] = osp.join(experiments_root, 'training_state')opt['path']['log'] = experiments_rootopt['path']['val_images'] = osp.join(experiments_root, 'val_images')else:# 测试模式:把结果写到 results/<name>/ ...results_root = osp.join(current_root, 'results', opt['name'])opt['path']['results_root'] = results_rootopt['path']['log'] = results_rootreturn optdef dict2str(opt, indent_level=0):"""Convert a nested dict into a readable string for logging."""msg = ''for k, v in opt.items():if isinstance(v, dict):msg += ' ' * indent_level + f"{k}:\n"msg += dict2str(v, indent_level + 1)else:msg += ' ' * indent_level + f"{k}: {v}\n"return msgdef main():"""Main entry for the SR test pipeline.1. Hardcode a config dict opt.2. parse_opt for GPU setup, path creation, etc.3. Convert dict to NoneDict to avoid KeyError.4. Initialize logging.5. Create dataset and model.6. Perform inference and save results."""# ========== 1. 硬编码配置 ==========opt = {'name': 'IRN_x4','model': 'IRN','suffix': '','distortion': 'sr','scale': 4,'gpu_ids': [0],'dist': False, # 是否分布式'datasets': {'test_1': {'name': 'test1','mode': 'LQ','dataroot_GT': None, # 如果无需 Ground Truth,可设为 None'dataroot_LQ': '/home/hxr/IRN/input','data_type': 'img'}},'network_G': {'which_model_G': {'subnet_type': 'DBNet','use_ConvDownsampling': False,'down_first': False},'in_nc': 3,'out_nc': 3,'block_num': [8, 8],'scale': 4,'init': 'xavier'# 如果需要gc,可手动加 'gc': 32},'path': {'pretrain_model_G': '/home/hxr/IRN/IRN_x4.pth','strict_load': False}}# ========== 2. 解析配置,处理 GPU, 路径等 ==========opt = parse_opt(opt, is_train=False)# ========== 3. dict -> NoneDict ==========opt = dict_to_nonedict(opt)# ========== 4. 创建输出目录、日志 ==========util.mkdirs([opt['path']['results_root'], opt['path']['log']])util.setup_logger('base',opt['path']['log'],'test_' + opt['name'],level=logging.INFO,screen=True,tofile=True)logger = logging.getLogger('base')# 输出配置信息logger.info("==== Test configuration ====\n%s", dict2str(opt))logger.info("Results will be saved to: %s", opt['path']['results_root'])# ========== 5. 创建数据集和 DataLoader ==========test_loaders = []for phase, dataset_opt in sorted(opt['datasets'].items()):test_set = create_dataset(dataset_opt)test_loader = create_dataloader(test_set, dataset_opt)logger.info("Test set [%s]: number of images = %d",dataset_opt['name'], len(test_set))test_loaders.append(test_loader)# ========== 6. 创建模型 ==========model = create_model(opt)# ========== 7. 进行推理并保存结果 ==========for test_loader in test_loaders:test_set_name = test_loader.dataset.opt['name']logger.info("\n>>> Testing [%s]...", test_set_name)start_time = time.time()# 对应数据集的输出目录dataset_dir = osp.join(opt['path']['results_root'], test_set_name)util.mkdir(dataset_dir)suffix = opt['suffix'] or ''scale = opt['scale'] or 4for data in test_loader:img_path = data['LQ_path'][0]img_name = osp.splitext(osp.basename(img_path))[0]# 无需自动求导with torch.no_grad():output_img = model.upscale(data['LQ'].cuda(), scale)output_img = output_img.detach()[0].float().cpu()# 转为可保存的图像格式sr_img = util.tensor2img(output_img)# 生成文件名并保存save_name = img_name + (suffix + '.png' if suffix else '.png')save_img_path = osp.join(dataset_dir, save_name)util.save_img(sr_img, save_img_path)elapsed = time.time() - start_timelogger.info("Finished testing [%s], time = %.3fs", test_set_name, elapsed)logger.info("All tests finished. Results are in: %s", opt['path']['results_root'])if __name__ == '__main__':main()
测试结果:
CRM.py
import os
import os.path as osp
import logging
import time
import cv2
import torch
from collections import OrderedDict# ====== 以下 import 需要结合你的实际项目结构调整 ======
# 比如 data、models、utils 这些包请根据你本地的情况更改路径
from data import create_dataset, create_dataloader
from models import create_model
import utils.util as util
from data.util import bgr2ycbcrclass NoneDict(dict):"""A dict class that returns None for missing keys instead of raising a KeyError."""def __missing__(self, key):return Nonedef dict_to_nonedict(opt):"""Recursively convert all nested dicts into NoneDict,so that missing keys return None instead of raising a KeyError."""if isinstance(opt, dict):new_opt = {}for k, v in opt.items():new_opt[k] = dict_to_nonedict(v)return NoneDict(**new_opt)elif isinstance(opt, list):return [dict_to_nonedict(e) for e in opt]else:return optdef dict2str(opt, indent_level=0):"""Convert a nested dict into a readable string for logging."""msg = ''for k, v in opt.items():if isinstance(v, dict):msg += ' ' * indent_level + f"{k}:\n"msg += dict2str(v, indent_level + 1)else:msg += ' ' * indent_level + f"{k}: {v}\n"return msgdef parse_opt(opt):"""解析配置字典,设置 GPU、创建输出目录等。这里简化了 SR7.py 中的 parse_opt,可根据需要自行扩展。"""# 设置 GPUgpu_list = ','.join(str(x) for x in opt.get('gpu_ids', []))os.environ['CUDA_VISIBLE_DEVICES'] = gpu_listprint('[Info] export CUDA_VISIBLE_DEVICES=' + gpu_list)# 标记训练或测试模式,这里默认是测试opt['is_train'] = False# 如果需要,可以在此处对路径、log 等做更灵活的处理。# 简化:所有结果写到 results/<name>/current_root = os.getcwd()if 'path' not in opt:opt['path'] = {}results_root = osp.join(current_root, 'results', opt['name'])opt['path']['results_root'] = results_rootopt['path']['log'] = results_rootreturn optdef main():"""主函数:加载配置 -> 创建数据集和模型 -> 进行推理(下采样+JPEG压缩)-> 计算/保存结果"""# ========== 2.1. 构造配置字典(将原 test_IRN-Compression_x2_q90.yml 融进来) ==========# 下例中只保留最关键的参数,可根据需要自行增删opt = {'name': 'IRN_Compress_x2_q90','model': 'IRN-CRM','gpu_ids': [0],'dist': False,# IRN 通常既可做下采样也可做重建,具体由模型结构与预训练权重决定'scale': 2, # 代表想把原图宽高各除以2'test': {'jpg_quality': 90,# 如果希望在评价时裁边,可加上 crop_border;若不裁边,则可设 0'crop_border': 0},'datasets': {# 假设只有一个测试集'test_1': {'name': 'CompressTestSet','mode': 'LQGT', # 模式可根据你的 create_dataset 实现来调整'dataroot_GT': '/kaggle/working/results/IRN_x4/test1', # 大图路径'dataroot_LQ': None, # 如无单独 LQ 数据集,可设 None'scale': 2,'data_type': 'img'}},'network_G': {'which_model_G': {'subnet_type': 'DBNet','use_ConvDownsampling': False,'down_first': False},'in_nc': 3,'out_nc': 3,'block_num': [8],'scale': 2,'init': 'xavier'# 如果需要gc,可手动加 'gc': 32},'network_R': {'in_nc': 3,'out_nc': 3,'nf': 64,'nb': 8,'gc': 32,},'path': {# 可在此指定预训练模型路径'pretrain_model_G': '/kaggle/working/Invertible-Image-Rescaling/pretrained_models/pretrained_models_IRN_all/IRN_x2_finetune.pth','pretrain_model_R': '/kaggle/working/Invertible-Image-Rescaling/pretrained_models/pretrained_models_IRN_all/CRM_x2_q90.pth','strict_load': False # 根据需要决定是否严格匹配权重}}# ========== 2.2. parse_opt:处理 GPU、创建日志/输出目录等 ==========opt = parse_opt(opt)opt = dict_to_nonedict(opt)# ========== 2.3. 初始化日志器 ==========util.mkdirs([opt['path']['results_root'], opt['path']['log']])util.setup_logger('base',opt['path']['log'],'test_' + opt['name'],level=logging.INFO,screen=True,tofile=True)logger = logging.getLogger('base')# 打印配置logger.info('[Config]')logger.info(dict2str(opt))logger.info('[Paths]')logger.info('Results will be saved to: {}'.format(opt['path']['results_root']))# ========== 3. 创建数据集和模型 ==========test_loaders = []for phase, dataset_opt in sorted(opt['datasets'].items()):test_set = create_dataset(dataset_opt)test_loader = create_dataloader(test_set, dataset_opt)logger.info('Test set [%s]: number of images = %d',dataset_opt['name'], len(test_set))test_loaders.append(test_loader)# 创建模型(IRN)model = create_model(opt)# ========== 4. 测试:遍历数据集,做下采样及统计 ==========for test_loader in test_loaders:test_set_name = test_loader.dataset.opt['name']logger.info('\n>>> Testing [%s]...', test_set_name)test_start_time = time.time()# 输出结果目录dataset_dir = osp.join(opt['path']['results_root'], test_set_name)util.mkdir(dataset_dir)# 用于统计平均指标test_results = OrderedDict()for k in ['psnr', 'ssim', 'psnr_y', 'ssim_y','psnr_lr', 'ssim_lr', 'psnr_y_lr', 'ssim_y_lr','bpp', 'psnr_restore']:test_results[k] = []# ---------- 推理循环 -----------for data in test_loader:# 将数据送入模型model.feed_data(data)img_path = data['GT_path'][0] # 假设存在 GTimg_name = osp.splitext(osp.basename(img_path))[0]# 推理model.test()visuals = model.get_current_visuals()# 获取 SR、GT、LR 等图像 (uint8)# 其中 SR 通常是重建后得到的图,LR 是下采样+JPEG压缩后的图sr_img = util.tensor2img(visuals['SR'])srgt_img = util.tensor2img(visuals['GT'])lr_img = util.tensor2img(visuals['LR'])lrgt_img = util.tensor2img(visuals['LR_ref']) # downsampled “GT”rlr_img = util.tensor2img(visuals['RLR']) # restored LR?# 保存 SR(对应模型重建结果)suffix = opt.get('suffix', '')save_sr_path = osp.join(dataset_dir, img_name + suffix + '.png')util.save_img(sr_img, save_sr_path)# 保存 SR 对应的 GT(原图)save_srgt_path = osp.join(dataset_dir, img_name + suffix + '_GT.png')util.save_img(srgt_img, save_srgt_path)# 保存 LR 图(即缩放后再 JPEG 压缩的结果)quality = opt['test']['jpg_quality']save_lr_path = osp.join(dataset_dir, img_name + suffix + '_LR.jpg')cv2.imwrite(save_lr_path, lr_img, [int(cv2.IMWRITE_JPEG_QUALITY), quality])# 获取 LR 图的大小,用于计算 bpplr_size = osp.getsize(save_lr_path)# 保存下采样“参考”图(若有需要做比对时,可用 PNG 不失真)save_lrgt_path = osp.join(dataset_dir, img_name + suffix + '_LR_ref.png')util.save_img(lrgt_img, save_lrgt_path)# ========== 计算指标 ==========# bpp(比特率)h, w = srgt_img.shape[:2]bpp = lr_size * 8.0 / (h * w)test_results['bpp'].append(bpp)# 转为 [0,1] 浮点计算 PSNR/SSIMsr_img_f = sr_img / 255.0srgt_img_f = srgt_img / 255.0lr_img_f = lr_img / 255.0lrgt_img_f = lrgt_img / 255.0crop_border = opt['test'].get('crop_border', 0)if crop_border > 0:cropped_sr_img = sr_img_f[crop_border:-crop_border, crop_border:-crop_border, :]cropped_srgt_img = srgt_img_f[crop_border:-crop_border, crop_border:-crop_border, :]else:cropped_sr_img = sr_img_fcropped_srgt_img = srgt_img_fpsnr = util.calculate_psnr(cropped_sr_img * 255, cropped_srgt_img * 255)ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_srgt_img * 255)test_results['psnr'].append(psnr)test_results['ssim'].append(ssim)# LR 图自身的 PSNR/SSIM(可作为对比)psnr_lr = util.calculate_psnr(lr_img_f * 255, lrgt_img_f * 255)ssim_lr = util.calculate_ssim(lr_img_f * 255, lrgt_img_f * 255)test_results['psnr_lr'].append(psnr_lr)test_results['ssim_lr'].append(ssim_lr)# PSNR of “restored LR” (如有相关过程)psnr_restore = util.calculate_psnr(lr_img_f * 255, rlr_img)test_results['psnr_restore'].append(psnr_restore)# 对 Y 通道做一次统计if sr_img_f.shape[2] == 3:sr_img_y = bgr2ycbcr(sr_img_f, only_y=True)srgt_img_y = bgr2ycbcr(srgt_img_f, only_y=True)if crop_border > 0:cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border]cropped_srgt_img_y = srgt_img_y[crop_border:-crop_border, crop_border:-crop_border]else:cropped_sr_img_y = sr_img_ycropped_srgt_img_y = srgt_img_ypsnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_srgt_img_y * 255)ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_srgt_img_y * 255)test_results['psnr_y'].append(psnr_y)test_results['ssim_y'].append(ssim_y)# LR 的 Y 通道lr_img_y = bgr2ycbcr(lr_img_f, only_y=True)lrgt_img_y = bgr2ycbcr(lrgt_img_f, only_y=True)psnr_y_lr = util.calculate_psnr(lr_img_y * 255, lrgt_img_y * 255)ssim_y_lr = util.calculate_ssim(lr_img_y * 255, lrgt_img_y * 255)test_results['psnr_y_lr'].append(psnr_y_lr)test_results['ssim_y_lr'].append(ssim_y_lr)logger.info('%s - PSNR: %.6f dB; SSIM: %.6f; PSNR_Y: %.6f dB; SSIM_Y: %.6f. ''LR PSNR: %.6f dB; SSIM: %.6f; LR PSNR_Y: %.6f dB; SSIM_Y: %.6f. ''PSNR Restore: %.6f, bpp: %.6f.',img_name, psnr, ssim, psnr_y, ssim_y,psnr_lr, ssim_lr, psnr_y_lr, ssim_y_lr,psnr_restore, bpp)else:logger.info('%s - PSNR: %.6f dB; SSIM: %.6f. LR PSNR: %.6f dB; SSIM: %.6f. PSNR Restore: %.6f, bpp: %.6f.',img_name, psnr, ssim, psnr_lr, ssim_lr, psnr_restore, bpp)# ---------- 统计平均指标 -----------def average(lst):return sum(lst) / len(lst) if lst else 0.0ave_psnr = average(test_results['psnr'])ave_ssim = average(test_results['ssim'])ave_psnr_lr = average(test_results['psnr_lr'])ave_ssim_lr = average(test_results['ssim_lr'])ave_psnr_restore = average(test_results['psnr_restore'])ave_bpp = average(test_results['bpp'])logger.info('----Average PSNR/SSIM for [%s]----\n\tPSNR: %.6f dB; SSIM: %.6f. ''LR PSNR: %.6f dB; LR SSIM: %.6f.',test_set_name, ave_psnr, ave_ssim, ave_psnr_lr, ave_ssim_lr)if test_results['psnr_y'] and test_results['ssim_y']:ave_psnr_y = average(test_results['psnr_y'])ave_ssim_y = average(test_results['ssim_y'])ave_psnr_y_lr = average(test_results['psnr_y_lr'])ave_ssim_y_lr = average(test_results['ssim_y_lr'])logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: %.6f dB; SSIM_Y: %.6f. ''LR PSNR_Y: %.6f dB; LR SSIM_Y: %.6f.',ave_psnr_y, ave_ssim_y, ave_psnr_y_lr, ave_ssim_y_lr)logger.info('----Average PSNR Restore and bpp----\n\tPSNR_Restore: %.6f dB; bpp: %.6f.\n',ave_psnr_restore, ave_bpp)cost_time = time.time() - test_start_timelogger.info('Finished testing [%s], time = %.3fs', test_set_name, cost_time)logger.info('All tests finished. See results in: %s', opt['path']['results_root'])if __name__ == '__main__':main()
p = average(test_results[‘bpp’])
logger.info(
'----Average PSNR/SSIM for [%s]----\n\tPSNR: %.6f dB; SSIM: %.6f. ’
‘LR PSNR: %.6f dB; LR SSIM: %.6f.’,
test_set_name, ave_psnr, ave_ssim, ave_psnr_lr, ave_ssim_lr
)
if test_results['psnr_y'] and test_results['ssim_y']:ave_psnr_y = average(test_results['psnr_y'])ave_ssim_y = average(test_results['ssim_y'])ave_psnr_y_lr = average(test_results['psnr_y_lr'])ave_ssim_y_lr = average(test_results['ssim_y_lr'])logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: %.6f dB; SSIM_Y: %.6f. ''LR PSNR_Y: %.6f dB; LR SSIM_Y: %.6f.',ave_psnr_y, ave_ssim_y, ave_psnr_y_lr, ave_ssim_y_lr)logger.info('----Average PSNR Restore and bpp----\n\tPSNR_Restore: %.6f dB; bpp: %.6f.\n',ave_psnr_restore, ave_bpp)cost_time = time.time() - test_start_timelogger.info('Finished testing [%s], time = %.3fs', test_set_name, cost_time)logger.info('All tests finished. See results in: %s', opt['path']['results_root'])
if name == ‘main’:
main()