def _validate(self, model, dataset): validation_dataset = dataset criterion = nn.MSELoss() losses = [] total_targets, total_predictions = [], [] for images in tqdm(validation_dataset, desc="validating", ncols=150): images = images.cuda() predictions = model.eval()(images, train_=False) loss = criterion(images, predictions) losses.append(loss.item()) for i in range(predictions.size()[0]): total_targets.append(images[i].cpu().detach().numpy().reshape( 32, 32, 3)) total_predictions.append( predictions[i].cpu().detach().numpy().reshape(32, 32, 3)) # calculate ssim ssim = calculate_ssim(total_targets, total_predictions) # calculate loss loss = np.mean(losses) return loss, ssim
def calculate_psnr_ssim_ESRGAN(): dir = "/home/jakaria/Super_Resolution/Filter_Enhance_Detect/saved_ESRGAN/val_images/*/*" HR_DIR = "/home/jakaria/Super_Resolution/Datasets/COWC/DetectionPatches_256x256/Potsdam_ISPRS/HR/x4/valid_img/" img_SR = sorted(glob.glob(dir + '_300000.png')) psnr_SR = 0 ssim_SR = 0 total = len(img_SR) print(total) i = 0 for im_SR in img_SR: print(os.path.basename(im_SR) + '--') im_gt = os.path.basename(im_SR) im_gt = im_gt.rsplit('_', 1)[0] + ".jpg" im_gt = os.path.join(HR_DIR, im_gt) print(im_gt) image_SR = cv2.imread(im_SR) image_SR = cv2.cvtColor(image_SR, cv2.COLOR_BGR2RGB) cv2.imwrite(im_SR, image_SR) image_gt = cv2.imread(im_gt) image_SR = cv2.imread(im_SR) psnr_SR += calculate_psnr(image_gt, image_SR) ssim_SR += calculate_ssim(image_gt, image_SR) i += 1 print(i) avg_psnr_SR = psnr_SR / total avg_ssim_SR = ssim_SR / total text_file = open( "/home/jakaria/Super_Resolution/Filter_Enhance_Detect/saved_ESRGAN/Output.txt", "a") print("SR PSNR: %4.2f" % avg_psnr_SR) text_file.write("SR PSNR: %4.2f \n" % avg_psnr_SR) print("SR SSIM: %5.4f" % avg_ssim_SR) text_file.write("SR SSIM: %5.4f \n" % avg_ssim_SR)
def train(self, continue_: bool = False): model = Model().cuda() if continue_: model.load_state_dict(torch.load(self.model_path)) criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=self.lr) train_loss, train_ssim, val_loss, val_ssim = [], [], [], [] for epoch in range(1, (self.epochs + 1)): epoch_loss = [] epoch_train_targets, epoch_train_predictions = [], [] for images in tqdm(self.train_set, desc="epoch", ncols=150): optimizer.zero_grad() images = images.float().cuda() predictions = model.train()(images, train_=True) loss = criterion(predictions, images) loss.backward() optimizer.step() epoch_loss.append(loss.item()) for i in range(predictions.size()[0]): epoch_train_targets.append( images[i].cpu().detach().numpy()) \ epoch_train_predictions.append(predictions[i].cpu().detach().numpy()) current_val_loss, current_val_ssim = self._validate( model, self.validation_set) current_train_loss = np.mean(epoch_loss) current_train_ssim = calculate_ssim(epoch_train_targets, epoch_train_predictions) show_progress(self.epochs, epoch, current_train_loss, current_train_ssim, current_val_loss, current_val_ssim) torch.save(model.state_dict(), self.model_path)
def evaluate(model, update_step, writer, bucket, engine): device = torch.device('cuda:0') model.eval() eval_paths = [os.path.join(args.eval_path, v) for v in ['Set14', 'Set5']] metrics_list = [] unnorm = UnNormalize(0.5, 0.5) for eval_path in eval_paths: eval_name = os.path.basename(eval_path) HQ_path = os.path.join(eval_path, eval_name) + '.lmdb' LQ_path = os.path.join(eval_path, eval_name) + '_LQ.lmdb' LQ_r_path = os.path.join(eval_path, eval_name) + '_LQ_restored.lmdb' eval_set = ValDataset(HQ_path, LQ_path, LQ_r_path, args.scale) eval_loader = DataLoader( eval_set, batch_size=1, shuffle=False, num_workers=4) psrn_rgb = 0.0 psrn_y = 0.0 ssim_rgb = 0.0 ssim_y = 0.0 for i, data_dict in enumerate(eval_loader): img_HQ = data_dict['img_GT'] img_LQ = data_dict['img_LQ'].to(device) img_LQ_r = data_dict['img_LQ_r'] with torch.no_grad(): # SR image range [-1, 1] img_SR = model(img_LQ) # SR image range [0, 1] img_SR = unnorm(img_SR) if i == 0: imgs = torch.cat([img_HQ, img_SR.detach().cpu(), img_LQ_r], dim=0) grid = vutils.make_grid(imgs, nrow=3, normalize=False) tmp_image = T.ToPILImage()(grid) tmp_image.save('images/tmp_image.png') upload_to_cloud(bucket, 'images/tmp_image.png', 'odesr01_04/image_progress/{}/gen_step_{}'. format(eval_name, update_step * args.update_freq)) if eval_name == 'Set5': writer.add_image('Set5', grid, update_step) crop_size = args.scale img_HQ_rgb = img_HQ[0].permute(2, 1, 0).cpu(). \ numpy()[crop_size:-crop_size, crop_size:-crop_size, :] img_SR_rgb = img_SR[0].permute(2, 1, 0).detach().cpu(). \ numpy()[crop_size:-crop_size, crop_size:-crop_size, :] img_HQ_y = rgb2ycbcr(img_HQ_rgb) img_SR_y = rgb2ycbcr(img_SR_rgb) psrn_rgb += calculate_psnr(img_HQ_rgb * 255, img_SR_rgb * 255) psrn_y += calculate_psnr(img_HQ_y * 255, img_SR_y * 255) ssim_rgb += calculate_ssim(img_HQ_rgb * 255, img_SR_rgb * 255) ssim_y += calculate_ssim(img_HQ_y * 255, img_SR_y * 255) psrn_rgb = psrn_rgb / len(eval_loader.dataset) psrn_y = psrn_y / len(eval_loader.dataset) ssim_rgb = ssim_rgb / len(eval_loader.dataset) ssim_y = ssim_y / len(eval_loader.dataset) metrics_list.extend([psrn_rgb, psrn_y, ssim_rgb, ssim_y]) if eval_name == 'Set5': writer.add_scalar('psrn_rgb', psrn_rgb, update_step) writer.add_scalar('psrn_y', psrn_y, update_step) writer.add_scalar('ssim_rgb', ssim_rgb, update_step) writer.add_scalar('ssim_y', ssim_y, update_step) query = ''' INSERT INTO odesr01_04_val (set14_psnr_rgb, set14_psnr_y, set14_ssim_rgb, set14_ssim_y, set5_psnr_rgb, set5_psnr_y, set5_ssim_rgb, set5_ssim_y) VALUES (%f, %f, %f, %f, %f, %f, %f, %f) ''' % tuple(metrics_list) engine.execute(query) model.train()
def admm_denoise(y, Phi, A, At, _lambda=1, gamma=0.01, denoiser='tv', iter_max=50, noise_estimate=True, sigma=None, tv_weight=0.1, tv_iter_max=5, multichannel=True, x0=None, X_orig=None, show_iqa=True): ''' Alternating direction method of multipliers (ADMM)[1]-based denoising regularization for snapshot compressive imaging (SCI). Parameters ---------- y : two-dimensional (2D) ndarray of ints, uints or floats Input single measurement of the snapshot compressive imager (SCI). Phi : three-dimensional (3D) ndarray of ints, uints or floats, omitted Input sensing matrix of SCI with the third dimension as the time-variant, spectral-variant, volume-variant, or angular-variant masks, where each mask has the same pixel resolution as the snapshot measurement. Phi_sum : 2D ndarray Sum of the sensing matrix `Phi` along the third dimension. A : function Forward model of SCI, where multiple encoded frames are collapsed into a single measurement. At : function Transpose of the forward model. proj_meth : {'admm' or 'gap'}, optional Projection method of the data term. Alternating direction method of multipliers (ADMM)[1] and generalizedv alternating projection (GAP)[2] are used, where ADMM for noisy data, especially real data and GAP for noise-free data. gamma : float, optional Parameter in the ADMM projection, where more noisy measurements require greater gamma. denoiser : string, optional Denoiser used as the regularization imposing on the prior term of the reconstruction. _lambda : float, optional Regularization factor balancing the data term and the prior term, where larger `_lambda` imposing more constrains on the prior term. iter_max : int or uint, optional Maximum number of iterations. accelerate : boolean, optional Enable acceleration in GAP. noise_estimate : boolean, optional Enable noise estimation in the denoiser. sigma : one-dimensional (1D) ndarray of ints, uints or floats Input noise standard deviation for the denoiser if and only if noise estimation is disabled(i.e., noise_estimate==False). The scale of sigma is [0, 255] regardless of the the scale of the input measurement and masks. tv_weight : float, optional weight in total variation (TV) denoising. x0 : 3D ndarray Start point (initialized value) for the iteration process of the reconstruction. Returns ------- x : 3D ndarray Reconstructed 3D scene captured by the SCI system. References ---------- .. [1] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein, "Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers," Foundations and Trends® in Machine Learning, vol. 3, no. 1, pp. 1-122, 2011. .. [2] X. Yuan, "Generalized alternating projection based total variation minimization for compressive sensing," in IEEE International Conference on Image Processing (ICIP), 2016, pp. 2539-2543. .. [3] Y. Liu, X. Yuan, J. Suo, D. Brady, and Q. Dai, "Rank Minimization for Snapshot Compressive Imaging," IEEE Transactions on Pattern Analysis and Machine Intelligence, doi:10.1109/TPAMI.2018.2873587, 2018. Code credit ----------- Xin Yuan, Bell Labs, [email protected], created Aug 7, 2018. Yang Liu, Tsinghua University, [email protected], updated Jan 22, 2019. See Also -------- gap_denoise ''' # [0] initialization if x0 is None: x0 = At(y, Phi) # default start point (initialized value) if not isinstance(sigma, list): sigma = [sigma] if not isinstance(iter_max, list): iter_max = [iter_max] * len(sigma) # [1] start iteration for reconstruction x = x0 # initialization theta = x0 Phi_sum = np.sum(Phi, 2) Phi_sum[Phi_sum == 0] = 1 b = np.zeros_like(x0) psnr_all = [] ssim_all = [] k = 0 device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu') model = net() model.load_state_dict( torch.load( r'/home/dgl/zhengsiming/self_train/check_points/best_smallsigma.pth' )) model.eval() for q, v in model.named_parameters(): v.requires_grad = False model = model.to(device) for idx, nsig in enumerate(sigma): # iterate all noise levels for it in range(iter_max[idx]): # Euclidean projection yb = A(theta + b, Phi) x = (theta + b) + _lambda * (At( (y - yb) / (Phi_sum + gamma), Phi)) # ADMM x1 = shift_back(x - b, step=2) #x1=x-b # switch denoiser if denoiser.lower() == 'tv': # total variation (TV) denoising #theta = denoise_tv_chambolle(x1, nsig/255, n_iter_max=tv_iter_max, multichannel=multichannel) theta = TV_denoiser(x1, tv_weight, n_iter_max=tv_iter_max) elif denoiser.lower() == 'wavelet': # wavelet denoising if noise_estimate or nsig is None: # noise estimation enabled theta = denoise_wavelet(x1, multichannel=multichannel) else: theta = denoise_wavelet(x1, sigma=nsig, multichannel=multichannel) elif denoiser.lower() == 'vnlnet': # Video Non-local net denoising theta = vnlnet(np.expand_dims((x1).transpose(2, 0, 1), 3), nsig) theta = np.transpose(theta.squeeze(3), (1, 2, 0)) elif denoiser.lower() == 'hsicnn': if k >= 89: tem = None for i in range(28): net_input = None if i < 3: if i == 0: net_input = np.dstack( (x1[:, :, i], x1[:, :, i], x1[:, :, i], x1[:, :, i:i + 4])) elif i == 1: net_input = np.dstack( (x1[:, :, i - 1], x1[:, :, i - 1], x1[:, :, i - 1], x1[:, :, i:i + 4])) elif i == 2: net_input = np.dstack( (x1[:, :, i - 2], x1[:, :, i - 2], x1[:, :, i - 1], x1[:, :, i:i + 4])) net_input = torch.from_numpy( np.ascontiguousarray(net_input)).permute( 2, 0, 1).float().unsqueeze(0) net_input = net_input.to(device) Nsigma = torch.full((1, 1, 1, 1), 10 / 255.).type_as(net_input) output = model(net_input, Nsigma) output = output.data.squeeze().float().cpu().numpy( ) if i == 0: tem = output else: tem = np.dstack((tem, output)) elif i > 24: if i == 25: net_input = np.dstack( (x1[:, :, i - 3:i + 1], x1[:, :, i + 1], x1[:, :, i + 2], x1[:, :, i + 2])) elif i == 26: net_input = np.dstack( (x1[:, :, i - 3:i + 1], x1[:, :, i + 1], x1[:, :, i + 1], x1[:, :, i + 1])) elif i == 27: net_input = np.dstack( (x1[:, :, i - 3:i + 1], x1[:, :, i], x1[:, :, i], x1[:, :, i])) net_input = torch.from_numpy( np.ascontiguousarray(net_input)).permute( 2, 0, 1).float().unsqueeze(0) net_input = net_input.to(device) Nsigma = torch.full((1, 1, 1, 1), 10 / 255.).type_as(net_input) output = model(net_input, Nsigma) output = output.data.squeeze().float().cpu().numpy( ) tem = np.dstack((tem, output)) else: net_input = x1[:, :, i - 3:i + 4] net_input = torch.from_numpy( np.ascontiguousarray(net_input)).permute( 2, 0, 1).float().unsqueeze(0) net_input = net_input.to(device) Nsigma = torch.full((1, 1, 1, 1), 10 / 255.).type_as(net_input) output = model(net_input, Nsigma) output = output.data.squeeze().float().cpu().numpy( ) tem = np.dstack((tem, output)) theta = tem else: #print('theta:', np.max(theta)) theta = denoise_tv_chambolle(x1, tv_weight, n_iter_max=tv_iter_max, multichannel=multichannel) else: raise ValueError('Unsupported denoiser {}!'.format(denoiser)) # [optional] calculate image quality assessment, i.e., PSNR for # every five iterations if show_iqa and X_orig is not None: psnr_all.append(psnr(X_orig, theta)) ssim_all.append(calculate_ssim(X_orig, theta)) if (k + 1) % 1 == 0: if not noise_estimate and nsig is not None: if nsig < 1: print( ' ADMM-{0} iteration {1: 3d}, sigma {2: 3g}/255, ' 'PSNR {3:2.2f} dB.'.format( denoiser.upper(), k + 1, nsig * 255, psnr_all[k]), 'SSIM:{}'.format(ssim_all[k])) else: print( ' ADMM-{0} iteration {1: 3d}, sigma {2: 3g}, ' 'PSNR {3:2.2f} dB.'.format( denoiser.upper(), k + 1, nsig, psnr_all[k]), 'SSIM:{}'.format(ssim_all[k])) else: print( ' ADMM-{0} iteration {1: 3d}, ' 'PSNR {2: 2.2f} dB.'.format( denoiser.upper(), k + 1, psnr_all[k]), 'SSIM:{}'.format(ssim_all[k])) theta = shift(theta, step=2) b = b - (x - theta) # update residual k = k + 1 return theta, psnr_all, ssim_all
def gap_denoise(y, Phi, A, At, _lambda=1, accelerate=True, denoiser='tv', iter_max=50, noise_estimate=True, sigma=None, tv_weight=0.1, tv_iter_max=5, multichannel=True, x0=None, X_orig=None, model=None, show_iqa=True): ''' Alternating direction method of multipliers (ADMM)[1]-based denoising regularization for snapshot compressive imaging (SCI). Parameters ---------- y : two-dimensional (2D) ndarray of ints, uints or floats Input single measurement of the snapshot compressive imager (SCI). Phi : three-dimensional (3D) ndarray of ints, uints or floats, omitted Input sensing matrix of SCI with the third dimension as the time-variant, spectral-variant, volume-variant, or angular-variant masks, where each mask has the same pixel resolution as the snapshot measurement. Phi_sum : 2D ndarray, Sum of the sensing matrix `Phi` along the third dimension. A : function Forward model of SCI, where multiple encoded frames are collapsed into a single measurement. At : function Transpose of the forward model. proj_meth : {'admm' or 'gap'}, optional Projection method of the data term. Alternating direction method of multipliers (ADMM)[1] and generalizedv alternating projection (GAP)[2] are used, where ADMM for noisy data, especially real data and GAP for noise-free data. gamma : float, optional Parameter in the ADMM projection, where more noisy measurements require greater gamma. denoiser : string, optional Denoiser used as the regularization imposing on the prior term of the reconstruction. _lambda : float, optional Regularization factor balancing the data term and the prior term, where larger `_lambda` imposing more constrains on the prior term. iter_max : int or uint, optional Maximum number of iterations. accelerate : boolean, optional Enable acceleration in GAP. noise_estimate : boolean, optional Enable noise estimation in the denoiser. sigma : one-dimensional (1D) ndarray of ints, uints or floats Input noise standard deviation for the denoiser if and only if noise estimation is disabled(i.e., noise_estimate==False). The scale of sigma is [0, 255] regardless of the the scale of the input measurement and masks. tv_weight : float, optional weight in total variation (TV) denoising. x0 : 3D ndarray Start point (initialized value) for the iteration process of the reconstruction. model : pretrained model for image/video denoising. Returns ------- x : 3D ndarray Reconstructed 3D scene captured by the SCI system. References ---------- .. [1] X. Liao, H. Li, and L. Carin, "Generalized Alternating Projection for Weighted-$\ell_{2,1}$ Minimization with Applications to Model-Based Compressive Sensing," SIAM Journal on Imaging Sciences, vol. 7, no. 2, pp. 797-823, 2014. .. [2] X. Yuan, "Generalized alternating projection based total variation minimization for compressive sensing," in IEEE International Conference on Image Processing (ICIP), 2016, pp. 2539-2543. .. [3] Y. Liu, X. Yuan, J. Suo, D. Brady, and Q. Dai, "Rank Minimization for Snapshot Compressive Imaging," IEEE Transactions on Pattern Analysis and Machine Intelligence, doi:10.1109/TPAMI.2018.2873587, 2018. Code credit ----------- Xin Yuan, Bell Labs, [email protected], created Aug 7, 2018. Yang Liu, Tsinghua University, [email protected], updated Jan 22, 2019. See Also -------- admm_denoise ''' # [0] initialization if x0 is None: print(At) x0 = At(y, Phi) # default start point (initialized value) if not isinstance(sigma, list): sigma = [sigma] if not isinstance(iter_max, list): iter_max = [iter_max] * len(sigma) y1 = np.zeros_like(y) Phi_sum = np.sum(Phi, 2) Phi_sum[Phi_sum == 0] = 1 # [1] start iteration for reconstruction x = x0 # initialization psnr_all = [] ssim_all = [] k = 0 device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu') model = net() model.load_state_dict(torch.load(r'./check_points/deep_denoiser.pth')) model.eval() for q, v in model.named_parameters(): v.requires_grad = False model = model.to(device) for idx, nsig in enumerate(sigma): # iterate all noise levels for it in range(iter_max[idx]): #print('max1_{0}_{1}:'.format(idx,it),np.max(x)) yb = A(x, Phi) if accelerate: # accelerated version of GAP y1 = y1 + (y - yb) x = x + _lambda * (At((y1 - yb) / Phi_sum, Phi)) # GAP_acc else: x = x + _lambda * (At((y - yb) / Phi_sum, Phi)) # GAP x = shift_back(x, step=1) # switch denoiser if denoiser.lower() == 'tv': # total variation (TV) denoising x = denoise_tv_chambolle(x, nsig / 255, n_iter_max=tv_iter_max, multichannel=multichannel) #x= TV_denoiser(x, tv_weight, n_iter_max=tv_iter_max) elif denoiser.lower() == 'hsicnn': l_ch = 10 m_ch = 10 h_ch = 10 if (k > 123 and k <= 125) or (k >= 119 and k <= 121) or ( k >= 115 and k <= 117 ) or (k >= 111 and k <= 113) or (k >= 107 and k <= 109) or ( k >= 103 and k <= 105) or (k >= 99 and k <= 101) or ( k >= 95 and k <= 97) or (k >= 91 and k <= 93) or ( k >= 87 and k <= 89) or (k >= 83 and k <= 85): tem = None for i in range(31): net_input = None if i < 3: ori_nsig = nsig if i == 0: net_input = np.dstack( (x[:, :, i], x[:, :, i], x[:, :, i], x[:, :, i:i + 4])) elif i == 1: net_input = np.dstack( (x[:, :, i - 1], x[:, :, i - 1], x[:, :, i - 1], x[:, :, i:i + 4])) elif i == 2: net_input = np.dstack( (x[:, :, i - 2], x[:, :, i - 2], x[:, :, i - 1], x[:, :, i:i + 4])) net_input = torch.from_numpy( np.ascontiguousarray(net_input)).permute( 2, 0, 1).float().unsqueeze(0) net_input = net_input.to(device) Nsigma = torch.full((1, 1, 1, 1), l_ch / 255.).type_as(net_input) output = model(net_input, Nsigma) output = output.data.squeeze().cpu().numpy() if k < 0: output = denoise_tv_chambolle( x[:, :, i], nsig / 255, n_iter_max=tv_iter_max, multichannel=False) nsig = ori_nsig if i == 0: tem = output else: tem = np.dstack((tem, output)) elif i > 27: ori_nsig = nsig if k >= 45: nsig /= 1 if i == 28: net_input = np.dstack( (x[:, :, i - 3:i + 1], x[:, :, i + 1], x[:, :, i + 2], x[:, :, i + 2])) elif i == 29: net_input = np.dstack( (x[:, :, i - 3:i + 1], x[:, :, i + 1], x[:, :, i + 1], x[:, :, i + 1])) elif i == 30: net_input = np.dstack( (x[:, :, i - 3:i + 1], x[:, :, i], x[:, :, i], x[:, :, i])) net_input = torch.from_numpy( np.ascontiguousarray(net_input)).permute( 2, 0, 1).float().unsqueeze(0) net_input = net_input.to(device) Nsigma = torch.full((1, 1, 1, 1), m_ch / 255.).type_as(net_input) output = model(net_input, Nsigma) output = output.data.squeeze().cpu().numpy() if k < 0: output = denoise_tv_chambolle( x[:, :, i], 10 / 255, n_iter_max=tv_iter_max, multichannel=False) tem = np.dstack((tem, output)) nsig = ori_nsig else: ori_nsig = nsig net_input = x[:, :, i - 3:i + 4] net_input = torch.from_numpy( np.ascontiguousarray(net_input)).permute( 2, 0, 1).float().unsqueeze(0) net_input = net_input.to(device) Nsigma = torch.full((1, 1, 1, 1), h_ch / 255.).type_as(net_input) output = model(net_input, Nsigma) output = output.data.squeeze().cpu().numpy() tem = np.dstack((tem, output)) nsig = ori_nsig #x = np.clip(tem,0,1) x = tem else: x = denoise_tv_chambolle(x, nsig / 255, n_iter_max=tv_iter_max, multichannel=multichannel) #x = TV_denoiser(x, tv_weight, n_iter_max=tv_iter_max) elif denoiser.lower() == 'bm3d': sigma = nsig / 255 v = np.zeros((15, 15)) for x1 in range(-7, 8, 1): for x2 in range(-7, 8, 1): v[x1 + 7, x2 + 7] = 1 / (x1**2 + x2**2 + 1) v = v / np.sum(v) for i in range(28): x[:, :, i] = bm3d_deblurring(np.atleast_3d(x[:, :, i]), sigma, v) else: raise ValueError('Unsupported denoiser {}!'.format(denoiser)) # [optional] calculate image quality assessment, i.e., PSNR for # every five iterations if show_iqa and X_orig is not None: ssim_all.append(calculate_ssim(X_orig, x)) psnr_all.append(psnr(X_orig, x)) if (k + 1) % 1 == 0: if not noise_estimate and nsig is not None: if nsig < 1: print( ' GAP-{0} iteration {1: 3d}, sigma {2: 3g}/255, ' 'PSNR {3:2.2f} dB.'.format( denoiser.upper(), k + 1, nsig * 255, psnr_all[k]), 'SSIM:{}'.format(ssim_all[k])) else: print( ' GAP-{0} iteration {1: 3d}, sigma {2: 3g}, ' 'PSNR {3:2.2f} dB.'.format( denoiser.upper(), k + 1, nsig, psnr_all[k]), 'SSIM:{}'.format(ssim_all[k])) else: print( ' GAP-{0} iteration {1: 3d}, ' 'PSNR {2:2.2f} dB.'.format(denoiser.upper(), k + 1, psnr_all[k]), 'SSIM:{}'.format(ssim_all[k])) x = shift(x, step=1) if k == 123: break k = k + 1 return x, psnr_all
def calculate_psnr_ssim(): dir = "/home/jakaria/Super_Resolution/Filter_Enhance_Detect/saved/" HR_DIR = "/home/jakaria/Super_Resolution/Datasets/COWC/DetectionPatches_256x256/Potsdam_ISPRS/HR/x4/valid_img/*" bicubic_DIR = "/home/jakaria/Super_Resolution/Datasets/COWC/DetectionPatches_256x256/Potsdam_ISPRS/Bic/x4/valid_img/*" img_GT = sorted(glob.glob(HR_DIR + '.jpg')) img_final_SR_enhanced_1 = sorted( glob.glob(dir + '/enhanced_SR_images_1/*.png')) img_final_SR_enhanced_2 = sorted( glob.glob(dir + '/enhanced_SR_images_2/*.png')) img_final_SR_enhanced_3 = sorted( glob.glob(dir + '/enhanced_SR_images_3/*.png')) img_final_SR = sorted(glob.glob(dir + '/final_SR_images_216000/*.png')) img_SR = sorted(glob.glob(dir + '/SR_images/*.png')) img_SR_combined = sorted( glob.glob(dir + '/combined_SR_images_216000/*.png')) img_Bic = sorted(glob.glob(bicubic_DIR + '.jpg')) psnr_enhanced_1 = 0 psnr_enhanced_2 = 0 psnr_enhanced_3 = 0 psnr_final = 0 psnr_SR = 0 psnr_SR_combined = 0 psnr_Bic = 0 ssim_enhanced_1 = 0 ssim_enhanced_2 = 0 ssim_enhanced_3 = 0 ssim_final = 0 ssim_SR = 0 ssim_SR_combined = 0 ssim_Bic = 0 total = len(img_SR) print(total) i = 0 for im_gt, im_enhanced_1, im_enhanced_2, im_enhanced_3, im_final, im_SR, \ im_SR_combined, im_Bic in zip(img_GT, img_final_SR_enhanced_1, img_final_SR_enhanced_2, img_final_SR_enhanced_3, img_final_SR, img_SR, img_SR_combined, img_Bic): print( os.path.basename(im_gt) + '--', os.path.basename(im_enhanced_1) + '--', os.path.basename(im_enhanced_2) + '--', os.path.basename(im_enhanced_3) + '--', os.path.basename(im_final) + '--', os.path.basename(im_SR) + '--', os.path.basename(im_SR_combined) + '--', os.path.basename(im_Bic)) image_gt = cv2.imread(im_gt) image_enhanced_1 = cv2.imread(im_enhanced_1) image_enhanced_2 = cv2.imread(im_enhanced_2) image_enhanced_3 = cv2.imread(im_enhanced_3) image_final = cv2.imread(im_final) image_SR = cv2.imread(im_SR) image_SR_combined = cv2.imread(im_SR_combined) image_Bic = cv2.imread(im_Bic) psnr_enhanced_1 += calculate_psnr(image_gt, image_enhanced_1) psnr_enhanced_2 += calculate_psnr(image_gt, image_enhanced_2) psnr_enhanced_3 += calculate_psnr(image_gt, image_enhanced_3) psnr_final += calculate_psnr(image_gt, image_final) psnr_SR += calculate_psnr(image_gt, image_SR) psnr_SR_combined += calculate_psnr(image_gt, image_SR_combined) psnr_Bic += calculate_psnr(image_gt, image_Bic) ssim_enhanced_1 += calculate_ssim(image_gt, image_enhanced_1) ssim_enhanced_2 += calculate_ssim(image_gt, image_enhanced_2) ssim_enhanced_3 += calculate_ssim(image_gt, image_enhanced_3) ssim_final += calculate_ssim(image_gt, image_final) ssim_SR += calculate_ssim(image_gt, image_SR) ssim_SR_combined += calculate_ssim(image_gt, image_SR_combined) ssim_Bic += calculate_ssim(image_gt, image_Bic) i += 1 print(i) avg_psnr_enhanced_1, avg_psnr_enhanced_2, avg_psnr_enhanced_3, avg_psnr_final, \ avg_psnr_SR, avg_psnr_SR_combined, avg_psnr_Bic = (psnr_enhanced_1 / total, psnr_enhanced_2 / total, psnr_enhanced_3 / total, psnr_final / total, psnr_SR / total, psnr_SR_combined / total, psnr_Bic / total) avg_ssim_enhanced_1, avg_ssim_enhanced_2, avg_ssim_enhanced_3, avg_ssim_final, \ avg_ssim_SR, avg_ssim_SR_combined, avg_ssim_Bic = (ssim_enhanced_1 / total, ssim_enhanced_2 / total, ssim_enhanced_3 / total, ssim_final / total, ssim_SR / total, ssim_SR_combined / total, ssim_Bic / total) text_file = open( "/home/jakaria/Super_Resolution/Filter_Enhance_Detect/saved/Output_216000.txt", "a") print("Enhanced PSNR_1: %4.2f" % avg_psnr_enhanced_1) text_file.write("Enhanced PSNR_1: %4.2f \n" % avg_psnr_enhanced_1) print("Enhanced PSNR_2: %4.2f" % avg_psnr_enhanced_2) text_file.write("Enhanced PSNR_2: %4.2f \n" % avg_psnr_enhanced_2) print("Enhanced PSNR_3: %4.2f" % avg_psnr_enhanced_3) text_file.write("Enhanced PSNR_3: %4.2f \n" % avg_psnr_enhanced_3) print("Final PSNR: %4.2f" % avg_psnr_final) text_file.write("Final PSNR: %4.2f \n" % avg_psnr_final) print("SR PSNR: %4.2f" % avg_psnr_SR) text_file.write("SR PSNR: %4.2f \n" % avg_psnr_SR) print("SR PSNR_combined: %4.2f" % avg_psnr_SR_combined) text_file.write("SR PSNR_combined: %4.2f \n" % avg_psnr_SR_combined) print("Bic PSNR: %4.2f" % avg_psnr_Bic) text_file.write("Bic PSNR: %4.2f \n" % avg_psnr_Bic) print("Enhanced SSIM_1: %5.4f" % avg_ssim_enhanced_1) text_file.write("Enhanced SSIM_1: %5.4f \n" % avg_ssim_enhanced_1) print("Enhanced SSIM_2: %5.4f" % avg_ssim_enhanced_2) text_file.write("Enhanced SSIM_2: %5.4f \n" % avg_ssim_enhanced_2) print("Enhanced SSIM_3: %5.4f" % avg_ssim_enhanced_3) text_file.write("Enhanced SSIM_3: %5.4f \n" % avg_ssim_enhanced_3) print("Final SSIM: %5.4f" % avg_ssim_final) text_file.write("Final SSIM: %5.4f \n" % avg_ssim_final) print("SR SSIM: %5.4f" % avg_ssim_SR) text_file.write("SR SSIM: %5.4f \n" % avg_ssim_SR) print("SR SSIM_combined: %5.4f" % avg_ssim_SR_combined) text_file.write("SR SSIM_combined: %5.4f \n" % avg_ssim_SR_combined) print("Bic SSIM: %5.4f" % avg_ssim_Bic) text_file.write("Bic SSIM: %5.4f \n" % avg_ssim_Bic) text_file.close()
sr_img = sr_img / 255.0 crop_border = opt["crop_border"] if opt["crop_border"] else opt["scale"] if crop_border == 0: cropped_sr_img = sr_img cropped_gt_img = gt_img else: cropped_sr_img = sr_img[ crop_border:-crop_border, crop_border:-crop_border, : ] cropped_gt_img = gt_img[ crop_border:-crop_border, crop_border:-crop_border, : ] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results["psnr"].append(psnr) test_results["ssim"].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) if crop_border == 0: cropped_sr_img_y = sr_img_y cropped_gt_img_y = gt_img_y else: cropped_sr_img_y = sr_img_y[ crop_border:-crop_border, crop_border:-crop_border ] cropped_gt_img_y = gt_img_y[
def test(test_loader, net, scale, scene_name): train_mode = False net.eval() count = 0 PSNR = 0 SSIM = 0 PSNR_t = 0 SSIM_t = 0 out = [] for image_num, data in enumerate(test_loader): x_input, target = data[0], data[1] # print(x_input.shape) B, _, _, _, _ = x_input.shape with torch.no_grad(): x_input = Variable(x_input).cuda() target = Variable(target).cuda() t0 = time.time() # prediction = net(x_input) # ensamble test if True: # x_input = Variable(x_input).cuda() x_input_list = [x_input] for tf in ['v', 'h', 't']: x_input_list.extend( [test_ensamble(t, tf) for t in x_input_list]) prediction_list = [net(aug) for aug in x_input_list] for i in range(len(prediction_list)): if i > 3: prediction_list[i] = test_ensamble_2( prediction_list[i], 't') if i % 4 > 1: prediction_list[i] = test_ensamble_2( prediction_list[i], 'h') if (i % 4) % 2 == 1: prediction_list[i] = test_ensamble_2( prediction_list[i], 'v') prediction_cat = torch.cat(prediction_list, dim=0) prediction = prediction_cat.mean(dim=0, keepdim=True) torch.cuda.synchronize() t1 = time.time() print("===> Timer: %.4f sec." % (t1 - t0)) prediction = prediction.unsqueeze(2) count += 1 prediction = prediction.squeeze(0).permute(1, 2, 3, 0) # [T,H,W,C] prediction = prediction.cpu().numpy( )[:, :, :, ::-1] # tensor -> numpy, rgb -> bgr target = target.squeeze(0).permute(1, 2, 3, 0) # [T,H,W,C] target = target.cpu().numpy()[:, :, :, :: -1] # tensor -> numpy, rgb -> bgr save_img(prediction[0], scene_name, 2) # test_Y______________________ # prediction_Y = bgr2ycbcr(prediction[0]) # target_Y = bgr2ycbcr(target[0]) # prediction_Y = prediction_Y * 255 # target_Y = target_Y * 255 # test_RGB _______________________________ prediction_Y = prediction[0] * 255 target_Y = target[0] * 255 # ________________________________ # calculate PSNR and SSIM print('PSNR: {:.6f} dB, \tSSIM: {:.6f}'.format( calculate_psnr(prediction_Y, target_Y), calculate_ssim(prediction_Y, target_Y))) PSNR += calculate_psnr(prediction_Y, target_Y) SSIM += calculate_ssim(prediction_Y, target_Y) out.append(calculate_psnr(prediction_Y, target_Y)) print('===>{} PSNR = {}'.format(scene_name, PSNR)) print('===>{} SSIM = {}'.format(scene_name, SSIM)) PSNR_t = PSNR SSIM_t = SSIM return PSNR_t, SSIM_t