IDR: Self-Supervised Image Denoising via Iterative Data Refinement

IDR 是一个无监督降噪模型。

1. noisy-clean pair 比较难获取

noisy-clean pair:
x: noidy image, y: clean image
但是 y比较难获取

noisr-noisy pair
x + n , x

对噪声图像再添加噪声，得到 噪声更大的图像。

这里的n表示的是sensor的噪声模型（也可以是采样得到的，参考作者另一篇论文rethinking noise）.

2. noiser-noisy pair 比较容易获取，但是训练效果呢？

作者的两个发现：

2.1 noiser-noisy 训练的模型，能够对 noisy 图像一定程度的降噪

如下图：

2.2 noiser-noisy数据 越接近 noisy-clean 数据，训练的效果越好。

3.通过训练 让noiser-noisy数据 更接近 noisy-clean 数据

1.训练F0，生成新的数据集

2.利用新的数据集训练F1.
由于 新的数据集 更接近 noisy-clean 数据，因此训练的结果对于noisy的表现会更好。

3.因此可以迭代训练，不断生成新的less biased数据集, 训练新的model

4.Fast Iterative Data Refinement

以上迭代训练需要生成多次数据集，训练多次model.

作者提出改进的方案：
a.每个epoch refine一次dataset, 不需要训练到完全收敛
b.利用上个epoch的model初始化下一个epoch的model
*
这样改进下来，和正常训练差别不大了，除了每个epoch要更新一次数据集。

实际的效果如下：
每次迭代，降噪效果都有改善。

5. sensenoise-500 dataset

IMX586, 3000x4000 pixels, low light conditions.

64 帧 = 4 帧 正常曝光noisy image + 60 帧 长曝光（1s-2s） use median value ad ground truth

正常曝光和长曝光的图像如何 保持亮度一致呢？需要设置 iso 和曝光时间：

图像示例：

最终图像数据是1010张（505pairs）:

dng是噪声图， npy是groundtruth

5.1 数据集 error

部分ground truth 高亮区域偏红色。

6.训练

4种方案训练sensenoise 500

1. pair
2. add GP noise
3. idr（本文）
4. noise2noise: add gp noise

由于不知道数据集的实际噪声参数。因此add noise都是添加的一定范围

k = np.random.uniform(0.8, 3)
scale = np.random.uniform(1, 30)
# k = torch.FloatTensor(k)
# scale = torch.FloatTensor(scale)
in_img1 = add_noise_torch(gt_img, k, scale).to(device)
in_img2 = add_noise_torch(gt_img, k, scale).to(device) # 是否需要转化为int16类型，因为实际raw图数据都是整数
gt_img = gt_img.to(device)
# print(in_img.min(), in_img.max(), in_img.mean(), in_img.var())
# print(gt_img.min(), gt_img.max(), gt_img.mean(), gt_img.var()) # more and more little
gt_img = gt_img / 1023
in_img = in_img1 / 1023
in_img = torch.clamp(in_img, 0, 1)
in_img2 = in_img2 / 1023
in_img2 = torch.clamp(in_img2, 0, 1)

idr训练：

import glob
import os.path

import cv2
import numpy as np
import rawpy
import torch
import torch.optim as optim
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from skimage.metrics import peak_signal_noise_ratio as compare_psnr
from skimage.metrics import structural_similarity as compare_ssim
from model import UNetSeeInDark
from sensenoise500 import add_noise_torch
from sid_dataset_sensenoise500 import sensenoise_dataset, apply_wb_ccm, sensenoise_dataset_2, \
    sensenoise_dataset_addnoise, sensenoise_dataset_addnoise_2, choose_k_sigma
import torchvision


if __name__ == "__main__":
    # 1.当前版本信息
    print(torch.__version__)
    print(torch.version.cuda)
    print(torch.backends.cudnn.version())
    print(torch.cuda.get_device_name(0))

    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed_all(0)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # 2. 设置device信息 和 创建model
    model = UNetSeeInDark()
    model._initialize_weights()

    gpus = [1]
    #model = nn.DataParallel(model, device_ids=gpus)
    device = torch.device('cuda:1')
    model = model.cuda(device=gpus[0])
    
    # 6. 是否恢复模型
    resume = 0
    last_epoch = 0
    lr_epoch = 1
    if resume and last_epoch > 1:
        model.load_state_dict(torch.load(os.path.join(save_model_dir, f'checkpoint_{last_epoch:04d}.pth'), map_location=device))
        lr_epoch = 0.5**(last_epoch // 500)
    
    # 3. dataset 和 data loader, num_workers设置线程数目，pin_memory设置固定内存
    # train_dataset = sensenoise_dataset_addnoise_2(mode='train')
    # train_dataset_loader = DataLoader(train_dataset, batch_size=4*len(gpus), shuffle=True, num_workers=8, pin_memory=True)

    eval_dataset = sensenoise_dataset_2(mode='eval')
    eval_dataset_loader = DataLoader(eval_dataset, batch_size=1, num_workers=8, pin_memory=True)
    print('load dataset !')
	files = glob.glob(os.path.join('/home/wangzhansheng/dataset/sidd/SenseNoise500/final_datasetv3/', '*.dng'))
    files = sorted(files)[:400]
    datas = []
    for file in files:
        input_path = file

        txt_path = input_path[:-4] + '.txt'
        para = np.loadtxt(txt_path)
        wb_gain = np.array(para[:3]).astype(np.float32)
        ccm = np.array(para[3:12]).astype(np.float32).reshape(3, 3)
        iso = para[-1]
        # gt_raw = np.load(gt_path).astype(np.int32)
        # gt_raw = np.dstack((gt_raw[0::2, 0::2], gt_raw[0::2, 1::2], gt_raw[1::2, 0::2], gt_raw[1::2, 1::2]))

        input_raw = rawpy.imread(input_path).raw_image_visible.astype(np.float32)
        input_raw = np.dstack((input_raw[0::2, 0::2], input_raw[0::2, 1::2], input_raw[1::2, 0::2], input_raw[1::2, 1::2]))

        datas.append([input_raw, wb_gain, ccm, input_path, iso])
        print(file, len(datas))

    # 4. 损失函数 和  优化器
    loss_fn = nn.L1Loss()
    learning_rate = 3*1e-4
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    lr_step = 500
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, lr_step, gamma=0.5)

    # 5. hyper para 设置
    epochs = 5000
    save_epoch = 100
    save_model_dir = 'saved_model_sensenoise500_addnoise_single_idr'

    eval_epoch = 100
    save_sample_dir = 'saved_sample_sensenoise500_addnoise_single_idr'
    if not os.path.exists(save_model_dir):
        os.makedirs(save_model_dir)


    # 7. 训练epoch
    epoch_infos = []
    eval_infos = []
    
    
    patch_size = 512
    for epoch in range(last_epoch+1, epochs + 1):
        print('current epoch:', epoch, 'current lr:', optimizer.state_dict()['param_groups'][0]['lr'])
        if epoch < 101:
            save_epoch = 10
            eval_epoch = 10
        else:
            save_epoch = 100
            eval_epoch = 100
        # 8. train loop
        model_copy = UNetSeeInDark().to(device)
        model_copy.load_state_dict(model.state_dict())
        model_copy.eval()


        model.train()
        g_loss = []
        g_psnr = []

        kk = 0
        for idx in tqdm(np.random.permutation(len(datas))):
            data = datas[idx]
        #for data in np.random.shuffle(datas):
            # gt_path = file
            # txt_path = gt_path[:-4] + '.txt'
            # para = np.loadtxt(txt_path)
            # wb_gain = np.array(para[:3]).astype(np.float32)
            # ccm = np.array(para[3:12]).astype(np.float32).reshape(3, 3)
            # iso = para[-1]
            #
            # gt_raw = np.load(gt_path).astype(np.int32)
            # iso, k, sigma = choose_k_sigma(iso/2)
            # k = k * np.random.uniform(0.8, 1.2)
            # sigma2 = np.sqrt(sigma) * np.random.uniform(0.8, 1.1)
            # short_raw = k * np.random.poisson(gt_raw / k) + np.random.normal(0., sigma2, gt_raw.shape)
            # gt_raw = gt_raw / 1023
            # short_raw = short_raw / 1023

            input_raw, wb_gain, ccm, gt_path, iso = data
            # crop
            h, w, c = input_raw.shape
            h1 = np.random.randint(0, h - patch_size)
            w1 = np.random.randint(0, w - patch_size)
            # short_raw = short_raw[h1:h1 + patch_size, w1:w1 + patch_size, :]
            short_raw = input_raw[h1:h1 + patch_size, w1:w1 + patch_size, :]

            # augment
            if np.random.randint(2, size=1)[0] == 1:  # random flip
                short_raw = np.flip(short_raw, axis=0)
                #gt_raw = np.flip(gt_raw, axis=0)
            if np.random.randint(2, size=1)[0] == 1:
                short_raw = np.flip(short_raw, axis=1)
                #gt_raw = np.flip(gt_raw, axis=1)
            if np.random.randint(2, size=1)[0] == 1:  # random transpose
                short_raw = np.transpose(short_raw, (1, 0, 2))
                #gt_raw = np.transpose(gt_raw, (1, 0, 2))
            #in_img = torch.permute(input_patch, (0,3,1,2)).cuda(device=gpus[0])

            short_raw = np.ascontiguousarray(short_raw[np.newaxis, ...])
            gt_img = torch.from_numpy(short_raw).permute(0, 3, 1, 2)


            if epoch > last_epoch + 1:
                model_copy.eval()
                with torch.no_grad():

                    gt_img_last = gt_img.to(device) / 1023
                    gt_img = model_copy(gt_img_last).cpu()
                    gt_img = torch.clamp(gt_img* 1023, 0, 1023)
                    # print(gt_img_last.min(), gt_img_last.max(), gt_img_last.mean(), gt_img_last.var())
                    # print(gt_img.min(), gt_img.max(), gt_img.mean(), gt_img.var()) # more and more little
                    if kk  > 50000:
                        im1 = gt_img_last.cpu().float().numpy().squeeze().transpose(1, 2, 0)
                        im2 = gt_img.float().numpy().squeeze().transpose(1, 2, 0) / 1023
                        im1 = im1[..., [0, 1, 3]] ** (1 / 2.2)
                        im2 = im2[..., [0, 1, 3]] ** (1 / 2.2)
                        im1 = np.clip(im1 * 255 + 0.5, 0, 255).astype(np.uint8)
                        im2 = np.clip(im2 * 255 + 0.5, 0, 255).astype(np.uint8)
                        save_sample_dir3 = save_sample_dir + f'/{epoch:04}dd/'
                        if not os.path.isdir(save_sample_dir3):
                            os.makedirs(save_sample_dir3)

                        filename_save = os.path.basename(gt_path)[:-4]
                        cv2.imwrite(os.path.join(save_sample_dir3, '%s_dddd1.png' % (filename_save)), im1[..., ::-1])
                        cv2.imwrite(os.path.join(save_sample_dir3, '%s_dddd2.png' % (filename_save)),
                                    im2[..., ::-1])
            iso, k, sigma = choose_k_sigma(iso/2)

            # k = np.random.uniform(0.8, 3)
            # scale = np.random.uniform(1, 30)
            # k = torch.FloatTensor(k)
            # scale = torch.FloatTensor(scale)
            scale = np.sqrt(sigma)
            in_img = add_noise_torch(gt_img, k, scale).to(device)
            gt_img = gt_img.to(device)
            # print(in_img.min(), in_img.max(), in_img.mean(), in_img.var())
            # print(gt_img.min(), gt_img.max(), gt_img.mean(), gt_img.var()) # more and more little
            gt_img = gt_img / 1023
            in_img = in_img / 1023
            in_img = torch.clamp(in_img, 0, 1)

            # print(gt_img.shape, gt_img.min(), gt_img.max())
            # print(in_img.shape, in_img.min(), in_img.max())
            # print(wb_gain, ccm, iso, gt_path)

            out = model(in_img)
            loss = loss_fn(out, gt_img)

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

            # training result
            g_loss.append(loss.data.detach().cpu())
            mse_value = np.mean((out.cpu().data.numpy() - gt_img.cpu().data.numpy()) ** 2)
            psnr = 10. * np.log10(1. / mse_value)
            g_psnr.append(psnr)

        mean_loss = np.mean(np.array(g_loss))
        mean_psnr = np.mean(np.array(g_psnr))
        print(f'epoch{epoch:04d} ,train loss: {mean_loss},train psnr: {mean_psnr}')
        epoch_infos.append([epoch, mean_loss, mean_psnr])

        # 9. save model
        if epoch % save_epoch == 0:
            save_model_path = os.path.join(save_model_dir, f'checkpoint_{epoch:04d}.pth')
            torch.save(model.state_dict(), save_model_path)
        # 10. eval test and save some samples if needed
        if epoch % eval_epoch == 0:
            model.eval()
            k = 0

            with torch.no_grad():
                psnr_12800_0 = []
                psnr_12800_1 = []

                ssim_12800_0 = []
                ssim_12800_1 = []

                for data in tqdm(eval_dataset_loader):
                    input_patch, gt_patch, wb_gain, ccm, gt_path, iso = data

                    in_img = input_patch.permute(0, 3, 1, 2).cuda(device=gpus[0])
                    gt_img = gt_patch.permute(0, 3, 1, 2).cuda(device=gpus[0])

                    out = model(in_img)

                    im1 = gt_img.detach().cpu().float().numpy().squeeze().transpose(1,2,0)
                    im2 = out.detach().cpu().float().numpy().squeeze().transpose(1,2,0)
                    im1 = np.clip(im1 * 255 + 0.5, 0, 255).astype(np.uint8)
                    im2 = np.clip(im2 * 255 + 0.5, 0, 255).astype(np.uint8)

                    temp_psnr = compare_psnr(im1, im2, data_range=255)
                    temp_ssim = compare_ssim(im1, im2, data_range=255, channel_axis=-1)
                    if iso <= 12800:
                        psnr_12800_0.append(temp_psnr)
                        ssim_12800_0.append(temp_ssim)
                    else:
                        psnr_12800_1.append(temp_psnr)
                        ssim_12800_1.append(temp_ssim)

                    # show training out
                    save_img = 1
                    if save_img and k<10:
                        k += 1
                        im_input = in_img.detach().permute(0, 2, 3, 1).cpu().float().numpy()[0]
                        im_gt = gt_img.detach().permute(0, 2, 3, 1).cpu().float().numpy()[0]
                        im_out = out.detach().permute(0, 2, 3, 1).cpu().float().numpy()[0]
                        wb_gain = wb_gain.data.cpu().numpy()[0]
                        ccm = ccm.data.cpu().numpy()[0]
                        gt_path = gt_path[0]

                        pattern_sensenoise500 = 'RGGB'
                        im_input_srgb = apply_wb_ccm(im_input[..., [0, 1, 3]], wb_gain, ccm, pattern_sensenoise500)
                        im_gt_srgb = apply_wb_ccm(im_gt[..., [0, 1, 3]],  wb_gain, ccm, pattern_sensenoise500)
                        im_out_srgb = apply_wb_ccm(im_out[..., [0, 1, 3]],  wb_gain, ccm, pattern_sensenoise500)

                        im_input_srgb = np.clip(im_input_srgb * 255 + 0.5, 0, 255).astype(np.uint8)
                        im_gt_srgb = np.clip(im_gt_srgb * 255 + 0.5, 0, 255).astype(np.uint8)
                        im_out_srgb = np.clip(im_out_srgb * 255 + 0.5, 0, 255).astype(np.uint8)

                        save_sample_dir2 = save_sample_dir + f'/{epoch:04}/'
                        if not os.path.isdir(save_sample_dir2):
                            os.makedirs(save_sample_dir2)

                        # save_sample_path = os.path.join(save_sample_dir2, os.path.basename(gt_path)[:-4]+'.png')
                        # cv2.imwrite(save_sample_path, np.hstack((im_gt_srgb,im_input_srgb, im_out_srgb))[..., ::-1])

                        filename_save = os.path.basename(gt_path)[:-4]
                        cv2.imwrite(os.path.join(save_sample_dir2, '%s_psnr_%.2f_out.png' % (filename_save, temp_psnr)), im_out_srgb[...,::-1])
                        cv2.imwrite(os.path.join(save_sample_dir2, '%s_NOISY.png' % (filename_save)), im_input_srgb[...,::-1])
                        cv2.imwrite(os.path.join(save_sample_dir2, '%s_GT.png' % (filename_save)), im_gt_srgb[...,::-1])

                print('eval dataset  psnr: ', np.array(psnr_12800_0).mean(), np.array(psnr_12800_1).mean())
                print('eval dataset  ssim: ', np.array(ssim_12800_0).mean(), np.array(ssim_12800_1).mean())
                eval_infos.append([epoch, np.array(psnr_12800_0).mean(), np.array(psnr_12800_1).mean(), np.array(ssim_12800_0).mean(), np.array(ssim_12800_1).mean()])
        scheduler.step() # 更新学习率

    np.savetxt('train_infos.txt',  epoch_infos, fmt='%.4f') # epoch loss psnr
    np.savetxt('eval_infos.txt', eval_infos, fmt='%.4f')    # epoch psnr, ssim