Zero-Shot Noise2Noise:无需数据集训练的有效图像去噪方法

文章介绍了一种新的图像去噪方法,它基于Zero-ShotNoise2Noise和neighbor2neighbor,使用单个噪声图像生成一对噪声图像来训练一个简单的两层神经网络。这种方法无需任何训练数据或噪声模型,能以较低成本提供高质量的去噪效果。通过下采样和正则化避免过拟合,实现对噪声图像的有效去噪。
摘要由CSDN通过智能技术生成

注!这个博客是CSDN该论文最详细的免费的解析,引用请注明出处!!!

一.简介

        近年来,图像开始使用监督网络进行去噪,并且表现出来良好的性能,但是当前的无数据集的方法要么计算成本高,要不就是需要噪声模型,或者就是图像质量较低,因此作者提出了一个简单的两层网络,以没有任何训练数据或者噪声分布的方法,用较低的成本去高质量的对图片进行去噪,该方法是一种以更低的成本并且优于现有的其他的无数据去噪的方法。

代码地址:https://colab.research.google.com/drive/1i82nyizTdszyHkaHBuKPbWnTzao8HF9b?usp=sharing

论文地址:

https://openaccess.thecvf.com/content/CVPR2023/papers/Mansour_Zero-Shot_Noise2Noise_Efficient_Image_Denoising_Without_Any_Data_CVPR_2023_paper.pdf

二.方法

        作者提出的方法建立在Noise2Noise和neighbor2neighbor的基础上,前者使用一对噪声图片训练神经网络,后者用单个噪声图片生成这样的一对噪声图。作者的方法是使用单个噪声图片生成一对噪声图片,并用这对噪声图片训练一个简单的(二层)神经网络。

三.简介Noise2Noise(N2N)和neighbor2neighbor(NB2NB)

        监督式的去噪方法通常是神经网络f_{\theta }将一个噪声图像y映射到干净图像x的似然估计函数f_{\theta }(y)中,监督式的去噪方法通常是对于干净图片x 和一个噪声图片y=x+e,e就是噪声,可以使用numpy库进行随机生成,这样的监督去噪称为Noise2Clean。

        神经网络也可以在同一张干净图像的不同噪声图上进行训练,就是说,第一幅干净图片施加两种不同的噪声。Noise2Noise假设访问了一对噪声图像y1=x+e1,y2=x+e2,e1和e2是两种独立的噪声向量,然后训练网络f_{\theta }以最小化均方损失误差\frac{1}{n} \sum_{n}^{i=1}\left \| f_{\theta}(y_{1}^{i})- y_{2}^{i}\right \|_{2}^{2},这是有道理的,假设有一张无噪声图片,生成了一对分别包含e1和e2噪声的实例图片,通过训练监督式神经网络的方式,将噪声图片映射到噪声图片和将噪声图片映射到干净图片是等价的,这句话就是说,因为噪声图片已经够脏了,映射到噪声图片不过也是脏脏联合 ,映射到干净图片也是一样的,不过就是弄脏了图片而已。所以就期望最小化参数θ

        理论上,如果数据集无限大,N2N训练与N2C训练达到相同的性能。在实践中,由于训练集的大小有限,N2N略低于N2C。例如,与使用UNet的N2C相比,在5万张图像上使用UNet进行N2N训练的性能下降仅约0.02 dB。

        尽管N2N的性能很好,但它的可用性往往受到限制,因为很难获得同一静态场景的一对噪声图像。例如,被捕获的物体可能是非静态的,或者光照条件变化很快。

        neighbor2neighbour改进了N2N,允许只在一组单个噪声图像上进行训练,通过对噪声图像进行子采样来生成一对噪声图像。与N2N相似,NB2NB在多幅图像上训练时表现出较强的去噪性能。

四.Zero-shot Noise2Noise

        作者的工作是改进了N2N和NB2NB,仅仅使用单个噪声图片进行训练,为了避免单个图像的过拟合,只使用了一个非常浅的网络和一个明确的正则化项来避免该现象。

        几乎所有的监督式或者非监督去噪方式,包括作者提出的方法,都依赖于一个前提,就是干净的自然图片与具有不同分布方式的噪声图片,噪声图像可以分解为一对下采样图像。基于干净图像附近高像素具有高相似性,相关性的值,而噪声像素是非结构化且独立的,对噪声图片进行下采样后的特征图具有相似的特征,但是噪声依然独立,因此,这一对噪声图片可以作为同一个场景的两个噪声观测值的近似值,其中一个观测值可以用作输入值,另一个用作目标。

        作者的方法是使用下采样算子先将图像分解为一对下采样图像,然后用正则化训练一个轻量级网络,将一个下采样图像映射到另一个下采样图像,然后将经过训练的网络应用于噪声图像得到去噪图像。

        下采样算子将一个形状为H*W*C的图片y作为输入,然后生成两张图片D1(y),D2(y),形状为H/2*W/2*C,下采样算子通过将图像分割成大小为2*2的不重叠小块来生成这两张图像,就是说,这个下采样算子是2*2的卷积核,步长为2,然后把每个小块的副对角线像素的平均值,将其赋值为第一个特征图D1(y)的第一个像素点,类似于平均池化。然后求主对角线的像素平均值,并将其赋值给另一个D2(y)的第一个像素点。

        这个2*2的不重叠小块可以用固定的卷积核K1表示,K_{1}=\bigl(\begin{smallmatrix} 0& 0.5\\ 0.5 & 0 \end{smallmatrix}\bigr),对于原始图像做如下卷积D_{1}(y)=y\circledast k_{1},就是矩阵做内积,用来获得第一个下采样图像;同样的第二个下采样图像使用K_{2}=\bigl(\begin{smallmatrix} 0.5&0\\ 0&0.5 \end{smallmatrix}\bigr)对原始图像进行内积,就是D_{2}(y)=y\circledast k_{2},卷积是根据通道数量实现的,因此下采样方案适用于任意数量的输入通道。

 

        通过给定的要去噪声的图像y,先拟合一个图像神经网络模型f_{\theta },通过损失最小化的方法将第一个下采样图像D1(y)映射到第二个下采样图像D2(y),损失计算如下

\iota (\theta )=\left \| f_{\theta }(D_{1}(y))-D_{2}(y) \right \|_{2}^{2}

        一旦我们拟合了这个网络,我们就可以将这个网络应用于原始噪声图像去估计去噪后的图像为

x\hat{}=f_{\theta \hat{}}(y)

        然而,实验表明,残差学习,对称损失和额外的一致性强化项(正则项)对模型的表现至关重要,在残差学习中,网络被优化用来适应噪声而不是图像,因为噪声下采样后可能更容易观察和体现出噪声。下采样是指在图像或信号上减少采样点的数量,从而降低其分辨率。当噪声存在于原始图像中时,下采样会对噪声进行传播和放大,导致噪声在降采样后更加明显,噪声在图像中通常以高频成分的形式存在,而下采样会丢失部分高频信息。因此,当图像中存在高频噪声时,下采样会导致噪声在降采样后更加明显,可能会使噪声变得更容易被观察和感知,因此我们在下采样的图片D1(y)和D2(y)中进行损失计算,因此损失可以表示成:

\iota (\theta )=\left \| D_{1}(y)-f_{\theta }(D_{1}(y))-D_{2}(y) \right \|_{2}^{2}

        在siamese网络自监督预训练中使用了对称损失,在去噪时,我们也用到了对称损失,这时候正则损失会变成
\iota (\theta )_{residual}=\frac{1}{2}(\left \| D_{1}(y)-f_{\theta }(D_{1}(y))-D_{2}(y) \right \|_{2}^{2}+\left \| D_{2}(y)-f_{\theta }(D_{2}(y))-D_{1}(y) \right \|_{2}^{2})

        此外,为了图像确保一致性,我们对原始图像y进行去噪后再进行下采样,类似于我们首先对原始图像y进行下采样然后再对其进行去噪,损失函数如下

\iota (\theta )=\left \| D(y)-f_{\theta }(D(y))-D(y) \right \|_{2}^{2}

        最后采用对称损失,则一致性损失变成:

\iota (\theta )_{cons.}=\frac{1}{2}(\left \| D_{1}(y)-f_{\theta }(D_{1}(y))-D_{1}(y) \right \|_{2}^{2}+\left \| D_{2}(y)-f_{\theta }(D_{2}(y))-D_{2}(y) \right \|_{2}^{2})

        注意,对于残差损失,网络只将下采样图像作为输入。只有在一致性损失的情况下,网络才能看到全空间分辨率的图像。包括一致性损失可以提高去噪性能,并有助于避免过拟合。因此,它可以被看作是一个正则化项。

        因此综上所述,我们为了最小化损失\iota (\theta )=\iota (\theta )_{residual}+\iota (\theta )_{cons.},使用了梯度下降,最后生成最佳网络参数θ,有了这些,我们把去噪图像设为x\hat{}=y-f_{\theta \hat{}}(y),在梯度下降更新过程中,只有网络参数θ被优化更新,因为下采样操作生成的图像D1和D2是固定的,融合大约需要1k到2K次迭代,这要归功于这个两层的轻量级网络,在GPU上大概需要半分钟,CPU上需要一分钟

        许多监督和自监督方法使用一个相对较大的网络,通常是一个UNet,作者使用了一个非常简单的两层图像网络,第一层是3*3的卷积核,第二层是1*1的卷积核,这个网络大约有20k参数,这是非常轻量级的去噪声网络了。没有标准化或池化层。参数数少,结构简单,即使部署在CPU上也能快速去噪。在消融实验中作者用UNet代替轻量级网络会导致过拟合和更差的去噪性能。

        性能在高斯噪声上,标准差为25的时候可以达到最好为29.07,在泊松噪声上,λ为50时最好为29.45。

五.复现代码

这里使用Anaconda的jupyter notebook进行。

我的电脑参数配置

Nvidia GTX1660S

cuda11.3

torch-1.10.0+cu113-cp39-cp39-win_amd64(gpu版本)

torchvision-0.11.0+cu113-cp39-cp39-win_amd64(gpu版本)

我下面会按照代码块来输入代码

#选择是cpu还是gpu, 选择'cuda' 就是 GPU, 选择 'cpu' 就是 CPU
device = 'cuda'

 

import numpy as np 
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

 如果出现报错

ImportError: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html

  这个错误是由于缺少所需的IProgress模块引起的,它通常是由于Jupyter和ipywidgets版本不兼容引起的。解决这个问题的步骤如下:

#首先确保你的Jupyter和ipywidgets是最新版本。可以通过在终端中运行以下命令来更新它们:

pip install --upgrade jupyter ipywidgets

#安装ipywidgets扩展。运行以下命令

jupyter nbextension enable --py widgetsnbextension

#如果你使用的是JupyterLab,还需要安装它的扩展。运行以下命令

jupyter labextension install @jupyter-widgets/jupyterlab-manager

     重新启动Jupyter Notebook或JupyterLab。   

     如果没有出现报错就直接执行下面这个代码块

# Load test image from URL

import requests
from io import BytesIO

# 从Kodak24 dataset下载图片

#url = "https://drive.google.com/uc?export=download&id=18LcKoV4SYusF16wKqwNBJwYpTXE9myie"
#url = "https://drive.google.com/uc?export=download&id=176lM7ONjvyC83GcllCod-j1RPqjLoRoG"
#url = "https://drive.google.com/uc?export=download&id=1UIh9CwXSCf01JmAXgJo0LPtw5TUkWUU-"
url = "https://drive.google.com/uc?export=download&id=1j1OOzvGhet_GHJCaXbfisiW8uGDxI7ty"

response = requests.get(url) #获取文件
path=BytesIO(response.content) #将字节数据视为文件对象。它通常用于在内存中读写二进制数据,而不需要实际创建物理文件

clean_img = torch.load(path).unsqueeze(0) #加载图片并扩充维度
print(clean_img.shape) #B C H W
noise_type = 'gauss' # Either 'gauss' or 'poiss'
noise_level = 25     # Pixel range is 0-255 for Gaussian, and 0-1 for Poission

def add_noise(x,noise_level):#添加噪声
    
    if noise_type == 'gauss':#添加高斯噪声
        noisy = x + torch.normal(0, noise_level/255, x.shape)
        noisy = torch.clamp(noisy,0,1)
        
    elif noise_type == 'poiss':#添加泊松噪声
        noisy = torch.poisson(noise_level * x)/noise_level
    
    return noisy

noisy_img = add_noise(clean_img, noise_level) 
clean_img = clean_img.to(device)#把原始图片传入选定的设备GPU or CPU
noisy_img = noisy_img.to(device)
class network(nn.Module):#定义神经网络(两层)
    def __init__(self,n_chan,chan_embed=48):
        super(network, self).__init__()
        
        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.conv1 = nn.Conv2d(n_chan,chan_embed,3,padding=1)
        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding = 1)
        self.conv3 = nn.Conv2d(chan_embed, n_chan, 1)

    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        x = self.conv3(x)
        
        return x

n_chan = clean_img.shape[1] #通道数
model = network(n_chan)    #定义模型
model = model.to(device)   #在选定的设备上进行计算
print("The number of parameters of the network is: ",  sum(p.numel() for p in model.parameters() if p.requires_grad))  #输出模型参数
def pair_downsampler(img):
    #img has shape B C H W
    c = img.shape[1] #获取通道数
    
    filter1 = torch.FloatTensor([[[[0 ,0.5],[0.5, 0]]]]).to(img.device)
    #定义副对角线平均下采样算子
    
    filter1 = filter1.repeat(c,1, 1, 1)
    #这行代码的作用是复制 filter1 指定的次数    
    #其形状为 (c, 1, 2, 2)
    #c 表示复制的次数,也就是输入图像的通道数。
    #2 表示滤波器的高度为 2。
    #2 表示滤波器的宽度为 2。
    
    filter2 = torch.FloatTensor([[[[0.5 ,0],[0, 0.5]]]]).to(img.device)
    filter2 = filter2.repeat(c,1, 1, 1)
    
    output1 = F.conv2d(img, filter1, stride=2, groups=c)
    output2 = F.conv2d(img, filter2, stride=2, groups=c)
    
    return output1, output2 #返回D1(y),D2(y)
'''
展示原始图像y和其对应的下采样图像对,展示这个下采样算子是如何对图像空间分辨率进行下采样的 
'''
img1, img2 = pair_downsampler(noisy_img) #使用前面定义好的下采样算子进行下采样

img0 = noisy_img.cpu().squeeze(0).permute(1,2,0)#改变维度成H*W*C,并把数据移动cpu
img1 = img1.cpu().squeeze(0).permute(1,2,0)#同上
img2 = img2.cpu().squeeze(0).permute(1,2,0)

fig, ax = plt.subplots(1, 3,figsize=(15, 15)) #使用matplotlib绘图,定义一个一行三列图标

ax[0].imshow(img0)
ax[0].set_title('Noisy Img')#展示原始图像y

ax[1].imshow(img1)
ax[1].set_title('First downsampled')#展示下采样图D2(y)

ax[2].imshow(img2)
ax[2].set_title('Second downsampled')#展示下采样图D1(y)

 

 

def mse(gt: torch.Tensor, pred:torch.Tensor)-> torch.Tensor:#定义均方误差
    loss = torch.nn.MSELoss()
    return loss(gt,pred)

def loss_func(noisy_img):  #重新按照论文中的代码进行损失函数定义
    noisy1, noisy2 = pair_downsampler(noisy_img)

    pred1 =  noisy1 - model(noisy1)
    pred2 =  noisy2 - model(noisy2)
    
    loss_res = 1/2*(mse(noisy1,pred2)+mse(noisy2,pred1))
    
    noisy_denoised =  noisy_img - model(noisy_img)
    denoised1, denoised2 = pair_downsampler(noisy_denoised)
    
    loss_cons=1/2*(mse(pred1,denoised1) + mse(pred2,denoised2))
    
    loss = loss_res + loss_cons

    return loss

 

def train(model, optimizer, noisy_img):#定义训练函数
  
  loss = loss_func(noisy_img)#对噪声图像进行损失计算
  
  optimizer.zero_grad()  #将优化器中模型参数的梯度清零,准备进行反向传播
  loss.backward()        #对损失值进行反向传播,计算模型参数的梯度
  optimizer.step()       #根据计算得到的梯度,使用优化器更新模型的参数,使损失值最小化

  return loss.item()     #返回训练得到的损失值

def test(model, noisy_img, clean_img):#定义一个测试函数用于测试模型性能
    
    with torch.no_grad():
        #使用 torch.no_grad() 上下文管理器,表示在测试过程中不需要计算梯度,节省内存并加速计算
        pred = torch.clamp(noisy_img - model(noisy_img),0,1)
        #使用模型对输入的噪声图像 noisy_img 进行去噪,表示对噪声图像进行模型的前向传播得到去噪后的            
        #结果。torch.clamp 函数将输出限制在 [0, 1] 的范围内,确保像素值不会超出有效范围
        MSE = mse(clean_img, pred).item()
        #计算去噪后的结果 pred 与真实干净图像 clean_img 之间的均方误差(MSE)。mse 是一个函数用于 
        #计算均方误差
        PSNR = 10*np.log10(1/MSE)
        #峰值性噪比,PSNR 是一个表示图像质量的常用指标,用于衡量去噪效果。
    
    return PSNR

def denoise(model, noisy_img):#义一个去噪函数,用于对输入的噪声图像 noisy_img,就是y进行去噪
    
    with torch.no_grad():
        pred = torch.clamp( noisy_img - model(noisy_img),0,1)
        #使用模型对输入的噪声图像 noisy_img 进行去噪,得到去噪后的结果 pred。同样使用 torch.clamp 
        #函数将输出限制在 [0, 1] 的范围内,确保像素值不会超出有效范围
    
    return pred 
max_epoch = 2000     # 训练轮数
lr = 0.001           # 学习率
step_size = 1500     # 学习步长
gamma = 0.5          # 学习率衰减值

optimizer = optim.Adam(model.parameters(), lr=lr)  
'''
定义一个Adam自适应学习率优化算法
model.parameters(): model 是神经网络模型对象,parameters() 是模型的方法,它返回一个包含所有可学习参数的迭代器。优化器将根据这些参数来更新权重。

'''
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
'''
这行代码定义了一个学习率调度器(scheduler),用于在训练过程中自动调整优化器的学习率。
它将在每个 step_size 的倍数(即每经过 step_size 个epoch)时,将优化器 optimizer 的学习率按照 gamma 的比例进行衰减。通过学习率的调整,可以帮助训练过程更加稳定和高效。
'''
for epoch in tqdm(range(max_epoch)):
    train(model, optimizer, noisy_img)
    scheduler.step()
'''
进行训练
在每个训练周期结束后,调用 scheduler 学习率调度器的 step 方法。这将根据之前设置的 step_size 和 gamma 来更新优化器 optimizer 的学习率。每经过 step_size 个epoch,学习率将乘以 gamma 的值进行衰减。
'''
PSNR = test(model, noisy_img, clean_img)
print(PSNR)
#输出峰值性噪比
denoised_img = denoise(model, noisy_img)

denoised = denoised_img.cpu().squeeze(0).permute(1,2,0)
clean = clean_img.cpu().squeeze(0).permute(1,2,0)
noisy = noisy_img.cpu().squeeze(0).permute(1,2,0)

#用于展示原始图片,以及去噪图片等
fig, ax = plt.subplots(1, 3,figsize=(15, 15))

ax[0].imshow(clean)
ax[0].set_xticks([])
ax[0].set_yticks([])
ax[0].set_title('Ground Truth')

ax[1].imshow(noisy)
ax[1].set_xticks([])
ax[1].set_yticks([])
ax[1].set_title('Noisy Img')
noisy_psnr = 10*np.log10(1/mse(noisy_img,clean_img).item())
ax[1].set(xlabel= str(round(noisy_psnr,2)) + ' dB')

ax[2].imshow(denoised)
ax[2].set_xticks([])
ax[2].set_yticks([])
ax[2].set_title('Denoised Img')
ax[2].set(xlabel= str(round(PSNR,2)) + ' dB')
#绘图展示原始图片,加入噪声后的图片,以及去噪声图片

 完结!!!

 

 

 

  • 16
    点赞
  • 59
    收藏
    觉得还不错? 一键收藏
  • 16
    评论
评论 16
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值