def bicubic_degradation(x, sf=3): ''' Args: x: HxWxC image, [0, 1] sf: down-scale factor Return: bicubicly downsampled LR image ''' x = util.imresize_np(x, scale=1 / sf) return x
def __getitem__(self, index): # ------------------------------------ # get H image # ------------------------------------ H_path = self.paths_H[index] H_resolution = [int(s) for s in self.sizes_H[index].split('_')] img_H = _read_img_lmdb(self.h_env, H_path, H_resolution) img_H = util.uint2single(img_H) img_H = util.modcrop(img_H, self.sf) # ---------------------------------------------- # get down4 sharp image(img_h0) to be target in train for deblur # ---------------------------------------------- img_h0 = util.imresize_np(img_H, 1 / 4, True) # ------------------------------------ # get L image # ------------------------------------ L_path = self.paths_L[index] L_resolution = [int(s) for s in self.sizes_L[index].split('_')] img_L = _read_img_lmdb(self.l_env, L_path, L_resolution) img_L = util.uint2single(img_L) # ------------------------------------ # L/H pairs, HWC to CHW, numpy to tensor # ------------------------------------ img_deblur = [img_L,img_h0] img_deblur = util.generate_pyramid(*img_deblur,n_scales=3) img_deblur =util.np2tensor(*img_deblur) img_L = img_deblur[0] img_H = util.single2tensor3(img_H) img_L0 = img_L[0] return {'L0':img_L0, 'L': img_L, 'H': img_H, 'L_path': L_path, 'H_path': H_path, 'ls': img_L, 'hs': img_deblur[1]}
def main(): # ---------------------------------------- # Preparation # ---------------------------------------- noise_level_img = 0 # default: 0, noise level for LR image noise_level_model = noise_level_img # noise level for model model_name = 'dpsr_x4_gan' # 'dpsr_x2' | 'dpsr_x3' | 'dpsr_x4' | 'dpsr_x4_gan' testset_name = 'set5' # test set, 'set5' | 'srbsd68' need_degradation = True # default: True x8 = False # default: False, x8 to boost performance sf = [int(s) for s in re.findall(r'\d+', model_name)][0] # scale factor show_img = False # default: False task_current = 'sr' # 'dn' for denoising | 'sr' for super-resolution n_channels = 3 # fixed nc = 96 # fixed, number of channels nb = 16 # fixed, number of conv layers model_pool = 'model_zoo' # fixed testsets = 'testsets' # fixed results = 'results' # fixed result_name = testset_name + '_' + model_name border = sf if task_current == 'sr' else 0 # shave boader to calculate PSNR and SSIM model_path = os.path.join(model_pool, model_name + '.pth') # ---------------------------------------- # L_path, E_path, H_path # ---------------------------------------- L_path = os.path.join(testsets, testset_name) # L_path, for Low-quality images H_path = L_path # H_path, for High-quality images E_path = os.path.join(results, result_name) # E_path, for Estimated images util.mkdir(E_path) if H_path == L_path: need_degradation = True logger_name = result_name utils_logger.logger_info(logger_name, log_path=os.path.join(E_path, logger_name + '.log')) logger = logging.getLogger(logger_name) need_H = True if H_path is not None else False device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # ---------------------------------------- # load model # ---------------------------------------- from models.network_dpsr import MSRResNet_prior as net model = net(in_nc=n_channels + 1, out_nc=n_channels, nc=nc, nb=nb, upscale=sf, act_mode='R', upsample_mode='pixelshuffle') model.load_state_dict(torch.load(model_path), strict=False) model.eval() for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) logger.info('Model path: {:s}'.format(model_path)) number_parameters = sum(map(lambda x: x.numel(), model.parameters())) logger.info('Params number: {}'.format(number_parameters)) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] logger.info('model_name:{}, model sigma:{}, image sigma:{}'.format( model_name, noise_level_img, noise_level_model)) logger.info(L_path) L_paths = util.get_image_paths(L_path) H_paths = util.get_image_paths(H_path) if need_H else None for idx, img in enumerate(L_paths): # ------------------------------------ # (1) img_L # ------------------------------------ img_name, ext = os.path.splitext(os.path.basename(img)) # logger.info('{:->4d}--> {:>10s}'.format(idx+1, img_name+ext)) img_L = util.imread_uint(img, n_channels=n_channels) img_L = util.uint2single(img_L) # degradation process, bicubic downsampling + Gaussian noise if need_degradation: img_L = util.modcrop(img_L, sf) img_L = util.imresize_np(img_L, 1 / sf) np.random.seed(seed=0) # for reproducibility img_L += np.random.normal(0, noise_level_img / 255., img_L.shape) util.imshow(util.single2uint(img_L), title='LR image with noise level {}'.format( noise_level_img)) if show_img else None img_L = util.single2tensor4(img_L) noise_level_map = torch.full((1, 1, img_L.size(2), img_L.size(3)), noise_level_model / 255.).type_as(img_L) img_L = torch.cat((img_L, noise_level_map), dim=1) img_L = img_L.to(device) # ------------------------------------ # (2) img_E # ------------------------------------ if not x8: img_E = model(img_L) else: img_E = utils_model.test_mode(model, img_L, mode=3, sf=sf) img_E = util.tensor2uint(img_E) if need_H: # -------------------------------- # (3) img_H # -------------------------------- img_H = util.imread_uint(H_paths[idx], n_channels=n_channels) img_H = img_H.squeeze() img_H = util.modcrop(img_H, sf) # -------------------------------- # PSNR and SSIM # -------------------------------- psnr = util.calculate_psnr(img_E, img_H, border=border) ssim = util.calculate_ssim(img_E, img_H, border=border) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format( img_name + ext, psnr, ssim)) util.imshow(np.concatenate([img_E, img_H], axis=1), title='Recovered / Ground-truth') if show_img else None if np.ndim(img_H) == 3: # RGB image img_E_y = util.rgb2ycbcr(img_E, only_y=True) img_H_y = util.rgb2ycbcr(img_H, only_y=True) psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border) ssim_y = util.calculate_ssim(img_E_y, img_H_y, border=border) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) # ------------------------------------ # save results # ------------------------------------ util.imsave(img_E, os.path.join(E_path, img_name + '.png')) if need_H: ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info( 'Average PSNR/SSIM(RGB) - {} - x{} --PSNR: {:.2f} dB; SSIM: {:.4f}' .format(result_name, sf, ave_psnr, ave_ssim)) if np.ndim(img_H) == 3: ave_psnr_y = sum(test_results['psnr_y']) / len( test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len( test_results['ssim_y']) logger.info( 'Average PSNR/SSIM( Y ) - {} - x{} - PSNR: {:.2f} dB; SSIM: {:.4f}' .format(result_name, sf, ave_psnr_y, ave_ssim_y))
def main(): # ---------------------------------------- # Preparation # ---------------------------------------- model_name = 'usrnet' # 'usrgan' | 'usrnet' | 'usrgan_tiny' | 'usrnet_tiny' testset_name = 'set5' # test set, 'set5' | 'srbsd68' need_degradation = True # default: True sf = 4 # scale factor, only from {2, 3, 4} show_img = False # default: False save_L = True # save LR image save_E = True # save estimated image # load approximated bicubic kernels #kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernels_bicubicx234.mat'))['kernels'] kernels = loadmat(os.path.join('kernels', 'kernels_bicubicx234.mat'))['kernels'] kernel = kernels[0, sf - 2].astype(np.float64) kernel = util.single2tensor4(kernel[..., np.newaxis]) task_current = 'sr' # fixed, 'sr' for super-resolution n_channels = 3 # fixed, 3 for color image model_pool = 'model_zoo' # fixed testsets = 'testsets' # fixed results = 'results' # fixed noise_level_img = 0 # fixed: 0, noise level for LR image noise_level_model = noise_level_img # fixed, noise level of model, default 0 result_name = testset_name + '_' + model_name + '_bicubic' border = sf if task_current == 'sr' else 0 # shave boader to calculate PSNR and SSIM model_path = os.path.join(model_pool, model_name + '.pth') # ---------------------------------------- # L_path, E_path, H_path # ---------------------------------------- L_path = os.path.join( testsets, testset_name) # L_path, fixed, for Low-quality images H_path = L_path # H_path, 'None' | L_path, for High-quality images E_path = os.path.join(results, result_name) # E_path, fixed, for Estimated images util.mkdir(E_path) if H_path == L_path: need_degradation = True logger_name = result_name utils_logger.logger_info(logger_name, log_path=os.path.join(E_path, logger_name + '.log')) logger = logging.getLogger(logger_name) need_H = True if H_path is not None else False device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # ---------------------------------------- # load model # ---------------------------------------- from models.network_usrnet import USRNet as net # for pytorch version <= 1.7.1 # from models.network_usrnet_v1 import USRNet as net # for pytorch version >=1.8.1 if 'tiny' in model_name: model = net(n_iter=6, h_nc=32, in_nc=4, out_nc=3, nc=[16, 32, 64, 64], nb=2, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose") else: model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=2, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose") model.load_state_dict(torch.load(model_path), strict=True) model.eval() for key, v in model.named_parameters(): v.requires_grad = False number_parameters = sum(map(lambda x: x.numel(), model.parameters())) logger.info('Params number: {}'.format(number_parameters)) model = model.to(device) logger.info('Model path: {:s}'.format(model_path)) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] logger.info('model_name:{}, image sigma:{}'.format(model_name, noise_level_img)) logger.info(L_path) L_paths = util.get_image_paths(L_path) H_paths = util.get_image_paths(H_path) if need_H else None for idx, img in enumerate(L_paths): # ------------------------------------ # (1) img_L # ------------------------------------ img_name, ext = os.path.splitext(os.path.basename(img)) logger.info('{:->4d}--> {:>10s}'.format(idx + 1, img_name + ext)) img_L = util.imread_uint(img, n_channels=n_channels) img_L = util.uint2single(img_L) # degradation process, bicubic downsampling if need_degradation: img_L = util.modcrop(img_L, sf) img_L = util.imresize_np(img_L, 1 / sf) # img_L = util.uint2single(util.single2uint(img_L)) # np.random.seed(seed=0) # for reproducibility # img_L += np.random.normal(0, noise_level_img/255., img_L.shape) w, h = img_L.shape[:2] if save_L: util.imsave( util.single2uint(img_L), os.path.join(E_path, img_name + '_LR_x' + str(sf) + '.png')) img = cv2.resize(img_L, (sf * h, sf * w), interpolation=cv2.INTER_NEAREST) img = utils_deblur.wrap_boundary_liu(img, [ int(np.ceil(sf * w / 8 + 2) * 8), int(np.ceil(sf * h / 8 + 2) * 8) ]) img_wrap = sr.downsample_np(img, sf, center=False) img_wrap[:w, :h, :] = img_L img_L = img_wrap util.imshow(util.single2uint(img_L), title='LR image with noise level {}'.format( noise_level_img)) if show_img else None img_L = util.single2tensor4(img_L) img_L = img_L.to(device) # ------------------------------------ # (2) img_E # ------------------------------------ sigma = torch.tensor(noise_level_model).float().view([1, 1, 1, 1]) [img_L, kernel, sigma] = [el.to(device) for el in [img_L, kernel, sigma]] img_E = model(img_L, kernel, sf, sigma) img_E = util.tensor2uint(img_E) img_E = img_E[:sf * w, :sf * h, :] if need_H: # -------------------------------- # (3) img_H # -------------------------------- img_H = util.imread_uint(H_paths[idx], n_channels=n_channels) img_H = img_H.squeeze() img_H = util.modcrop(img_H, sf) # -------------------------------- # PSNR and SSIM # -------------------------------- psnr = util.calculate_psnr(img_E, img_H, border=border) ssim = util.calculate_ssim(img_E, img_H, border=border) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format( img_name + ext, psnr, ssim)) util.imshow(np.concatenate([img_E, img_H], axis=1), title='Recovered / Ground-truth') if show_img else None if np.ndim(img_H) == 3: # RGB image img_E_y = util.rgb2ycbcr(img_E, only_y=True) img_H_y = util.rgb2ycbcr(img_H, only_y=True) psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border) ssim_y = util.calculate_ssim(img_E_y, img_H_y, border=border) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) # ------------------------------------ # save results # ------------------------------------ if save_E: util.imsave( img_E, os.path.join( E_path, img_name + '_x' + str(sf) + '_' + model_name + '.png')) if need_H: ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info( 'Average PSNR/SSIM(RGB) - {} - x{} --PSNR: {:.2f} dB; SSIM: {:.4f}' .format(result_name, sf, ave_psnr, ave_ssim)) if np.ndim(img_H) == 3: ave_psnr_y = sum(test_results['psnr_y']) / len( test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len( test_results['ssim_y']) logger.info( 'Average PSNR/SSIM( Y ) - {} - x{} - PSNR: {:.2f} dB; SSIM: {:.4f}' .format(result_name, sf, ave_psnr_y, ave_ssim_y))
def main(): # -------------------------------- # let's start! # -------------------------------- utils_logger.logger_info('test_dpsr', log_path='test_dpsr.log') logger = logging.getLogger('test_dpsr') # basic setting # ================================================ sf = 4 # scale factor noise_level_img = 0 / 255.0 # noise level of low quality image, default 0 noise_level_model = noise_level_img # noise level of model, default 0 show_img = True use_srganplus = True # 'True' for SRGAN+ (x4) and 'False' for SRResNet+ (x2,x3,x4) testsets = 'testsets' testset_current = 'BSD68' if use_srganplus and sf == 4: model_prefix = 'DPSRGAN' save_suffix = 'dpsrgan' else: model_prefix = 'DPSR' save_suffix = 'dpsr' model_path = os.path.join('DPSR_models', model_prefix + 'x%01d.pth' % (sf)) iter_num = 15 # number of iterations, fixed n_channels = 3 # only color images, fixed border = sf**2 # shave boader to calculate PSNR, fixed # k_type = ('d', 'm', 'g') k_type = ('m') # motion blur kernel # ================================================ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # -------------------------------- # load model # -------------------------------- model = SRResNet(in_nc=4, out_nc=3, nc=96, nb=16, upscale=sf, act_mode='R', upsample_mode='pixelshuffle') model.load_state_dict(torch.load(model_path), strict=True) model.eval() for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) logger.info('Model path {:s}. Testing...'.format(model_path)) # -------------------------------- # read image (img) and kernel (k) # -------------------------------- test_results = OrderedDict() for k_type_n in range(len(k_type)): # --1--> L_folder, folder of Low-quality images testsubset_current = 'x%01d_%01s' % (sf, k_type[k_type_n]) L_folder = os.path.join(testsets, testset_current, testsubset_current) # --2--> E_folder, folder of Estimated images E_folder = os.path.join(testsets, testset_current, testsubset_current + '_' + save_suffix) util.mkdir(E_folder) # --3--> H_folder, folder of High-quality images H_folder = os.path.join(testsets, testset_current, 'GT') test_results['psnr_' + k_type[k_type_n]] = [] logger.info(L_folder) idx = 0 for im in os.listdir(os.path.join(L_folder)): if im.endswith('.jpg') or im.endswith('.bmp') or im.endswith( '.png'): # -------------------------------- # (1) img_L # -------------------------------- idx += 1 img_name, ext = os.path.splitext(im) img_L = util.imread_uint(os.path.join(L_folder, im), n_channels=n_channels) util.imshow(img_L) if show_img else None np.random.seed(seed=0) # for reproducibility img_L = util.unit2single(img_L) + np.random.normal( 0, noise_level_img, img_L.shape) # -------------------------------- # (2) kernel # -------------------------------- k = loadmat(os.path.join(L_folder, img_name + '.mat'))['kernel'] k = k.astype(np.float32) k /= np.sum(k) # -------------------------------- # (3) get upperleft, denominator # -------------------------------- upperleft, denominator = utils_deblur.get_uperleft_denominator( img_L, k) # -------------------------------- # (4) get rhos and sigmas # -------------------------------- rhos, sigmas = utils_deblur.get_rho_sigma(sigma=max( 0.255 / 255., noise_level_model), iter_num=iter_num) # -------------------------------- # (5) main iteration # -------------------------------- z = img_L rhos = np.float32(rhos) sigmas = np.float32(sigmas) for i in range(iter_num): # -------------------------------- # step 1, Eq. (9) // FFT # -------------------------------- rho = rhos[i] if i != 0: z = util.imresize_np(z, 1 / sf, True) z = np.real( np.fft.ifft2( (upperleft + rho * np.fft.fft2(z, axes=(0, 1))) / (denominator + rho), axes=(0, 1))) # imsave('LR_deblurred_%02d.png'%i, np.clip(z, 0, 1)) # -------------------------------- # step 2, Eq. (12) // super-resolver # -------------------------------- sigma = torch.from_numpy(np.array(sigmas[i])) img_L = util.single2tensor4(z) noise_level_map = torch.ones( (1, 1, img_L.size(2), img_L.size(3)), dtype=torch.float).mul_(sigma) img_L = torch.cat((img_L, noise_level_map), dim=1) img_L = img_L.to(device) # with torch.no_grad(): z = model(img_L) z = util.tensor2single(z) # -------------------------------- # (6) img_E # -------------------------------- img_E = util.single2uint(z) # np.uint8((z * 255.0).round()) # -------------------------------- # (7) img_H # -------------------------------- img_H = util.imread_uint(os.path.join(H_folder, img_name[:7] + '.png'), n_channels=n_channels) util.imshow( np.concatenate([img_E, img_H], axis=1), title='Recovered / Ground-truth') if show_img else None psnr = util.calculate_psnr(img_E, img_H, border=border) logger.info('{:->4d}--> {:>10s}, {:.2f}dB'.format( idx, im, psnr)) test_results['psnr_' + k_type[k_type_n]].append(psnr) util.imsave(img_E, os.path.join(E_folder, img_name + ext)) ave_psnr = sum(test_results['psnr_' + k_type[k_type_n]]) / len( test_results['psnr_' + k_type[k_type_n]]) logger.info( '------> Average PSNR(RGB) of ({} - {}) is : {:.2f} dB'.format( testset_current, testsubset_current, ave_psnr))
def __getitem__(self, index): # ------------------------------------ # get H image # ------------------------------------ H_path = self.paths_H[index] img_H = util.imread_uint(H_path, self.n_channels) img_H = util.uint2single(img_H) # ------------------------------------ # modcrop for SR # ------------------------------------ img_H = util.modcrop(img_H, self.sf) # ------------------------------------ # sythesize L image via matlab's bicubic # ------------------------------------ H, W, _ = img_H.shape img_L = util.imresize_np(img_H, 1 / self.sf, True) if self.opt['phase'] == 'train': """ # -------------------------------- # get L/H patch pairs # -------------------------------- """ H, W, C = img_L.shape # -------------------------------- # randomly crop L patch # -------------------------------- rnd_h = random.randint(0, max(0, H - self.L_size)) rnd_w = random.randint(0, max(0, W - self.L_size)) img_L = img_L[rnd_h:rnd_h + self.L_size, rnd_w:rnd_w + self.L_size, :] # -------------------------------- # crop corresponding H patch # -------------------------------- rnd_h_H, rnd_w_H = int(rnd_h * self.sf), int(rnd_w * self.sf) img_H = img_H[rnd_h_H:rnd_h_H + self.patch_size, rnd_w_H:rnd_w_H + self.patch_size, :] # -------------------------------- # augmentation - flip and/or rotate # -------------------------------- mode = np.random.randint(0, 8) img_L, img_H = util.augment_img( img_L, mode=mode), util.augment_img(img_H, mode=mode) # -------------------------------- # get patch pairs # -------------------------------- img_H, img_L = util.single2tensor3(img_H), util.single2tensor3( img_L) # -------------------------------- # select noise level and get Gaussian noise # -------------------------------- if random.random() < 0.1: noise_level = torch.zeros(1).float() else: noise_level = torch.FloatTensor([ np.random.uniform(self.sigma_min, self.sigma_max) ]) / 255.0 # noise_level = torch.rand(1)*50/255.0 # noise_level = torch.min(torch.from_numpy(np.float32([7*np.random.chisquare(2.5)/255.0])),torch.Tensor([50./255.])) else: img_H, img_L = util.single2tensor3(img_H), util.single2tensor3( img_L) noise_level = torch.FloatTensor([self.sigma_test]) # ------------------------------------ # add noise # ------------------------------------ noise = torch.randn(img_L.size()).mul_(noise_level).float() img_L.add_(noise) # ------------------------------------ # get noise level map M # ------------------------------------ M_vector = noise_level.unsqueeze(1).unsqueeze(1) M = M_vector.repeat(1, img_L.size()[-2], img_L.size()[-1]) """ # ------------------------------------- # concat L and noise level map M # ------------------------------------- """ img_L = torch.cat((img_L, M), 0) L_path = H_path return {'L': img_L, 'H': img_H, 'L_path': L_path, 'H_path': H_path}
def main(): # ---------------------------------------- # Preparation # ---------------------------------------- noise_level_img = 0/255.0 # set AWGN noise level for LR image, default: 0, noise_level_model = noise_level_img # setnoise level of model, default 0 model_name = 'drunet_color' # set denoiser, | 'drunet_color' | 'ircnn_gray' | 'drunet_gray' | 'ircnn_color' testset_name = 'srbsd68' # set test set, 'set5' | 'srbsd68' x8 = True # default: False, x8 to boost performance test_sf = [2] # set scale factor, default: [2, 3, 4], [2], [3], [4] iter_num = 24 # set number of iterations, default: 24 for SISR modelSigma1 = 49 # set sigma_1, default: 49 classical_degradation = True # set classical degradation or bicubic degradation show_img = False # default: False save_L = True # save LR image save_E = True # save estimated image save_LEH = False # save zoomed LR, E and H images task_current = 'sr' # 'sr' for super-resolution n_channels = 1 if 'gray' in model_name else 3 # fixed model_zoo = 'model_zoo' # fixed testsets = 'testsets' # fixed results = 'results' # fixed result_name = testset_name + '_' + task_current + '_' + model_name model_path = os.path.join(model_zoo, model_name+'.pth') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') torch.cuda.empty_cache() # ---------------------------------------- # L_path, E_path, H_path # ---------------------------------------- L_path = os.path.join(testsets, testset_name) # L_path, for Low-quality images E_path = os.path.join(results, result_name) # E_path, for Estimated images util.mkdir(E_path) logger_name = result_name utils_logger.logger_info(logger_name, log_path=os.path.join(E_path, logger_name+'.log')) logger = logging.getLogger(logger_name) # ---------------------------------------- # load model # ---------------------------------------- if 'drunet' in model_name: from models.network_unet import UNetRes as net model = net(in_nc=n_channels+1, out_nc=n_channels, nc=[64, 128, 256, 512], nb=4, act_mode='R', downsample_mode="strideconv", upsample_mode="convtranspose") model.load_state_dict(torch.load(model_path), strict=True) model.eval() for _, v in model.named_parameters(): v.requires_grad = False model = model.to(device) elif 'ircnn' in model_name: from models.network_dncnn import IRCNN as net model = net(in_nc=n_channels, out_nc=n_channels, nc=64) model25 = torch.load(model_path) former_idx = 0 logger.info('model_name:{}, image sigma:{:.3f}, model sigma:{:.3f}'.format(model_name, noise_level_img, noise_level_model)) logger.info('Model path: {:s}'.format(model_path)) logger.info(L_path) L_paths = util.get_image_paths(L_path) # -------------------------------- # load kernel # -------------------------------- # kernels = hdf5storage.loadmat(os.path.join('kernels', 'Levin09.mat'))['kernels'] if classical_degradation: kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernels_12.mat'))['kernels'] else: kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernel_bicubicx234.mat'))['kernels'] test_results_ave = OrderedDict() test_results_ave['psnr_sf_k'] = [] test_results_ave['psnr_y_sf_k'] = [] for sf in test_sf: border = sf modelSigma2 = max(sf, noise_level_model*255.) k_num = 8 if classical_degradation else 1 for k_index in range(k_num): logger.info('--------- sf:{:>1d} --k:{:>2d} ---------'.format(sf, k_index)) test_results = OrderedDict() test_results['psnr'] = [] test_results['psnr_y'] = [] if not classical_degradation: # for bicubic degradation k_index = sf-2 k = kernels[0, k_index].astype(np.float64) util.surf(k) if show_img else None for idx, img in enumerate(L_paths): # -------------------------------- # (1) get img_L # -------------------------------- img_name, ext = os.path.splitext(os.path.basename(img)) img_H = util.imread_uint(img, n_channels=n_channels) img_H = util.modcrop(img_H, sf) # modcrop if classical_degradation: img_L = sr.classical_degradation(img_H, k, sf) util.imshow(img_L) if show_img else None img_L = util.uint2single(img_L) else: img_L = util.imresize_np(util.uint2single(img_H), 1/sf) np.random.seed(seed=0) # for reproducibility img_L += np.random.normal(0, noise_level_img, img_L.shape) # add AWGN # -------------------------------- # (2) get rhos and sigmas # -------------------------------- rhos, sigmas = pnp.get_rho_sigma(sigma=max(0.255/255., noise_level_model), iter_num=iter_num, modelSigma1=modelSigma1, modelSigma2=modelSigma2, w=1) rhos, sigmas = torch.tensor(rhos).to(device), torch.tensor(sigmas).to(device) # -------------------------------- # (3) initialize x, and pre-calculation # -------------------------------- x = cv2.resize(img_L, (img_L.shape[1]*sf, img_L.shape[0]*sf), interpolation=cv2.INTER_CUBIC) if np.ndim(x)==2: x = x[..., None] if classical_degradation: x = sr.shift_pixel(x, sf) x = util.single2tensor4(x).to(device) img_L_tensor, k_tensor = util.single2tensor4(img_L), util.single2tensor4(np.expand_dims(k, 2)) [k_tensor, img_L_tensor] = util.todevice([k_tensor, img_L_tensor], device) FB, FBC, F2B, FBFy = sr.pre_calculate(img_L_tensor, k_tensor, sf) # -------------------------------- # (4) main iterations # -------------------------------- for i in range(iter_num): # -------------------------------- # step 1, FFT # -------------------------------- tau = rhos[i].float().repeat(1, 1, 1, 1) x = sr.data_solution(x.float(), FB, FBC, F2B, FBFy, tau, sf) if 'ircnn' in model_name: current_idx = np.int(np.ceil(sigmas[i].cpu().numpy()*255./2.)-1) if current_idx != former_idx: model.load_state_dict(model25[str(current_idx)], strict=True) model.eval() for _, v in model.named_parameters(): v.requires_grad = False model = model.to(device) former_idx = current_idx # -------------------------------- # step 2, denoiser # -------------------------------- if x8: x = util.augment_img_tensor4(x, i % 8) if 'drunet' in model_name: x = torch.cat((x, sigmas[i].float().repeat(1, 1, x.shape[2], x.shape[3])), dim=1) x = utils_model.test_mode(model, x, mode=2, refield=32, min_size=256, modulo=16) elif 'ircnn' in model_name: x = model(x) if x8: if i % 8 == 3 or i % 8 == 5: x = util.augment_img_tensor4(x, 8 - i % 8) else: x = util.augment_img_tensor4(x, i % 8) # -------------------------------- # (3) img_E # -------------------------------- img_E = util.tensor2uint(x) if save_E: util.imsave(img_E, os.path.join(E_path, img_name+'_x'+str(sf)+'_k'+str(k_index)+'_'+model_name+'.png')) if n_channels == 1: img_H = img_H.squeeze() # -------------------------------- # (4) img_LEH # -------------------------------- img_L = util.single2uint(img_L).squeeze() if save_LEH: k_v = k/np.max(k)*1.0 if n_channels==1: k_v = util.single2uint(k_v) else: k_v = util.single2uint(np.tile(k_v[..., np.newaxis], [1, 1, n_channels])) k_v = cv2.resize(k_v, (3*k_v.shape[1], 3*k_v.shape[0]), interpolation=cv2.INTER_NEAREST) img_I = cv2.resize(img_L, (sf*img_L.shape[1], sf*img_L.shape[0]), interpolation=cv2.INTER_NEAREST) img_I[:k_v.shape[0], -k_v.shape[1]:, ...] = k_v img_I[:img_L.shape[0], :img_L.shape[1], ...] = img_L util.imshow(np.concatenate([img_I, img_E, img_H], axis=1), title='LR / Recovered / Ground-truth') if show_img else None util.imsave(np.concatenate([img_I, img_E, img_H], axis=1), os.path.join(E_path, img_name+'_x'+str(sf)+'_k'+str(k_index)+'_LEH.png')) if save_L: util.imsave(img_L, os.path.join(E_path, img_name+'_x'+str(sf)+'_k'+str(k_index)+'_LR.png')) psnr = util.calculate_psnr(img_E, img_H, border=border) test_results['psnr'].append(psnr) logger.info('{:->4d}--> {:>10s} -- sf:{:>1d} --k:{:>2d} PSNR: {:.2f}dB'.format(idx+1, img_name+ext, sf, k_index, psnr)) if n_channels == 3: img_E_y = util.rgb2ycbcr(img_E, only_y=True) img_H_y = util.rgb2ycbcr(img_H, only_y=True) psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border) test_results['psnr_y'].append(psnr_y) # -------------------------------- # Average PSNR for all kernels # -------------------------------- ave_psnr_k = sum(test_results['psnr']) / len(test_results['psnr']) logger.info('------> Average PSNR(RGB) of ({}) scale factor: ({}), kernel: ({}) sigma: ({:.2f}): {:.2f} dB'.format(testset_name, sf, k_index, noise_level_model, ave_psnr_k)) test_results_ave['psnr_sf_k'].append(ave_psnr_k) if n_channels == 3: # RGB image ave_psnr_y_k = sum(test_results['psnr_y']) / len(test_results['psnr_y']) logger.info('------> Average PSNR(Y) of ({}) scale factor: ({}), kernel: ({}) sigma: ({:.2f}): {:.2f} dB'.format(testset_name, sf, k_index, noise_level_model, ave_psnr_y_k)) test_results_ave['psnr_y_sf_k'].append(ave_psnr_y_k) # --------------------------------------- # Average PSNR for all sf and kernels # --------------------------------------- ave_psnr_sf_k = sum(test_results_ave['psnr_sf_k']) / len(test_results_ave['psnr_sf_k']) logger.info('------> Average PSNR of ({}) {:.2f} dB'.format(testset_name, ave_psnr_sf_k)) if n_channels == 3: ave_psnr_y_sf_k = sum(test_results_ave['psnr_y_sf_k']) / len(test_results_ave['psnr_y_sf_k']) logger.info('------> Average PSNR of ({}) {:.2f} dB'.format(testset_name, ave_psnr_y_sf_k))
def generate_LR(x, sf): lr = utils_image.imresize_np(x, scale=1 / sf).clip(0, 255) return lr
def __getitem__(self, index): # ------------------------------------ # get H image # ------------------------------------ H_path = self.paths_H[index] H_resolution = [int(s) for s in self.sizes_H[index].split('_')] img_H = _read_img_lmdb(self.h_env, H_path, H_resolution) img_H = util.uint2single(img_H) img_H = util.modcrop(img_H, self.sf) # ---------------------------------------------- # get down4 sharp image(img_h0) to be target in train for deblur # ---------------------------------------------- img_h0 = util.imresize_np(img_H, 1 / 4, True) # ------------------------------------ # get L image # ------------------------------------ L_path = self.paths_L[index] L_resolution = [int(s) for s in self.sizes_L[index].split('_')] img_L = _read_img_lmdb(self.l_env, L_path, L_resolution) img_L = util.uint2single(img_L) # # ---------------------------------------------- # # get blur image to be input in train for deblur # # ---------------------------------------------- # img_l0 = img_L # ------------------------------------ # if train, get L/H patch pair # ------------------------------------ if self.opt['phase'] == 'train': H, W, C = img_L.shape # -------------------------------- # randomly crop the L patch # -------------------------------- rnd_h = random.randint(0, max(0, H - self.L_size)) rnd_w = random.randint(0, max(0, W - self.L_size)) img_L = img_L[rnd_h:rnd_h + self.L_size, rnd_w:rnd_w + self.L_size, :] img_h0 = img_h0[rnd_h:rnd_h + self.L_size, rnd_w:rnd_w + self.L_size, :] # -------------------------------- # crop corresponding H patch # -------------------------------- rnd_h_H, rnd_w_H = int(rnd_h * self.sf), int(rnd_w * self.sf) img_H = img_H[rnd_h_H:rnd_h_H + self.patch_size, rnd_w_H:rnd_w_H + self.patch_size, :] # -------------------------------- # augmentation - flip and/or rotate # -------------------------------- mode = np.random.randint(0, 8) img_L, img_H, img_h0 = util.augment_img( img_L, mode=mode), util.augment_img( img_H, mode=mode), util.augment_img(img_h0, mode=mode) # ------------------------------------ # L/H pairs, HWC to CHW, numpy to tensor # ------------------------------------ img_deblur = [img_L, img_h0] img_deblur = util.generate_pyramid(*img_deblur, n_scales=3) img_deblur = util.np2tensor(*img_deblur) img_ls = img_deblur[0] img_hs = img_deblur[1] img_L = img_deblur[0] img_H = util.single2tensor3(img_H) img_L0 = img_L[0] #print(img_L[0].shape,img_H.shape,img_ls[0].shape,img_hs[0].shape) return { 'L0': img_L0, 'L': img_L, 'H': img_H, 'L_path': L_path, 'H_path': H_path, 'ls': img_ls, 'hs': img_hs }
def main(): # -------------------------------- # let's start! # -------------------------------- utils_logger.logger_info('test_dpsr_real', log_path='test_dpsr_real.log') logger = logging.getLogger('test_dpsr_real') global arg arg = parser.parse_args() # basic setting # ================================================ sf = arg.sf show_img = False noise_level_img = 8. / 255. #testsets = '/home/share2/wutong/DPSR/testsets/test/' #im = '0000115_01031_d_0000082.jpg' # chip.png colour.png # if 'chip' in im: # noise_level_img = 8./255. # elif 'colour' in im: #noise_level_img = 0.5/255. use_srganplus = False if use_srganplus and sf == 4: model_prefix = 'DPSRGAN' save_suffix = 'dpsrgan' else: model_prefix = 'DPSR' save_suffix = 'dpsr' model_path = os.path.join('DPSR_models', model_prefix + 'x%01d.pth' % (sf)) iter_num = 15 # number of iterations n_channels = 3 # only color images, fixed # ================================================ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # -------------------------------- # (1) load trained model # -------------------------------- model = SRResNet(in_nc=4, out_nc=3, nc=96, nb=16, upscale=sf, act_mode='R', upsample_mode='pixelshuffle') model.load_state_dict(torch.load(model_path), strict=True) model.eval() for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) logger.info('Model path {:s}. Testing...'.format(model_path)) # -------------------------------- # (2) L_folder, E_folder # -------------------------------- # --1--> L_folder, folder of Low-quality images L_folder = os.path.join(arg.load) # L: Low quality # --2--> E_folder, folder of Estimated images E_folder = os.path.join(arg.save) util.mkdir(E_folder) logger.info(L_folder) # for im in os.listdir(os.path.join(L_folder)): # if (im.endswith('.jpg') or im.endswith('.bmp') or im.endswith('.png')) and 'kernel' not in im: # -------------------------------- # (3) load low-resolution image # -------------------------------- img_list = os.listdir(L_folder) for im in img_list: img_path, ext = os.path.splitext(im) img_name = img_path.split('/')[-1] img = util.imread_uint(os.path.join(L_folder, im), n_channels=n_channels) h, w = img.shape[:2] util.imshow(img, title='Low-resolution image') if show_img else None img = util.unit2single(img) # -------------------------------- # (4) load blur kernel # -------------------------------- # if os.path.exists(os.path.join(L_folder, img_name+'_kernel.mat')): # k = loadmat(os.path.join(L_folder, img_name+'.mat'))['kernel'] # k = k.astype(np.float64) # k /= k.sum() # elif os.path.exists(os.path.join(L_folder, img_name+'_kernel.png')): # k = cv2.imread(os.path.join(L_folder, img_name+'_kernel.png'), 0) # k = np.float64(k) # float64 ! # k /= k.sum() #else: k = utils_deblur.fspecial('gaussian', 5, 0.25) iter_num = 5 # -------------------------------- # (5) handle boundary # -------------------------------- img = utils_deblur.wrap_boundary_liu( img, utils_deblur.opt_fft_size( [img.shape[0] + k.shape[0] + 1, img.shape[1] + k.shape[1] + 1])) # -------------------------------- # (6) get upperleft, denominator # -------------------------------- upperleft, denominator = utils_deblur.get_uperleft_denominator(img, k) # -------------------------------- # (7) get rhos and sigmas # -------------------------------- rhos, sigmas = utils_deblur.get_rho_sigma(sigma=max( 0.255 / 255.0, noise_level_img), iter_num=iter_num) # -------------------------------- # (8) main iteration # -------------------------------- z = img rhos = np.float32(rhos) sigmas = np.float32(sigmas) for i in range(iter_num): logger.info('Iter: {:->4d}--> {}'.format(i + 1, im)) # -------------------------------- # step 1, Eq. (9) // FFT # -------------------------------- rho = rhos[i] if i != 0: z = util.imresize_np(z, 1 / sf, True) z = np.real( np.fft.ifft2((upperleft + rho * np.fft.fft2(z, axes=(0, 1))) / (denominator + rho), axes=(0, 1))) # -------------------------------- # step 2, Eq. (12) // super-resolver # -------------------------------- sigma = torch.from_numpy(np.array(sigmas[i])) img_L = util.single2tensor4(z) noise_level_map = torch.ones((1, 1, img_L.size(2), img_L.size(3)), dtype=torch.float).mul_(sigma) img_L = torch.cat((img_L, noise_level_map), dim=1) img_L = img_L.to(device) # with torch.no_grad(): z = model(img_L) z = util.tensor2single(z) # -------------------------------- # (9) img_E # -------------------------------- img_E = util.single2uint( z[:h * sf, :w * sf]) # np.uint8((z[:h*sf, :w*sf] * 255.0).round()) logger.info('saving: sf = {}, {}.'.format( sf, img_name + '_x{}'.format(sf) + ext)) util.imsave(img_E, os.path.join(E_folder, img_name + ext)) util.imshow(img_E, title='Recovered image') if show_img else None
def __getitem__(self, index): L_path = None # ------------------------------------ # get H image # ------------------------------------ H_path = self.paths_H[index] img_H = util.imread_uint(H_path, self.n_channels) img_H = util.uint2single(img_H) # ------------------------------------ # modcrop # ------------------------------------ img_H = util.modcrop(img_H, self.sf) # ------------------------------------ # get L image # ------------------------------------ if self.paths_L: # -------------------------------- # directly load L image # -------------------------------- L_path = self.paths_L[index] img_L = util.imread_uint(L_path, self.n_channels) img_L = util.uint2single(img_L) else: # -------------------------------- # sythesize L image via matlab's bicubic # -------------------------------- H, W = img_H.shape[:2] img_L = util.imresize_np(img_H, 1 / self.sf, True) # ------------------------------------ # if train, get L/H patch pair # ------------------------------------ if self.opt['phase'] == 'train': H, W, C = img_L.shape # -------------------------------- # randomly crop the L patch # -------------------------------- rnd_h = random.randint(0, max(0, H - self.L_size)) rnd_w = random.randint(0, max(0, W - self.L_size)) img_L = img_L[rnd_h:rnd_h + self.L_size, rnd_w:rnd_w + self.L_size, :] # -------------------------------- # crop corresponding H patch # -------------------------------- rnd_h_H, rnd_w_H = int(rnd_h * self.sf), int(rnd_w * self.sf) img_H = img_H[rnd_h_H:rnd_h_H + self.patch_size, rnd_w_H:rnd_w_H + self.patch_size, :] # -------------------------------- # augmentation - flip and/or rotate # -------------------------------- mode = np.random.randint(0, 8) img_L, img_H = util.augment_img(img_L, mode=mode), util.augment_img(img_H, mode=mode) # ------------------------------------ # L/H pairs, HWC to CHW, numpy to tensor # ------------------------------------ img_H, img_L = util.single2tensor3(img_H), util.single2tensor3(img_L) if L_path is None: L_path = H_path return {'L': img_L, 'H': img_H, 'L_path': L_path, 'H_path': H_path}
def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None): """ This is the degradation model of BSRGAN from the paper "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" ---------- img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf) sf: scale factor isp_model: camera ISP model Returns ------- img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] """ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 sf_ori = sf h1, w1 = img.shape[:2] img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop h, w = img.shape[:2] if h < lq_patchsize * sf or w < lq_patchsize * sf: raise ValueError(f'img size ({h1}X{w1}) is too small!') hq = img.copy() if sf == 4 and random.random() < scale2_prob: # downsample1 if np.random.rand() < 0.5: img = cv2.resize( img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])), interpolation=random.choice([1, 2, 3])) else: img = util.imresize_np(img, 1 / 2, True) img = np.clip(img, 0.0, 1.0) sf = 2 shuffle_order = random.sample(range(7), 7) idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) if idx1 > idx2: # keep downsample3 last shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[ idx2], shuffle_order[idx1] for i in shuffle_order: if i == 0: img = add_blur(img, sf=sf) elif i == 1: img = add_blur(img, sf=sf) elif i == 2: a, b = img.shape[1], img.shape[0] # downsample2 if random.random() < 0.75: sf1 = random.uniform(1, 2 * sf) img = cv2.resize( img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3])) else: k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) k_shifted = shift_pixel(k, sf) k_shifted = k_shifted / k_shifted.sum( ) # blur with shifted kernel img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror') img = img[0::sf, 0::sf, ...] # nearest downsampling img = np.clip(img, 0.0, 1.0) elif i == 3: # downsample3 img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) img = np.clip(img, 0.0, 1.0) elif i == 4: # add Gaussian noise img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) elif i == 5: # add JPEG noise if random.random() < jpeg_prob: img = add_JPEG_noise(img) elif i == 6: # add processed camera sensor noise if random.random() < isp_prob and isp_model is not None: with torch.no_grad(): img, hq = isp_model.forward(img.copy(), hq) # add final JPEG compression noise img = add_JPEG_noise(img) # random crop img, hq = random_crop(img, hq, sf_ori, lq_patchsize) return img, hq