基于pytorch的噪声估计网络

版权声明: https://blog.csdn.net/gwplovekimi/article/details/84937991

给出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:

然后打开文件 LRHR_dataset.py给出修改代码如下:

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)。

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

运行matlab

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

$  matlab -nodesktop -nosplash -r matlabfile

加噪后的效果

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

Experiment

增加网络的接口放于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

训练

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

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

结果

test

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

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

猜你喜欢

转载自blog.csdn.net/gwplovekimi/article/details/84937991