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