实验——基于pytorch的噪声估计网络

给出DNCNN的代码：https://github.com/cszn/DnCNN/tree/master/TrainingCodes/dncnn_pytorch（里面有给出pytorch实现）

代码的框架基于xintao前辈的框架（https://github.com/xinntao/BasicSR）

先修改代码框架，进入将train_sr.json改为train_ne.json

给出该文件的setting：

setting：

{
  "name": "noise_estimation_gt" //"001_RRDB_PSNR_x4_DIV2K" //  please remove "debug_" during training or tensorboard wounld not work
  ,
  "use_tb_logger": true,
  "model": "sr",
  "crop_scale": 0
    , "scale": 1//it must be 1
  , "gpu_ids": [5]

  , "datasets": {
    "train": {
      "name": "DIV2K800",
      "mode": "LRHR" //it must be this, and the detail would be shown in LRHR_dataset.py
      , "noise_get": true///
      , "dataroot_HR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub"///must be sub
      , "dataroot_LR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub_noise"
      , "subset_file": null
      , "use_shuffle": true
      , "n_workers": 8
      , "batch_size": 24//how many samples in each iters
      , "HR_size": 128 // 128 | 192
      , "use_flip": false//true//
      , "use_rot": false//true
    }
    , "val": {
      "name": "val_set5",
      "mode": "LRHR",
      "dataroot_HR": "/home/guanwp/BasicSR_datasets/val_set5/Set5",
      "dataroot_LR": "/home/guanwp/BasicSR_datasets/val_set5/Set5_noise"
       , "noise_get": true///this is important
    }
  }

  , "path": {
    "root": "/home/guanwp/BasicSR-master/",
    "pretrain_model_G": null
     ,"experiments_root": "/home/guanwp/BasicSR-master/experiments/",
    "models": "/home/guanwp/BasicSR-master/experiments/noise_estimation_gt/models",
    "log": "/home/guanwp/BasicSR-master/experiments/noise_estimation_gt",
    "val_images": "/home/guanwp/BasicSR-master/experiments/noise_estimation_gt/val_images"
  }

  , "network_G": {
    "which_model_G": "noise_estimation"//"espcn"//"srresnet"//"sr_resnet"//"fsrcnn"//"sr_resnet" // RRDB_net | sr_resnet
    , "norm_type": null
    , "mode": "CNA"
    , "nf": 64//56//64
    , "nb": 23
    , "in_nc": 3
    , "out_nc": 3
    , "gc": 32
    , "group": 1
  }

  , "train": {
    "lr_G": 1e-4//1e-3//2e-4
    , "lr_scheme": "MultiStepLR"
    , "lr_steps": [200000, 300000, 400000,500000]
    , "lr_gamma": 0.5

    , "pixel_criterion": "l2"//"l2_tv"//"l1"//'l2'//huber//Cross   //should be MSE LOSS
    , "pixel_weight": 1.0
    , "val_freq": 1e3

    , "manual_seed": 0
    , "niter": 6e5//2e6//1e6
  }

  , "logger": {
    "print_freq": 200
    , "save_checkpoint_freq": 1e3
  }
}

 

Loss

其中关于TVloss（对图像的梯度求平方和）

在GitHub上也有TVLoss的实现https://github.com/jxgu1016/Total_Variation_Loss.pytorch/blob/master/TVLoss.py

给出TVloss的代码

#########################################
#TV loss(total variation regularizer)
class TVLoss(nn.Module):
    def __init__(self,TVLoss_weight=1):
        super(TVLoss,self).__init__()
        self.TVLoss_weight = TVLoss_weight

    def forward(self,x):
        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h = self._tensor_size(x[:,:,1:,:])
        count_w = self._tensor_size(x[:,:,:,1:])
        h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum()
        w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum()
        return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size

    def _tensor_size(self,t):
        return t.size()[1]*t.size()[2]*t.size()[3]


##########################################

注意，使用TVLOSS前应该先对其进行初始化，如

关于total variation

The total variation (TV) loss encourages spatial smoothness in the generated image.（总变差（TV）损失促进了生成的图像中的空间平滑性。）

Rubin等人在1990年左右观察到受噪声污染的图像的TV比无噪图像的总变分明显的大。 那么最小化TV理论上就可以最小化噪声。图片中相邻像素值的差异可以通过降低TV loss来一定程度上解决。比如降噪，对抗checkerboard等等。

wiki上关于TVLoss的描述https://en.wikipedia.org/wiki/Total_variation

在图像复原过程中，图像上的一点点噪声可能就会对复原的结果产生非常大的影响，因为很多复原算法都会放大噪声。这时候我们就需要在最优化问题的模型中添加一些正则项来保持图像的光滑性，TV是常用的一种正则项（https://www.zhihu.com/question/47162419/answer/251800162）。

另外参考的TVloss如下

# TV loss
class TVLoss(nn.Module):
    def __init__(self):
        super(TVLoss, self).__init__()

    def __call__(self, input):
        im_size = input.size()
        batchSize = input.size()[0]
        x_tv = torch.pow((input[:, :, 1:, :] - input[:, :, :im_size[2] - 1, :]), 2).sum()
        y_tv = torch.pow((input[:, :, :, 1:] - input[:, :, :, :im_size[3] - 1]), 2).sum()
        tv_x_size = self._tensor_size(input[:, :, 1:, :])
        tv_y_size = self._tensor_size(input[:, :, :, 1:])
        loss_tv = (x_tv / tv_x_size + y_tv / tv_y_size) / batchSize
        return loss_tv

    def _tensor_size(self, t):
        return t.size()[1] * t.size()[2] * t.size()[3]

TVLOSS一般是用于fine-tune，用于tradeoff

 

Data Prepare

heteroscedastic Gaussian（异方差高斯）

参考论文《Benchmarking Denoising Algorithms with Real Photographs》有：（改论文的代码https://github.com/lbasek/image-denoising-benchmark）

一般添加噪声都是采用adding noise to a latent noise-free image，但此处提到：

The main idea there is to model the noise distribution as a heteroscedastic Gaussian, whose variance is intensity-dependent.

给出代码

代码是从（https://github.com/GuoShi28/CBDNet）中修改的

function Test_Realistic_Noise_Model()
%%% Test Code for realistic noise model
addpath('./utils');   %%%%%%%%%%%%%this is the function

%% load CRF parameters, the pipeline of the camera, is should be turn off
load('201_CRF_data.mat');
load('dorfCurvesInv.mat');
I_gl = I;
B_gl = B;
I_inv_gl = invI;
B_inv_gl = invB;
mod_scale = 1;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% set parameters
% comment the unnecessary line
input_folder = '/home/guanwp/BasicSR_datasets/DIV2K_train_HR';
% save_mod_folder = '';
save_LR_folder = '/home/guanwp/BasicSR_datasets/DIV2K_train_noise';
% save_bic_folder = '';
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if exist('save_mod_folder', 'var')
    if exist(save_mod_folder, 'dir')
        disp(['It will cover ', save_mod_folder]);
    else
        mkdir(save_mod_folder);
    end
end
if exist('save_LR_folder', 'var')
    if exist(save_LR_folder, 'dir')
        disp(['It will cover ', save_LR_folder]);
    else
        mkdir(save_LR_folder);
    end
end
if exist('save_bic_folder', 'var')
    if exist(save_bic_folder, 'dir')
        disp(['It will cover ', save_bic_folder]);
    else
        mkdir(save_bic_folder);
    end
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

idx = 0;
filepaths = dir(fullfile(input_folder,'*.*'));
for i = 1 : length(filepaths)
    [paths,imname,ext] = fileparts(filepaths(i).name);
    if isempty(imname)
        disp('Ignore . folder.');
    elseif strcmp(imname, '.')
        disp('Ignore .. folder.');
    else
        idx = idx + 1;
        str_rlt = sprintf('%d\t%s.\n', idx, imname);
        fprintf(str_rlt);
        % read image
        img = imread(fullfile(input_folder, [imname, ext]));
        img = im2double(img);
        % modcrop
        %%%img = modcrop(img, mod_scale);
        if exist('save_mod_folder', 'var')
            imwrite(img, fullfile(save_mod_folder, [imname, '.png']));
        end
        % LR%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%add noise
        %%%sigma_s = [0.08 0.08 0.08]; % recommend 0~0.16
        %%%%sigma_c = [0.03 0.03 0.03]; % recommend 0~0.06
        %%%%%%%%%%random
        channel = size(img,3);
        sigma_s = 0.16*rand(1,channel,'single'); % original 0.16
        sigma_c = 0.06*rand(1,channel,'single'); % original 0.06
        %%%used the default value
        CRF_index = 5;  % 1~201
        pattern = 5;    % 1: 'gbrg', 2: 'grbg', 3: 'bggr', 4: 'rggb', 5: no mosaic

        im_LR = AddNoiseMosai(img,I_gl,B_gl,I_inv_gl,B_inv_gl, sigma_s,sigma_c, CRF_index, pattern);
        %%%%%%%%%%%%%%%%%%im_LR = AddNoiseMosai(img,I_gl,B_gl,I_inv_gl,B_inv_gl, CRF_index, pattern);

        %% JPEG compression
        %%% If JPEG compression is not considered, just commented out the following 
        %%% Code or just set "quality" equal to 100
        qality = 100; % image quality, recommend 60~100

        if exist('save_LR_folder', 'var')
            %%%%%%%%%%%%%%%%%%imwrite(im_LR, fullfile(save_LR_folder, [imname, '_bicLRx4.png']),'Quality', qality);
            imwrite(im_LR, fullfile(save_LR_folder, [imname, '_bicLRx4.png']));
        end
        % Bicubic
        if exist('save_bic_folder', 'var')
            im_B = imresize(im_LR, up_scale, 'bicubic');
            imwrite(im_B, fullfile(save_bic_folder, [imname, '_bicx4.png']));
        end
    end
end
end


%% modcrop
function img = modcrop(img, modulo)
if size(img,3) == 1
    sz = size(img);
    sz = sz - mod(sz, modulo);
    img = img(1:sz(1), 1:sz(2));
else
    tmpsz = size(img);
    sz = tmpsz(1:2);
    sz = sz - mod(sz, modulo);
    img = img(1:sz(1), 1:sz(2),:);
end
end

%//
%                             AddNoiseMosai
%function: add realistic noise 
% If this Code is helpful to you, please Cite: https://arxiv.org/abs/1807.04686
%%% realistic noise model:
%%% I = M^{-1}(M(f(L + n_s + n_c) + n_q))
% n_s: shot noise, depend on L, E[n_s]= 0, var[n_s] = L*sigma_s
%      in [1], n_s is a Possion Shot
%      in [2], n_s is GMM, sigma_s: 0~0.16
%      we choose [2] here
% n_c: other type of noise, i.e. read noise, dark current
%      in [2], E[n_c] = 0, var[n_c] = sigma_c, sigma_c: 0.01~0.06
% n_q: ignore in [2]
%//
%inpus:
%-------x: image data, double or single datatype, [0,1]
%-------I, B: CRF parameters provided by [3]
%-------Iinv, Binv: inverse CRF parameters, created by "inverseCRF" for
%                   faster computation
%**(optional): if not defined by user,the following parameters are set
%              ramdomly
%-------sigma_s: signal-dependent noise level, [0, 0.16]
%-------sigma_c: sigma_independent noise level, [0, 0.06]
%-------crf_index: CRF index, [1, 201]
%-------pattern: 1: 'gbrg', 2: 'grbg', 3: 'bggr', 4: 'rggb', 5: no mosaic
%
%output:
%-------y: noisy image
%//
% reference:
% [1] G.E. Healey and R. Kondepudy, Radiometric CCD Camera Calibration and Noise Estimation,
%     IEEE Trans. Pattern Analysis and Machine Intelligence
% [2] Liu, Ce et al. Automatic Estimation and Removal of Noise from a Single Image. 
%     IEEE Transactions on Pattern Analysis and Machine Intelligence 30 (2008): 299-314.
% [3] Grossberg, M.D., Nayar, S.K.: Modeling the space of camera response functions.
%     IEEE Transactions on Pattern Analysis and Machine Intelligence 26 (2004)
%

function [y] = AddNoiseMosai(x,I,B,Iinv,Binv,sigma_s,sigma_c,crf_index, pattern)
% default value
channel = size(x,3);
rng('default')
if nargin < 6
    rng('shuffle');
    sigma_s = 0.16*rand(1,channel,'single'); % original 0.16
    rng('shuffle');
    sigma_c = 0.06*rand(1,channel,'single'); % original 0.06
    rng('shuffle');
    rand_index = randperm(201);
    crf_index = rand_index(1);
    rng('shuffle');
    pattern = randperm(5);
end  
temp_x = x;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% x -> L
%%%%%%temp_x = ICRF_Map(temp_x,Iinv(crf_index,:),Binv(crf_index,:)); 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%noise
noise_s_map = bsxfun(@times,permute(sigma_s,[3 1 2]),temp_x);
noise_s = randn(size(temp_x),'single').* noise_s_map;
temp_x = temp_x + noise_s;

noise_c_map = repmat(permute(sigma_c,[3 1 2]),[size(x,1),size(x,2)]);
noise_c = noise_c_map .* randn(size(temp_x),'single');
temp_x = temp_x + noise_c;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% add Mosai
if pattern(1) == 1
    [B_b, ~, ~] = mosaic_bayer(temp_x, 'gbrg', 0);
elseif pattern(1) == 2
    [B_b, ~, ~] = mosaic_bayer(temp_x, 'grbg', 0);
elseif pattern(1) == 3
    [B_b, ~, ~] = mosaic_bayer(temp_x, 'bggr', 0);
elseif pattern(1) == 4
    [B_b, ~, ~] = mosaic_bayer(temp_x, 'rggb', 0);
else
    B_b = temp_x;
end
temp_x = single(B_b);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% DeMosai
if pattern(1) == 1
    Ba = uint16(temp_x*(2^16));
    lin_rgb = double(demosaic(Ba,'gbrg'))/2^16;
elseif pattern(1) == 2
    Ba = uint16(temp_x*(2^16));
    lin_rgb = double(demosaic(Ba,'grbg'))/2^16;
elseif pattern(1) == 3
    Ba = uint16(temp_x*(2^16));
    lin_rgb = double(demosaic(Ba,'bggr'))/2^16;
elseif pattern(1) == 4
    Ba = uint16(temp_x*(2^16));
    lin_rgb = double(demosaic(Ba,'rggb'))/2^16;
else
    lin_rgb = temp_x;
end
temp_x= single(lin_rgb);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% L -> x
%%%%%%%%%%%%%temp_x = CRF_Map(temp_x,I(crf_index,:),B(crf_index,:));
y = temp_x;
end

保存residual，直接处理成residual而不在代码里面处理

function Test_Realistic_Noise_Model()
%%% Test Code for realistic noise model
addpath('./utils');   %%%%%%%%%%%%%this is the function

%% load CRF parameters, the pipeline of the camera, is should be turn off
load('201_CRF_data.mat');
load('dorfCurvesInv.mat');
I_gl = I;
B_gl = B;
I_inv_gl = invI;
B_inv_gl = invB;
mod_scale = 1;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% set parameters
% comment the unnecessary line
input_folder = '/home/guanwp/BasicSR_datasets/val_set14/Set14';
%%%%%%%%%%%%%%%%%%%%% save_mod_folder = '';
save_LR_folder = '/home/guanwp/BasicSR_datasets/val_set14/Set14_noise';
%%%%%%%%%%%%%5% save_bic_folder = '';
save_residual_folder='/home/guanwp/BasicSR_datasets/val_set14/Set14_residual'%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if exist('save_mod_folder', 'var')
    if exist(save_mod_folder, 'dir')
        disp(['It will cover ', save_mod_folder]);
    else
        mkdir(save_mod_folder);
    end
end
if exist('save_LR_folder', 'var')
    if exist(save_LR_folder, 'dir')
        disp(['It will cover ', save_LR_folder]);
    else
        mkdir(save_LR_folder);
    end
end
if exist('save_bic_folder', 'var')
    if exist(save_bic_folder, 'dir')
        disp(['It will cover ', save_bic_folder]);
    else
        mkdir(save_bic_folder);
    end
end

if exist('save_residual_folder', 'var')
    if exist(save_residual_folder, 'dir')
        disp(['It will cover ', save_residual_folder]);
    else
        mkdir(save_residual_folder);
    end
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

idx = 0;
filepaths = dir(fullfile(input_folder,'*.*'));
for i = 1 : length(filepaths)
    [paths,imname,ext] = fileparts(filepaths(i).name);
    if isempty(imname)
        disp('Ignore . folder.');
    elseif strcmp(imname, '.')
        disp('Ignore .. folder.');
    else
        idx = idx + 1;
        str_rlt = sprintf('%d\t%s.\n', idx, imname);
        fprintf(str_rlt);
        % read image
        img = imread(fullfile(input_folder, [imname, ext]));
        img = im2double(img);
        % modcrop
        %%%img = modcrop(img, mod_scale);
        if exist('save_mod_folder', 'var')
            imwrite(img, fullfile(save_mod_folder, [imname, '.png']));
        end
        % LR%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%add noise
        %%%sigma_s = [0.08 0.08 0.08]; % recommend 0~0.16
        %%%%sigma_c = [0.03 0.03 0.03]; % recommend 0~0.06
        %%%%%%%%%%random
        channel = size(img,3);
        sigma_s = 0.16*rand(1,channel,'single'); % original 0.16
        sigma_c = 0.06*rand(1,channel,'single'); % original 0.06
        %%%used the default value
        CRF_index = 5;  % 1~201
        pattern = 5;    % 1: 'gbrg', 2: 'grbg', 3: 'bggr', 4: 'rggb', 5: no mosaic

        im_LR = AddNoiseMosai(img,I_gl,B_gl,I_inv_gl,B_inv_gl, sigma_s,sigma_c, CRF_index, pattern);
        im_residual=im_LR-img;
        %%%%%%%%%%%%%%%%%%im_LR = AddNoiseMosai(img,I_gl,B_gl,I_inv_gl,B_inv_gl, CRF_index, pattern);

        %% JPEG compression
        %%% If JPEG compression is not considered, just commented out the following 
        %%% Code or just set "quality" equal to 100
        qality = 100; % image quality, recommend 60~100

        if exist('save_LR_folder', 'var')
            %%%%%%%%%%%%%%%%%%imwrite(im_LR, fullfile(save_LR_folder, [imname, '_bicLRx4.png']),'Quality', qality);
            imwrite(im_LR, fullfile(save_LR_folder, [imname, '_bicLRx4.png']));
        end
        % Bicubic
        if exist('save_bic_folder', 'var')
            im_B = imresize(im_LR, up_scale, 'bicubic');
            imwrite(im_B, fullfile(save_bic_folder, [imname, '_bicx4.png']));
        end

        if exist('save_residual_folder', 'var')
            %%%%%%%%%%%%%%%%%%imwrite(im_LR, fullfile(save_residual_folder, [imname, '_bicLRx4.png']),'Quality', qality);
            imwrite(im_residual, fullfile(save_residual_folder, [imname, '_bicLRx4.png']));
        end
    end
end
end


%% modcrop
function img = modcrop(img, modulo)
if size(img,3) == 1
    sz = size(img);
    sz = sz - mod(sz, modulo);
    img = img(1:sz(1), 1:sz(2));
else
    tmpsz = size(img);
    sz = tmpsz(1:2);
    sz = sz - mod(sz, modulo);
    img = img(1:sz(1), 1:sz(2),:);
end
end

结果

运行matlab

进入m文件所在目录后，运行

$  matlab -nodesktop -nosplash -r matlabfile

加噪后的效果

高清的图片效果不明显，给张set5的图

 

Experiment

data处理（打开文件 LRHR_dataset.py给出修改代码如下：）

import os.path
import random
import numpy as np
import cv2
import torch
import torch.utils.data as data
import data.util as util


class LRHRDataset(data.Dataset):
    '''
    Read LR and HR image pairs.
    If only HR image is provided, generate LR image on-the-fly.
    The pair is ensured by 'sorted' function, so please check the name convention.
    '''

    def __init__(self, opt):
        super(LRHRDataset, self).__init__()
        self.opt = opt
        self.paths_LR = None
        self.paths_HR = None
        self.LR_env = None  # environment for lmdb
        self.HR_env = None

        # read image list from subset list txt
        if opt['subset_file'] is not None and opt['phase'] == 'train':
            with open(opt['subset_file']) as f:
                self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \
                        for line in f])
            if opt['dataroot_LR'] is not None:
                raise NotImplementedError('Now subset only supports generating LR on-the-fly.')
        else:  # read image list from lmdb or image files
            self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_HR'])
            self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])

        assert self.paths_HR, 'Error: HR path is empty.'
        if self.paths_LR and self.paths_HR:
            assert len(self.paths_LR) == len(self.paths_HR), \
                'HR and LR datasets have different number of images - {}, {}.'.format(\
                len(self.paths_LR), len(self.paths_HR))

        self.random_scale_list = [1]

    def __getitem__(self, index):###Objects can be iterated
        HR_path, LR_path = None, None
        scale = self.opt['scale']#the upscale
        HR_size = self.opt['HR_size']#the HR patch size

        # get HR image
        HR_path = self.paths_HR[index]
        img_HR = util.read_img(self.HR_env, HR_path)###feed the image
        # modcrop in the validation / test phase
        if self.opt['phase'] != 'train':
            img_HR = util.modcrop(img_HR, scale)
        # change color space if necessary
        if self.opt['color']:
            img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]


        #you just should know that this is the process of training
        # get LR image
        if self.paths_LR:
            LR_path = self.paths_LR[index]
            img_LR = util.read_img(self.LR_env, LR_path)
        else:  # down-sampling on-the-fly
            # randomly scale during training
            if self.opt['phase'] == 'train':
                random_scale = random.choice(self.random_scale_list)
                H_s, W_s, _ = img_HR.shape

                def _mod(n, random_scale, scale, thres):
                    rlt = int(n * random_scale)
                    rlt = (rlt // scale) * scale
                    return thres if rlt < thres else rlt

                H_s = _mod(H_s, random_scale, scale, HR_size)
                W_s = _mod(W_s, random_scale, scale, HR_size)
                img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
                # force to 3 channels
                if img_HR.ndim == 2:
                    img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)

            H, W, _ = img_HR.shape
            # using matlab imresize
            img_LR = util.imresize_np(img_HR, 1 / scale, True)
            if img_LR.ndim == 2:
                img_LR = np.expand_dims(img_LR, axis=2)


        if self.opt['phase'] == 'train':
            # if the image size is too small
            H, W, _ = img_HR.shape
            if H < HR_size or W < HR_size:
                img_HR = cv2.resize(
                    np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
                # using matlab imresize
                img_LR = util.imresize_np(img_HR, 1 / scale, True)
                if img_LR.ndim == 2:
                    img_LR = np.expand_dims(img_LR, axis=2)

            H, W, C = img_LR.shape
            LR_size = HR_size // scale###this step make sure the size of the LR is match the size of HR


            ############################this is the augmentation#####################################
            # randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
            img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]

            # augmentation - flip, rotate
            img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
                self.opt['use_rot'])
            #########################################################################################

        #######################################################################################################################################
        #the residual as the target
        #if self.opt['noise_get']:
            #img_HR=img_LR-img_HR##########the noise level map as the target
        #############################################################################################################################################


        # change color space if necessary
        if self.opt['color']:
            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_HR.shape[2] == 3:
            img_HR = img_HR[:, :, [2, 1, 0]]
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        if LR_path is None:
            LR_path = HR_path
        return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path}

    def __len__(self):
        return len(self.paths_HR)

SR_model

import os
from collections import OrderedDict

import torch
import torch.nn as nn
from torch.optim import lr_scheduler

import models.networks as networks#you should read the network
from .base_model import BaseModel#base_model.py just like the base class
from models.modules.loss import TVLoss


class SRModel(BaseModel):
    def __init__(self, opt):
        super(SRModel, self).__init__(opt)
        train_opt = opt['train']

        # define network and load pretrained models
        self.netG = networks.define_G(opt).to(self.device)#you should read the network
        self.load()
        ###############################################
        self.l_pix_tv=0

        if self.is_train:
            self.netG.train()

            # loss
            loss_type = train_opt['pixel_criterion']
            if loss_type == 'l1':
                self.cri_pix = nn.L1Loss().to(self.device)
            elif loss_type == 'l2':
                self.cri_pix = nn.MSELoss().to(self.device)
                #self.l_pix_tv=0##############
                #self.tvloss=TVLoss().to(self.device)###################
            #######################################################################
            elif loss_type=='huber':
                self.cri_pix = nn.SmoothL1Loss().to(self.device)

            elif loss_type=='Cross':
                self.cri_pix = nn.CrossEntropyLoss().to(self.device)

            elif loss_type=='l2_tv':
                self.cri_pix = nn.MSELoss().to(self.device)
                #self.l_pix_tv=1e-5
                #self.tvloss=TVLoss().to(self.device)
            ####################################################################3###
            else:
                raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
            self.l_pix_w = train_opt['pixel_weight']

            # optimizers
            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
            optim_params = []
            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
                if v.requires_grad:
                    optim_params.append(v)
                else:
                    print('WARNING: params [{:s}] will not optimize.'.format(k))
            self.optimizer_G = torch.optim.Adam(
                optim_params, lr=train_opt['lr_G'], weight_decay=wd_G)
            self.optimizers.append(self.optimizer_G)

            # schedulers
            if train_opt['lr_scheme'] == 'MultiStepLR':
                for optimizer in self.optimizers:
                    self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
                        train_opt['lr_steps'], train_opt['lr_gamma']))
            else:
                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')

            self.log_dict = OrderedDict()

        print('---------- Model initialized ------------------')
        self.print_network()
        print('-----------------------------------------------')

    def feed_data(self, data, need_HR=True):#feed the data,
        self.var_L = data['LR'].to(self.device)  # LR
        if need_HR:
            self.real_H = data['HR'].to(self.device)  # HR

    def optimize_parameters(self, step):
        self.optimizer_G.zero_grad()
        self.fake_H = self.netG(self.var_L)
        l_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.real_H)
        ################################################################################################################################
        #####read my CSDN for the TVLoss
        #l_pix=l_pix+self.l_pix_tv * self.tvloss(self.fake_H)#####################################
        ##########################################################

        l_pix.backward()
        self.optimizer_G.step()

        # set log
        self.log_dict['l_pix'] = l_pix.item()

    def test(self):
        self.netG.eval()
        with torch.no_grad():
            self.fake_H = self.netG(self.var_L)
        self.netG.train()

    def test_x8(self):
        # from https://github.com/thstkdgus35/EDSR-PyTorch
        self.netG.eval()
        for k, v in self.netG.named_parameters():
            v.requires_grad = False

        def _transform(v, op):
            # if self.precision != 'single': v = v.float()
            v2np = v.data.cpu().numpy()
            if op == 'v':
                tfnp = v2np[:, :, :, ::-1].copy()
            elif op == 'h':
                tfnp = v2np[:, :, ::-1, :].copy()
            elif op == 't':
                tfnp = v2np.transpose((0, 1, 3, 2)).copy()

            ret = torch.Tensor(tfnp).to(self.device)
            # if self.precision == 'half': ret = ret.half()

            return ret

        lr_list = [self.var_L]
        for tf in 'v', 'h', 't':
            lr_list.extend([_transform(t, tf) for t in lr_list])
        sr_list = [self.netG(aug) for aug in lr_list]
        for i in range(len(sr_list)):
            if i > 3:
                sr_list[i] = _transform(sr_list[i], 't')
            if i % 4 > 1:
                sr_list[i] = _transform(sr_list[i], 'h')
            if (i % 4) % 2 == 1:
                sr_list[i] = _transform(sr_list[i], 'v')

        output_cat = torch.cat(sr_list, dim=0)
        self.fake_H = output_cat.mean(dim=0, keepdim=True)

        for k, v in self.netG.named_parameters():
            v.requires_grad = True
        self.netG.train()

    def get_current_log(self):
        return self.log_dict

    def get_current_visuals(self, need_HR=True):
        out_dict = OrderedDict()
        out_dict['LR'] = self.var_L.detach()[0].float().cpu()
        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
        if need_HR:
            out_dict['HR'] = self.real_H.detach()[0].float().cpu()
        return out_dict

    def print_network(self):
        s, n = self.get_network_description(self.netG)
        print('Number of parameters in G: {:,d}'.format(n))
        if self.is_train:
            message = '-------------- Generator --------------\n' + s + '\n'
            network_path = os.path.join(self.save_dir, '../', 'network.txt')
            with open(network_path, 'w') as f:
                f.write(message)

    def load(self):
        load_path_G = self.opt['path']['pretrain_model_G']
        if load_path_G is not None:
            print('loading model for G [{:s}] ...'.format(load_path_G))
            self.load_network(load_path_G, self.netG)

    def save(self, iter_label):
        self.save_network(self.save_dir, self.netG, 'G', iter_label)

增加网络的接口放于networks.py

# Generator
def define_G(opt):
    gpu_ids = opt['gpu_ids']
    opt_net = opt['network_G']
    which_model = opt_net['which_model_G']#hear decide which model, and thia para is in .json. if you add a new model, this part must be modified

    if which_model == 'sr_resnet':  # SRResNet
        netG = arch.SRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
            nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
            act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')

#############################################################################################################
    elif which_model=='fsrcnn':#FSRCNN
        netG=arch.FSRCNN(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
            nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
            act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')
#############################################################################################################
#############################################################################################################
    elif which_model=='espcn':#ESPCN
        netG=arch.ESPCN(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
            nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
            act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')
#############################################################################################################
#############################################################################################################
    elif which_model=='srresnet':#SRResNet, the Original version
        netG=arch.OSRRESNET(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
            nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
            act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')
#############################################################################################################
#############################################################################################################
    elif which_model=='noise_estimation':#SRResNet, the Original version
        netG=arch.NENET(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
            nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
            act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')
#############################################################################################################

    elif which_model == 'sft_arch':  # SFT-GAN
        netG = sft_arch.SFT_Net()##the SFTNET

    elif which_model == 'RRDB_net':  # RRDB,this is ESRGAN
        netG = arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'],
            nb=opt_net['nb'], gc=opt_net['gc'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'],
            act_type='leakyrelu', mode=opt_net['mode'], upsample_mode='upconv')
    else:
        raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model))

    if opt['is_train']:
        init_weights(netG, init_type='kaiming', scale=0.1)###the weight initing. you can change this to change the method of init_weight
    if gpu_ids:
        assert torch.cuda.is_available()
        netG = nn.DataParallel(netG)
    return netG

网络结构在此处添加architecture.py

#######################################################################################################3
#define the block of the noise estimation network
class Block_of_NENET(nn.Module):
    def __init__(self):
        super(Block_of_NENET,self).__init__()

        self.conv1=nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=True)
        self.bn1=nn.BatchNorm2d(64, affine=True)
        self.relu=nn.ReLU()

    def forward(self, x):
        output = self.relu(self.bn1(self.conv1(x)))
        return output 

#noise estimation network
class NENET(nn.Module):
    def __init__(self, in_nc, out_nc, nf, nb, upscale=2, norm_type='batch', act_type='relu', \
            mode='NAC', res_scale=1, upsample_mode='upconv'):##play attention the upscales
        super(NENET,self).__init__() 

        self.conv1=nn.Conv2d(in_channels=in_nc,out_channels=64,kernel_size=3,stride=1,padding=1,bias=True)
        self.bn=nn.BatchNorm2d(64, affine=True)
        self.relu=nn.ReLU()
        self.block=self.make_layer(Block_of_NENET,15)
        self.last_part= nn.Conv2d(in_channels=64, out_channels=in_nc, kernel_size=3, stride=1, padding=1, bias=True)
 

    def make_layer(self, block, num_of_layer):
        layers = []
        for _ in range(num_of_layer):
            layers.append(block())
        return nn.Sequential(*layers)
 

    def forward(self, x):#
         out = self.conv1(x)
         out=self.bn(out)
         out=self.relu(out)
         out = self.block(out)
         out = self.last_part(out)
 
         return out

训练

python train.py -opt options/train/train_ne.json.json

忘记做subimage了。。。。补上，然后把学习率也改一下

 

结果

loss飞掉，并且网络输出是黑色的。。。。

接下来选择不用residual作为target，并修改相关setting

{
  "name": "n_g" //"001_RRDB_PSNR_x4_DIV2K" //  please remove "debug_" during training or tensorboard wounld not work
  ,
  "use_tb_logger": true,
  "model": "sr",
  //"crop_scale": 0,
   "scale": 1//it must be 1
  ,
  "gpu_ids": [3],
  "datasets": {
    "train": {
      "name": "DIV2K800",
      "mode": "LRHR" //it must be this, and the detail would be shown in LRHR_dataset.py
      //, "noise_get": true///
      ,
      "dataroot_HR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub" ///must be sub
      ,
      "dataroot_LR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub_n",
      "subset_file": null,
      "use_shuffle": true,
      "n_workers": 8,
      "batch_size": 16//32 //how many samples in each iters
      ,
      "HR_size": 192 // 128 | 192
      ,
      "use_flip": false //true//
      ,
      "use_rot": false //true
    },
    "val": {
      "name": "val_set5",
      "mode": "LRHR",
      "dataroot_HR": "/home/guanwp/BasicSR_datasets/val_set5/Set5",
      "dataroot_LR": "/home/guanwp/BasicSR_datasets/val_set5/Set5_n"
      //, "noise_get": true///this is important
    }
  },
  "path": {
    "root": "/home/guanwp/BasicSR-master/",
    "pretrain_model_G": null,
    "experiments_root": "/home/guanwp/BasicSR-master/experiments/",
    "models": "/home/guanwp/BasicSR-master/experiments/n_g/models",
    "log": "/home/guanwp/BasicSR-master/experiments/n_g",
    "val_images": "/home/guanwp/BasicSR-master/experiments/n_g/val_images"
  },
  "network_G": {
    "which_model_G": "noise_estimation" //"espcn"//"srresnet"//"sr_resnet"//"fsrcnn"//"sr_resnet" // RRDB_net | sr_resnet
    ,
    "norm_type": null,
    "mode": "CNA",
    "nf": 64 //56//64
    ,
    "nb": 23,
    "in_nc": 3,
    "out_nc": 3,
    "gc": 32,
    "group": 1
  },
  "train": {
    "lr_G": 8e-4 //1e-3//2e-4
    ,
    "lr_scheme": "MultiStepLR",
    "lr_steps": [131000,170000,190000,210000,230000],
    "lr_gamma": 0.5,
    "pixel_criterion": "l2" //"l2_tv"//"l1"//'l2'//huber//Cross   //should be MSE LOSS
    ,
    "pixel_weight": 1.0,
    "val_freq": 1e3,
    "manual_seed": 0,
    "niter": 2.6e5 //2e6//1e6
  },
  "logger": {
    "print_freq": 200,
    "save_checkpoint_freq": 1e3
  }
}

test

python test.py -opt options/test/test_ne.json

经过代码调试后，去噪效果如下图所示：从左到右分别是ground truth、noise image、denoise image。可以看出去噪效果还算不错，但是会导致图像变得过平滑，感觉这样的话，可能加上TVLoss就会好点了。以及本博文的网络其实跟DnCNN是差不多的，应该就是多了一个BN层，所以PSNR的效果是32dB左右，也可以达到原文的效果了~

 

Reference resources

关于Signal Dependent Noise Level Estimation

参考（https://ww2.mathworks.cn/matlabcentral/fileexchange/43224-signal-dependent-noise-level-estimation）

关于total variation

https://blog.csdn.net/weixin_42447651/article/details/82990941

关于matlab中的bsxfun

https://blog.csdn.net/ddreaming/article/details/52176790

matlab中permute

matlab中repmat

那么，bsxfun和repmat的作用应该是一样的？