【论文复现】Noise2Noise图像去噪

Eternity._

发布于 2024-12-15 08:59:37

1.6K0

概述

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

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

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

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

算法原理

对于干净的图片分布，我们用

x_{clean}

表示。在现实的生产生活中,干净的图片经常会被噪声所污染，变为

x_{noise}

，污染的过程我们可以使用如下过程表示：

其中noise最常见的分布便是高斯分布与泊松分布。在深度学习中，图像的去噪任务就是尽可能学习到一个模型使得:

其中f为模型，θ为训练过程中学习到的模型参数。当然在现实中是不可能的，我们只能使我们的模型尽可能去逼近理想效果，即:

为了达到以上效果,有监督学习采取的方法是是最小化训练数据与标签的损失函数来学习模型参数即:

其中L代表我们选取的损失函数,在图像训练中通常为

L _2

损失或

L_1

损失。具体到某个数据集上，有监督学习的过程实际上就是如下过程：

其中i代表数据集中的样本标号。

演示效果

核心逻辑

#获得噪声对
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路径下即可