传知代码-Noise2Noise图像去噪（论文复现）

代码以及视频讲解

本文所涉及所有资源均在传知代码平台可获取

概述

本文复现论文Noise2Noise: Learning Image Restoration without Clean Data中提出的图像去噪方法，并使用最广泛应用的高斯噪声与泊松噪声来做实验。原文连接https://arxiv.org/abs/1803.04189

该论文证明了，对于同一张干净图片，如果分两次污染它所用的噪声同分布且0均值的情况下，那么用这一对噪声图像进行网络训练即noise2noise的训练方法就能达到与用干净-噪声图像对即noise2clean的训练方法接近的去噪效果。

该方法在无监督与自监督方法中效果几乎达到了最好，相比ZS-Noise2Noise这样的方法效果要优秀非常多，与有监督方法的效果相近！

论文在多种噪声和测试数据集上进行了实验，本文只选取高斯噪声与泊松噪声并在官方测试数据集中选取了一个进行测试。

算法原理

对于干净的图片分布，我们用$x_{clean}$表示。在现实的生产生活中,干净的图片经常会被噪声所污染，变为$x_{noise}$，污染的过程我们可以使用如下过程表示：
$$x_{noise}=x_{clean}+noise$$
其中$noise$最常见的分布便是高斯分布与泊松分布。
在深度学习中，图像的去噪任务就是尽可能学习到一个模型使得:
$$x_{clean}=f_{\theta }(x_{noise})$$
其中$f$为模型，$\theta$为训练过程中学习到的模型参数。当然在现实中是不可能的，我们只能使我们的模型尽可能去逼近理想效果，即:
$$x_{clean}\approx f_{\theta }(x_{noise})$$
为了达到以上效果,有监督学习采取的方法是是最小化训练数据与标签的损失函数来学习模型参数即:
argmin ⁡ θ E ( x n o i s e , x c l e a n ) { L ( f θ ( x n o i s e ) , x c l e a n ) } \underset{\theta}{\operatorname{argmin}} \mathbb{E}_{(x_{noise}, x_{clean})}\left\{L\left(f_{\theta}(x_{noise}),x_{clean}\right)\right\} θargmin​E(xnoise​,xclean​)​{L(fθ​(xnoise​),xclean​)}
其中$L$代表我们选取的损失函数,在图像训练中通常为$L_2$损失或$L_1$损失。

具体到某个数据集上，有监督学习的过程实际上就是如下过程：
$$\underset{\theta }{\mathrm{argmin}}\sum _{i}L\left(f_{\theta }\left(x_{noise}^i\right),x_{clean}^i\right)$$
其中i代表数据集中的样本标号。

当然，有监督学习的效果毋庸置疑，但是在现实生活中($x_{noise}^i$,$x_{clean}^i$)这样的图像对获取并不容易，这也成为了制约有监督学习的一大因素。

本文算法的亮点就是不用传统有监督学习算法中使用的($x_{noise}^i$,$x_{clean}^i$)图像对，而是改用($x_{noise1}^i$,$x_{noise2}^i$)这样的图像对进行训练却能达到有监督方法的效果，该方法将有监督学习训练模型的过程改为了：
$$\underset{\theta }{\mathrm{argmin}}\sum _{i}L\left(f_{\theta }\left(x_{noise1}^i\right),x_{noise2}^i\right)$$
其中$noise1$与$noise2$代表独立同分布且都是0均值的噪声。$x_{noise1}$代表用$noise1$污染$x_{clean}$后的图片，$x_{noise2}$代表用$noise2$污染$x_{clean}$后的图片，在本文算法中$noise1$于$noise2$可选择高斯分布与泊松分布两种方案。

尽管对于$noise1$与$noise2$噪声分布的限制略显苛刻，但本文算法无疑为解决有监督学习的去噪任务中标签数据的难以获取问题提供了一种突破性的思路。

在模拟数据的实验结果上证明Noise2Noise的训练方法性能非常接近有监督的训练方法。

演示效果

核心逻辑

#获得噪声对
def get_noisePair(image, noise_type='gaussian', mode=None):
    # Convert tensor to numpy array
    image_np = image.numpy()

    # Determine noise parameters based on mode
    if mode == 'train':
        if noise_type == 'gaussian':
            std = np.random.randint(0, 51)  # Random standard deviation between 0 and 50
        elif noise_type == 'poisson':
            lam = np.random.randint(0, 51)  # Random λ between 0 and 50
    else:
        std = 25  # Default standard deviation for Gaussian noise
        lam = 30  # Default λ for Poisson noise

    # Add noise based on noise_type
    if noise_type == 'gaussian':
        noise = np.random.normal(0, std, size=image_np.shape)
        source = image_np + noise

        noise = np.random.normal(0, std, size=image_np.shape)
        target = image_np + noise
    elif noise_type == 'poisson':
        noise = np.random.poisson(lam, size=image_np.shape)
        source = image_np + noise

        noise = np.random.poisson(lam, size=image_np.shape)
        target = image_np + noise

    # Clip noisy image values to ensure they are within [0, 255] range
    source = np.clip(source, 0, 255)

    # Convert back to tensor
    source = torch.tensor(source, dtype=torch.float32)

    target = np.clip(target, 0, 255)

    target = torch.tensor(target, dtype=torch.float32)

    return source, target

#用噪声对训练以获取去噪模型
def train(train_loader, model, optimizer, scheduler, writer, epoch):

    batch_time = AverageMeter()
    loss_img = AverageMeter()
    model.train()
    end = time.time()

    step = 0

    for data in train_loader:

        source=data["source"]#噪声图像1
        source=source.cuda()

        target=data["target"]#噪声图像2
        target=target.cuda()

        pred=model(source)

        loss=mse(pred,target)
        # loss = 0.01 * loss1 + loss3

        loss_img.update(loss.item(), target.size(0))
        batch_time.update(time.time() - end)
        end = time.time()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        step += 1

    writer.add_scalars('train_loss', {'loss_img': loss_img.avg}, epoch + 1)
    scheduler.step()

    print('Train Epoch: {}\t train_loss: {:.6f}\t'.format(epoch + 1, loss_img.avg))

使用方式

系统及基础镜像

本代码于Ubuntu系统上编写与测试，在windows11系统上运行需要自行配置相应的版本的pytorch
基础镜像:python3.8.13

配置环境

进入noise2noise-pytorch-myown项目路径
运行以下命令：

pip install -r requirements.txt

不训练仅用已经训好的模型查看效果

可直接运行以下命令

python example.py

即可看到视频中演示的去噪效果

训练集及验证集获取

简便获取

百度云盘链接（下载后可在ReadMe文件同位置获得）
以上链接中存储了训练、验证、测试所需要的数据，下载后直接解压到项目根目录下，如下图所示即可

完整获取过程(想用全部数据集进行训练时推荐)

http://image-net.org/download
需认证资格并自行下载imagenet数据集，本次复现使用的是ILSVRC2012的验证集。
下载后解压到data路径下，解压完成后路径格式如下所示：

data/ILSVRC2012_img_val/...

…代表图片，格式为JPEG

为了简便训练，我们只用其中的5000张做训练集，1000张做验证集即可。完成解压后，运行下列命令即可获得训练集与验证集，分别存储在train与valid路径下

mkdir train && mkdir valid
find data/ILSVRC2012_img_val/ -maxdepth 1 -type f | head -n 5000 | xargs -I {} mv {} ./train/
find data/ILSVRC2012_img_val/ -maxdepth 1 -type f | tail -n +5001 | head -n 1000 | xargs -I {} mv {} ./valid/

当然，你也可以自行调整训练集的大小，甚至使用整个数据集作为训练集以更加逼近原论文的结果

测试集获取

运行下列命令可直接获取BSD300数据集，并将其测试集放在我们项目的test路径下

wget https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/segbench/BSDS300-images.tgz
tar -zxvf BSDS300-images.tgz
mkdir test
mv BSDS300/images/test/* ./test/

如果您想在windows系统上使用，或是命令的方式觉得不太清楚，以上工作均可以手动完成，只要最终的结果保证训练集在./train路径下，验证集在./valid路径下，测试集在./test路径下即可

训练

如果您想自行训练模型

用高斯噪声训练

直接运行以下命令：

python train.py

用泊松噪声训练

需要修改train.py中下列的参数，将gaussian改为poisson即可

然后再运行如下命令：

python train.py

##测试
根据你要测试的模型，修改test.py中如下参数

然后根据你想要测试的高斯或泊松模型选择在runs/Noise2Noisegaussian/checkpoints或runs/Noise2Noisepoisson/checkpoints路径下存储的你的训练模型路径
将test.py下列的参数

替换为你的模型路径
最后，运行如下命令

python test.py

即可得到测试结果，且会将你的测试集的第一张图的训练与输出结果放到result路径下，如下图所示

源码下载