def get_all_comparison_metrics(denoised,
source,
noisy=None,
return_title_string=False,
clamp=True):
metrics = {}
metrics['psnr'] = np.zeros(len(denoised))
metrics['ssim'] = np.zeros(len(denoised))
if noisy is not None:
metrics['psnr_delta'] = np.zeros(len(denoised))
metrics['ssim_delta'] = np.zeros(len(denoised))
if clamp:
denoised = torch.clamp(denoised, 0.0, 1.0)
metrics['psnr'] = psnr(source, denoised)
metrics['ssim'] = ssim(source, denoised)
if noisy is not None:
metrics['psnr_delta'] = metrics['psnr'] - psnr(source, noisy)
metrics['ssim_delta'] = metrics['ssim'] - ssim(source, noisy)
if return_title_string:
return convert_dict_to_string(metrics)
else:
return metrics
def gen_sr_img(self,
test_dst=Dataset('./test_image/'),
image_name='Spidy.jpg',
save=False,
verbose=0):
"""
Generate the high-resolution picture with trained model.
Input:
test_dst: Dataset, Instance of Dataset.
image_name : str, name of image.
save : Bool, whether to save the sr-image to local or not.
verbose, int.
Return:
orig_img, bicubic_img, sr_img and psnr of the hr_img and sr_img.
"""
stride = test_dst.stride
scale = test_dst.scale
downsample_flag = test_dst.downsample_flag
assert test_dst.slice_mode == 'normal', 'Cannot be merged if blocks are not completed. '
data, label, _, size_merge = test_dst._data_label_(image_name)
output = self.model.predict(data, verbose=verbose)
# merge all subimages.
hr_img = merge_to_whole(label, size_merge, stride=stride)
lr_img = downsample(hr_img,
scale=scale,
downsample_flag=downsample_flag)
sr_img = merge_to_whole(output, size_merge, stride=stride)
if verbose == 1:
print('PSNR is %f' % (psnr(sr_img / 255., hr_img / 255.)))
if save:
scipy.misc.imsave('./example/%s_SR.png' % (image_name), sr_img)
return (hr_img, lr_img, sr_img), psnr(sr_img / 255., hr_img / 255.)
def test_luma(model):
psnr_before = Averagvalue()
psnr_after = Averagvalue()
for i, (image, label, qpmap) in enumerate(testloader):
image, label, qpmap = image.cuda(), label.cuda(), qpmap.cuda()
outputs = model(torch.cat([image, qpmap], 1))
psnr_1 = psnr(F.mse_loss(image, label).item())
psnr_2 = psnr(F.mse_loss(outputs, label).item())
info = '[{}]'.format(i) + 'PSNR from {:.4f} to {:.4f}'.format(
psnr_1, psnr_2) + ' Delta:{:.4f}'.format(psnr_2 - psnr_1)
psnr_before.update(psnr_1)
psnr_after.update(psnr_2)
return psnr_after.avg - psnr_before.avg
def bilinear_upsampling(opt, dataloader, scale):
for batch_no, data in enumerate(dataloader['test']):
high_img, _ = data
inputs = torch.FloatTensor(opt.batchSize, 3, opt.imageSize,
opt.imageSize)
for j in range(opt.batchSize):
inputs[j] = scale(high_img[j])
high_img[j] = normalize(high_img[j])
outputs = F.upsample(inputs,
scale_factor=opt.upSampling,
mode='bilinear',
align_corners=True)
transform = transforms.Compose([
transforms.Normalize(mean=[-2.118, -2.036, -1.804],
std=[4.367, 4.464, 4.444]),
transforms.ToPILImage()
])
transform(outputs[0]).save('output/train/bilinear_fake/' +
str(batch_no) + '.png')
transform(high_img[0]).save('output/train/bilinear_real/' +
str(batch_no) + '.png')
# for output, himg in zip (outputs, high_img):
# psnr_val = psnr(output,himg)
#mssim = avg_msssim(himg, output)
print(psnr(un_normalize(outputs), un_normalize(high_img)))
def testPSNR(self, testdata, testlabel):
psnrsum = []
msesum = []
for i in range(len(testdata)):
Limage = testdata[i]
Himage = testlabel[i]
height, width, channel = Himage.shape
Limage = cv2.cvtColor(Limage, cv2.COLOR_BGR2YCrCb)
Limage = cv2.resize(Limage, (width, height),
interpolation=cv2.INTER_CUBIC)
net_input = np.reshape((Limage / 255)[:, :, 0],
(1, height, width, 1))
Cr = np.reshape((Limage / 255)[:, :, 1], (1, height, width, 1))
Cb = np.reshape((Limage / 255)[:, :, 2], (1, height, width, 1))
Y = self.session.run(self.output, feed_dict={self.data: net_input})
B = (Cb - 0.5) / 0.564 + Y
R = (Cr - 0.5) / 0.713 + Y
G = (Y - 0.299 * R - 0.114 * B) / 0.587
new_image = np.concatenate((B, G, R), axis=3)
new_image = new_image * 255
new_image = np.reshape(new_image, (height, width, channel))
psnr, mse = utils.psnr(new_image, Himage)
psnrsum.append(psnr)
msesum.append(mse / 255 / 255)
return np.mean(psnrsum), np.mean(msesum)
def evaluate(args, epoch, model, data_loader):
model.eval()
losses = []
start_epoch = time.time()
psnr_tot = []
with torch.no_grad():
for iter, data in enumerate(data_loader):
target, _ = data
input, mask = apply_random_mask(target, args['rate'])
target = torch.tensor(target).to(args['device'])
input = torch.tensor(input).to(args['device'])
mask = torch.tensor(mask).to(args['device'])
output = model(input, mask) #.squeeze(1)
loss = F.mse_loss(output, target, reduction='sum')
losses.append(loss.item())
psnr_tot.append(
np.mean([
psnr(t.cpu().numpy(),
o.cpu().numpy()) for t, o in zip(target, output)
]))
#if iter > 10:
# break
return np.mean(losses), np.mean(psnr_tot), time.time() - start_epoch
def test(net, test_load):
p_sum = 0
num = 0
with torch.no_grad():
for i, imgs in enumerate(test_load, 0):
imgs = imgs.to(main.device)
start_time = time.time()
re_imgs = net(imgs)
#re_imgs.cpu()
end_time = time.time()
time_cost = (end_time - start_time)/ 256
print("cost time:{0}".format(time_cost))
# restrict the value of images in 0 to 1
re_imgs = re_imgs.clamp(min=0, max=1)
psnr_batch = psnr(imgs, re_imgs)
if torch.isnan(psnr_batch).sum()>0 or torch.isinf(psnr_batch).sum()>0:
print(psnr_batch)
continue
p_sum += torch.sum(psnr_batch)
num += psnr_batch.size(0)
# imgs_display(imgs.cpu(), re_imgs.cpu())
mean_psnr = p_sum / num
return mean_psnr
def train_step(self, batch, batch_idx):
inputs, targets = batch
outputs = self.model(inputs)
loss = criterion(outputs, targets)
with torch.no_grad():
if batch_idx + 1 == len(train_loader):
indices = torch.randperm(inputs.shape[0])[:2]
baseline = F.interpolate(inputs[indices],
scale_factor=scale_factor,
mode='bicubic')
example = targets[indices], baseline, outputs[indices]
grid = make_grid(torch.cat(example, dim=-1)).clamp(0, 1)
train_writer.add_image('Random Images From Last Batch', grid,
self.current_epoch)
ssim = pytorch_msssim.ms_ssim(outputs, targets, data_range=1)
psnr = utils.psnr(outputs, targets)
size = len(targets)
self.loss.update(loss.item(), size)
self.psnr.update(psnr.item(), size)
self.ssim.update(ssim.item(), size)
return loss, {
'loss': f'{self.loss.compute():.4g}',
'psnr': f'{self.psnr.compute():.2f}',
'ssim': f'{self.ssim.compute():.4f}'
}
def ADMM_TV_rec(y, Phi, A, At, Phi_sum, maxiter, step_size, weight, row, col,
ColT, eta, X_ori):
#y1 = np.zeros((row,col))
begin_time = time.time()
theta = At(y, Phi)
v = theta
b = np.zeros((row, col, ColT))
for ni in range(maxiter):
yb = A(theta + b, Phi)
#y1 = y1+ (y-fb)
v = (theta + b) + np.multiply(
step_size, At(np.divide(y - yb, Phi_sum + eta), Phi))
#vmb = v-b
theta = denoise_tv_chambolle(v - b,
weight,
n_iter_max=30,
multichannel=True)
b = b - (v - theta)
weight = 0.999 * weight
eta = 0.998 * eta
if (ni + 1) % 5 == 0:
# mse = np.mean(np.sum((y-A(v,Phi))**2,axis=(0,1)))
end_time = time.time()
print("ADMM-TV: Iteration %3d, PSNR = %2.2f dB,"
" time = %3.1fs." %
(ni + 1, psnr(v, X_ori), end_time - begin_time))
return v
def manipImage(strImgName, lbd1, lbd2, seed=None):
"""
"""
# 1 Img read
mat_img = mpimg.imread(FOLDER_IMG + strImgName +'.png')
# 2 Apply noise
std_noise = PARAM_STDNOISE
mat_noise = np.random.normal(0, std_noise, mat_img.shape)
mat_Y = utils.im2Patches(mat_img + mat_noise, 8, 8, PARAM_LEARNINGOVERLAPPING)
# 3. Experiment
print('Debut manip')
cSva = cSvaIbpDl.CSvaDl()
cSva.run(mat_Y, PARAM_NBITER, lbd1, lbd2)
print('Fin manip')
matD_hat = cSva.get_dic()
matW_hat = cSva.get_coefs()
# print
psnr = utils.psnr(mat_img,
utils.patches2Ima(np.dot(matD_hat, matW_hat), mat_img.shape, 8, 8, PARAM_LEARNINGOVERLAPPING)
)
print('psnr ' + str(round(psnr, 2)) + ' dB')
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-f", default="images/createrLake.tiff") # image file
parser.add_argument("-q", default=90, type=int) # image quality
parser.add_argument("-p", default=100, type=int) # Percent of Downsampling
parser.add_argument("--plot", default=0, type=bool) # test suites
args = parser.parse_args()
fname = args.f
quality = args.q
percent = args.p
image = misc.imread(fname)
data = JPEG_compression(image,quality,percent)
utils.save_to_gzip(data,fname)
print("Original File Size = {0} B".format(utils.get_file_size(fname)))
print("New File Size = {0} B".format(utils.get_file_size(fname + ".gz")))
#Load File
data2 = utils.gzip_to_data(fname + ".gz")
im2 = JPEG_decompression(data2)
print("PSNR = {0}".format(utils.psnr(image,im2)))
if(args.plot):
plt.figure()
plt.subplot(1,2,1)
plt.imshow(image)
plt.subplot(1,2,2)
plt.imshow(im2)
plt.show()
def main(unused_argv):
# Load training and test data
train_labels = utils.load_dataset(DIR_TRAIN_LABELS)
print("reading train_labels done")
print(train_labels.shape)
train_inputs = utils.load_dataset(DIR_TRAIN_INPUTS)
print("reading train_inputs done")
print(train_inputs.shape)
test_labels = utils.load_dataset(DIR_TEST_LABELS)
print("reading test_labels done")
print(test_labels.shape)
test_inputs = utils.load_dataset(DIR_TEST_INPUTS)
print("reading test_inputs done")
print(test_inputs.shape)
# Create the Estimator
mnist_estimator = tf.estimator.Estimator(model_fn=cnn_model_fn)
# Train the model
train_input_fn = tf.estimator.inputs.numpy_input_fn(x={"x": train_inputs},
y=train_labels,
batch_size=100,
num_epochs=None,
shuffle=True)
mnist_estimator.train(input_fn=train_input_fn, steps=STEPS)
test_input_fn = tf.estimator.inputs.numpy_input_fn(x={"x": test_inputs},
shuffle=False)
preds = mnist_estimator.predict(input_fn=test_input_fn)
if not os.path.exists(os.path.dirname(DIR_TO_WRITE_RESULTS)):
os.makedirs(os.path.dirname(DIR_TO_WRITE_RESULTS))
count = 0
for inp in test_inputs:
count += 1
imageio.imwrite(DIR_TO_WRITE_RESULTS + "input" + str(count) + ".png",
inp.reshape((IMAGE_HEIGHT, IMAGE_WIDTH), order='C'))
count = 0
total_psnr = 0
for pred in preds:
count += 1
pred = pred.reshape((IMAGE_HEIGHT, IMAGE_WIDTH), order='C')
for i in range(pred.shape[0]):
for j in range(pred.shape[1]):
pred[i][j] = int(pred[i][j])
if (pred[i][j] < 0):
pred[i][j] = 0
else:
pred[i][j] = int(pred[i][j])
#pred[i][j] = (pred[i][j]-MIN)*255/(MAX-MIN)
#predicted_image = utils.fill_dark_part(test_inputs[count-1].reshape((28,28),order='C'), pred, 28, 28, 5)
total_psnr += utils.psnr(
test_labels[count - 1].reshape((IMAGE_HEIGHT, IMAGE_WIDTH),
order='C'), pred)
imageio.imwrite(DIR_TO_WRITE_RESULTS + str(count) + ".png", pred)
print("psnr = " + str(total_psnr / count))
def driver(args):
image = cv2.imread(args.input_image, 0)
psf = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(args.psf_diam, args.psf_diam))
psf = psf / np.sum(psf)
degraded = cv2.filter2D(image, -1, psf)
degraded = gaussian_noise(degraded, args.noise_sd)
if args.desired_spec == 'const':
max_psnr = -1 * float('inf')
optimal_k, best_restored = None, None
for k in range(0, 10000, 50):
restored = weiner_filter(degraded, psf, k, args.noise_sd**2)
p = psnr(image, restored)
print(f'[*] k: {k}\tPSNR: {p}')
if p > max_psnr:
max_psnr = p
optimal_k = k
best_restored = restored
print(f'[*] Optimal K: {optimal_k}, Max PSNR: {max_psnr}')
elif args.desired_spec == 'truth':
best_restored = weiner_filter(degraded, psf,
np.abs(np.fft.fft2(image)),
args.noise_sd**2)
print(f'[*] PSNR: {psnr(image, best_restored)}')
cv2.imshow('', np.hstack((image, degraded, best_restored)))
cv2.waitKey()
def manipImage(strImgName, lbd1, lbd2, seed=None):
"""
"""
# 1 Img read
mat_img = mpimg.imread(FOLDER_IMG + strImgName +'.png')
matImgMask = np.full(mat_img.shape, np.int32(1))
# 2 Apply noise
std_noise = PARAM_STDNOISE
mat_noise = np.random.normal(0, std_noise, mat_img.shape)
matYobs = utils.im2Patches(mat_img + mat_noise, 8, 8, PARAM_LEARNINGOVERLAPPING)
matPatchMask = utils.im2Patches(matImgMask + mat_noise, 8, 8, PARAM_LEARNINGOVERLAPPING)
# 3. Experiment
print('Debut manip')
cSvaInpainting = cSvaIbpDl.CSvaDlInpainting()
#cSvaInpainting.run(matYobs, matPatchMask.astype(np.int32), PARAM_NBITER, lbd1, lbd2)
print('Fin manip')
np.savetxt("test.csv", matYobs, delimiter=",", newline='\n')
matD_hat = cSvaInpainting.get_dic()
matW_hat = cSvaInpainting.get_coefs()
# print
psnr = utils.psnr(mat_img,
utils.patches2Ima(np.dot(matD_hat, matW_hat), mat_img.shape, 8, 8, PARAM_LEARNINGOVERLAPPING)
)
print('psnr ' + str(round(psnr, 2)) + ' dB')
def validation_step(self, batch, batch_idx):
inputs, targets = batch
outputs = self.model(inputs)
loss = criterion(outputs, targets)
if self.current_epoch is not None and batch_idx + 1 == len(val_loader):
baseline = F.interpolate(inputs,
scale_factor=scale_factor,
mode='bicubic')
grid = make_grid(torch.cat([targets, baseline, outputs], dim=-1),
nrow=2).clamp(0, 1)
val_writer.add_image('Images From Last Batch', grid,
self.current_epoch)
ssim = pytorch_msssim.ms_ssim(outputs, targets, data_range=1)
psnr = utils.psnr(outputs, targets)
size = len(targets)
self.val_loss.update(loss.item(), size)
self.val_psnr.update(psnr.item(), size)
self.val_ssim.update(ssim.item(), size)
return {
'val_loss': f'{self.val_loss.compute():.4g}',
'psnr': f'{self.val_psnr.compute():.2f}',
'ssim': f'{self.val_ssim.compute():.4f}'
}
def test(self, config):
test_data = yuv_data.YuvData(config.QP)
with tf.Session(config=tf.ConfigProto(
device_count={'cpu': 0})) as sess:
if config.blu:
if self.load(sess, config.blu_dir + str(config.QP)):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
exit(1)
else:
if self.load(sess, config.ppro_dir + str(config.QP)):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
exit(1)
#self.quant_w(sess)
for j in range(len(test_data.testData)):
ori = test_data.testData[j][1]
rec = test_data.testData[j][2]
output = np.empty((1, ori.shape[1], ori.shape[2]))
start_time = time.time()
for i in range(ori.shape[0]):
input = rec[i].reshape(1, ori.shape[1], ori.shape[2], 1)
if rec[i].shape[0] > 1500 or rec[i].shape[1] > 1500:
#print("memory limited, divided running")
pred_image = self.divided_run(sess, input, self.pred)
else:
pred_image = sess.run(self.pred,
feed_dict={self.images: input})
pred_image = pred_image * 255 + 128
output = np.concatenate(
(output,
pred_image.reshape(1, ori.shape[1], ori.shape[2])), 0)
run_time = time.time() - start_time
output = np.delete(output, 0, 0)
result = np.clip(output, 0, 255)
psnr_rec = utils.psnr(rec, ori)
psnr_net = utils.psnr(result, ori)
with open("psnr.data", "ab") as f:
f.write(struct.pack("<d", psnr_net))
with open("psnr_ori.data", "ab") as f:
f.write(struct.pack("<d", psnr_rec))
print(test_data.testData[j][0], test_data.testData[j][1].shape,
'time:%.3f' % run_time)
print("PSNR: before net %.3f\tafter net %.3f" %
(psnr_rec, psnr_net))
def test_chroma(model, testloader):
psnr_before = Averagvalue()
psnr_after = Averagvalue()
for i, (luma, chroma_rec, chroma_pad, chroma_gd) in enumerate(testloader):
luma, chroma_rec, chroma_pad, chroma_gd = luma.cuda(), chroma_rec.cuda(), chroma_pad.cuda(), chroma_gd.cuda()
outputs = model(luma, chroma_pad)
losslist = []
losslist.append(F.mse_loss(chroma_rec, chroma_gd).item())
loss = F.mse_loss(outputs, chroma_gd - chroma_rec)
losslist.append(loss.item())
info = '[{}]'.format(i) + 'PSNR from {:.4f} to {:.4f}'.format(psnr(losslist[0]), psnr(losslist[-1])) + ' Delta:{:.4f}'.format(psnr(losslist[-1])- psnr(losslist[0]))
psnr_before.update(psnr(losslist[0]))
psnr_after.update(psnr(losslist[-1]))
#print('PSNR from {:.4f} to {:.4f}'.format(psnr_before.avg, psnr_after.avg))
return psnr_after.avg - psnr_before.avg
def test_luma(model, testloader):
psnr_before = Averagvalue()
psnr_after = Averagvalue()
for i, (image, label) in enumerate(testloader):
image, label = image.cuda(), label.cuda()
outputs = model(image)
losslist = []
losslist.append(F.mse_loss(label, image).item())
loss = F.mse_loss(outputs, label)
losslist.append(loss.item())
info = '[{}]'.format(i) + 'PSNR from {:.4f} to {:.4f}'.format(psnr(losslist[0]), psnr(losslist[-1])) + ' Delta:{:.4f}'.format(psnr(losslist[-1])- psnr(losslist[0]))
psnr_before.update(psnr(losslist[0]))
psnr_after.update(psnr(losslist[-1]))
#print('PSNR from {:.4f} to {:.4f}'.format(psnr_before.avg, psnr_after.avg))
return psnr_after.avg - psnr_before.avg
def get_grid_psnrs(img, denoised_img, gt_img, grid_size=250):
psnr_list = []
psnr_noisy_list = []
for lower_x in range(0, img.shape[1], grid_size):
if lower_x + grid_size >= img.shape[1]:
lower_x = img.shape[1] - grid_size
upper_x = lower_x + grid_size
for lower_y in range(0, img.shape[0], grid_size):
if lower_y + grid_size >= img.shape[0]:
lower_y = img.shape[0] - grid_size
upper_y = lower_y + grid_size
crop = img[lower_y:upper_y, lower_x:upper_x]
denoised_crop = denoised_img[lower_y:upper_y, lower_x:upper_x]
gt_crop = gt_img[lower_y:upper_y, lower_x:upper_x]
psnr_list.append(psnr(denoised_crop, gt_crop))
psnr_noisy_list.append(psnr(crop, gt_crop))
return psnr_list, psnr_noisy_list
def train(trainloader, model, optimizer, epoch, save_dir):
global_step = epoch * len(trainloader) // args.print_freq
batch_time = Averagvalue()
loss_list = Averagvalue()
model.train()
end = time.time()
for i, (luma, chroma_rec, chroma_en, chroma_gd,
qpmap) in enumerate(trainloader):
luma, chroma_rec, chroma_en, chroma_gd, qpmap = luma.cuda(
), chroma_rec.cuda(), chroma_en.cuda(), chroma_gd.cuda(), qpmap.cuda()
outputs = model(torch.cat([chroma_en, qpmap], 1), luma)
psnr_1 = psnr(F.mse_loss(chroma_rec, chroma_gd).item())
psnr_2 = psnr(F.mse_loss(outputs, chroma_gd - chroma_rec).item())
loss = F.mse_loss(outputs, chroma_gd - chroma_rec)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_list.update(loss.item(), luma.size(0))
if i % args.print_freq == args.print_freq - 1:
info = 'Epoch: [{0}/{1}][{2}/{3}] '.format(epoch, args.maxepoch, i, len(trainloader)) + \
'Time {batch_time.val:.3f} (avg:{batch_time.avg:.3f})' .format(batch_time = batch_time) + \
'Loss {loss.val:f} (avg:{loss.avg:f})'.format(loss = loss_list) + ' PSNR {:.4f}'.format(psnr_2 - psnr_1)
print(info)
global_step += 1
writer = SummaryWriter(args.show_path)
writer.add_scalar('scalar/loss', loss_list.avg, global_step)
delta_psnr = test_chroma(model)
writer.add_scalar('scalar/psnr', delta_psnr, global_step)
loss_list.reset()
writer.close()
if not isdir(save_dir):
os.makedirs(save_dir)
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
},
filename=join(save_dir, "epoch-%d-checkpoint.pth" % epoch))
def test(self, config):
"""
Testing process.
"""
print("Testing...")
# Load checkpoint
if self.load(self.checkpoint_dir, config.scale):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
nx, ny = input_setup(self.sess, config)
data_dir = os.path.join('./{}'.format(config.checkpoint_dir),
"test.h5")
test_data, test_label = read_data(data_dir)
result = self.pred.eval({
self.images: test_data,
self.labels: test_label
})
result = merge(result, [nx, ny])
result = result.squeeze()
# Save output image
output_path = os.path.join(os.getcwd(), config.output_dir)
image_path = os.path.join(output_path, "test_img.png")
imsave(result, image_path)
# PSNR
label_path = os.path.join(output_path, "test_org_img.png")
bicubic_path = os.path.join(output_path, "test_bicubic_img.png")
bicubic_img = imread(bicubic_path, is_grayscale=True)
label_img = imread(label_path, is_grayscale=True)
output_img = imread(image_path, is_grayscale=True)
bicubic_psnr_value = psnr(label_img, bicubic_img)
srcnn_psnr_value = psnr(label_img, output_img)
print("Bicubic PSNR: [{}]".format(bicubic_psnr_value))
print("SRCNN PSNR: [{}]".format(srcnn_psnr_value))
def test_chroma(model):
psnr_before = Averagvalue()
psnr_after = Averagvalue()
for i, (luma, chroma_rec, chroma_en, chroma_gd,
qpmap) in enumerate(trainloader):
luma, chroma_rec, chroma_en, chroma_gd, qpmap = luma.cuda(
), chroma_rec.cuda(), chroma_en.cuda(), chroma_gd.cuda(), qpmap.cuda()
outputs = model(torch.cat([chroma_en, qpmap], 1), luma)
psnr_1 = psnr(F.mse_loss(chroma_rec, chroma_gd).item())
psnr_2 = psnr(F.mse_loss(outputs, chroma_gd - chroma_rec).item())
info = '[{}]'.format(i) + 'PSNR from {:.4f} to {:.4f}'.format(
psnr_1, psnr_2) + ' Delta:{:.4f}'.format(psnr_2 - psnr_1)
psnr_before.update(psnr_1)
psnr_after.update(psnr_2)
return psnr_after.avg - psnr_before.avg
def get_psnr(fake, real):
cpu = torch.device("cpu")
psnr_list = []
for i in range(len(fake)):
np_fake = fake[i].to(cpu).detach().clone().numpy().transpose([1, 2, 0])
np_real = real[i].to(cpu).detach().clone().numpy().transpose([1, 2, 0])
psnr_list.append(psnr(np_fake, np_real))
return statistics.mean(psnr_list)
def predict(images, session=None, model=None, model_name=None, targets=None):
session_passed = session is not None
if not session_passed:
session = tf.Session()
if model is None:
model = load_model(model_name, session)
predictions = []
if targets is not None:
psnr = []
for i in range(len(images)):
image = images[i].copy()
assert str(image.dtype).startswith('float') or str(
image.dtype).startswith('uint')
if image.dtype == 'uint8':
image = image / 255.0
elif image.dtype == 'uint16':
image = image / 65535.0
if len(image.shape) == 2:
image = np.expand_dims(rgb2gray(image), axis=2)
prediction = model.outputs.eval(
feed_dict={model.inputs: np.array([image])}, session=session)[0]
if targets is not None:
target = targets[i].copy()
assert str(target.dtype).startswith('float') or str(
target.dtype).startswith('uint')
if target.dtype == 'uint8':
target = target / 255.0
elif target.dtype == 'uint16':
target = target / 65535.0
psnr.append(utils.psnr(prediction, target, maximum=1.0))
predictions.append(prediction)
if not session_passed:
session.close()
if targets is not None:
return predictions, psnr
else:
return predictions
def predict(images, session=None, network=None, targets=None, border=0):
session_passed = session is not None
if not session_passed:
session = tf.Session()
if network is None:
network = load_model(session)
predictions = []
if targets is not None:
psnr = []
for i in range(len(images)):
image = images[i]
if len(image.shape) == 3:
image_ycbcr = color.rgb2ycbcr(image)
image_y = image_ycbcr[:, :, 0]
else:
image_y = image.copy()
image_y = image_y.astype(np.float) / 255
reshaped_image_y = np.array([np.expand_dims(image_y, axis=2)])
prediction = network.output.eval(feed_dict={network.input: reshaped_image_y}, session=session)[0]
prediction *= 255
if targets is not None:
if len(targets[i].shape) == 3:
target_y = color.rgb2ycbcr(targets[i])[:, :, 0]
else:
target_y = targets[i].copy()
psnr.append(utils.psnr(prediction[border:-border, border:-border, 0],
target_y[border:-border, border:-border], maximum=255.0))
if len(image.shape) == 3:
prediction = color.ycbcr2rgb(np.concatenate((prediction, image_ycbcr[:, :, 1:3]), axis=2)) * 255
else:
prediction = prediction[:, :, 0]
prediction = np.clip(prediction, 0, 255).astype(np.uint8)
predictions.append(prediction)
if not session_passed:
session.close()
if targets is not None:
return predictions, psnr
else:
return predictions
def train(self, epoch):
for i, (noisy, image) in enumerate(self.dataloader):
noisy = Variable(noisy)
image = Variable(image)
if self.args.cuda:
noisy = noisy.cuda()
image = image.cuda()
denoised = self.net(noisy)
loss = (denoised - image).pow(2).mean()
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
cur_loss = float(loss.data)
if cur_loss < self.best_loss:
self.best_loss = cur_loss
self.best_state_dict = self.net.state_dict()
if i % self.args.log_interval == 0:
noisy_psnr = psnr(image.data, noisy.data)
denoised_psnr = psnr(image.data, denoised.data)
print('{}: {}, {}: {:.6f}, {}: {:.6f}, {}: {:.6f}.'.format(
'Train epoch', epoch, 'loss', cur_loss, 'noisy psnr',
noisy_psnr, 'denoised psnr', denoised_psnr))
# save current training results
path = os.path.join(self.args.model_dir,
'image_{}_{}.jpg'.format(epoch, i))
print('Saving image {}'.format(path))
cat_tensor = torch.cat([image.data, noisy.data, denoised.data])
save_image(cat_tensor, path, nrow=self.args.batch_size)
# save best model
if cur_loss != self.best_loss:
self.save_model(epoch, i)
def validation(sess, vdsr, epoch, scale):
if not os.path.exists('./validation'):
os.makedirs('./validation')
validation_result_path = {
2: './validation/2.csv',
3: './validation/3.csv',
4: './validation/4.csv'
}
s = scale
if not os.path.exists('./validation/%d' % s):
os.makedirs('./validation/%d' % s)
lr, gt = data.load_lr_gt_mat('./data/test_data/mat/Set5', s)
v_len = len(gt)
psnr = []
for i in range(v_len):
lr_image = lr[i]['data']
gt_image = gt[i]['data']
residual, sr = sess.run([vdsr.residual, vdsr.inference],
feed_dict={
vdsr.lr:
lr_image.reshape((1, ) + lr_image.shape +
(1, ))
})
sr = sr.reshape(sr.shape[1:3])
residual = residual.reshape(residual.shape[1:3])
utils.save_image(
sr, './validation/%d/%s_sr_scale_%d_epoch_%d.png' %
(s, lr[i]['name'], s, epoch))
residual = utils.normalize(residual)
utils.save_image(
residual, './validation/%d/%s_residual_scale_%d_epoch_%d.png' %
(s, lr[i]['name'], s, epoch))
sr_ = utils.shave(sr, s)
gt_image_ = utils.shave(gt_image, s)
psnr.append(utils.psnr(gt_image_, sr_))
with open(validation_result_path[s], 'a') as f:
f.write('%d, %s, %f\n' %
(epoch, ', '.join(str(e) for e in psnr), float(np.mean(psnr))))
def for_display(black, rg, output, dir): if not os.path.exists(dir): os.makedirs(dir) output_reshape = np.reshape(output, [FLAGS.picture_size, FLAGS.picture_size, 2]) output_con = np.zeros([FLAGS.picture_size, FLAGS.picture_size, 3]) output_con[:, :, 0] = black output_con[:, :, 1:] = output_reshape save_LABimage(output_con , dir + "/output_color.jpg") output_con[:, :, 0] = 0.5 save_LABimage(output_con , dir + "/ab.jpg") real_con = np.zeros([FLAGS.picture_size, FLAGS.picture_size, 3]) real_con[:, :, 0] = black real_con[:, :, 1:] = rg save_LABimage(real_con , dir + "/real_color.jpg") save_image(black , dir + "/black_white.jpg") logging.info(dir) logging.info(utils.psnr(output_con * 255, real_con * 255))
def GAP_TV_rec(y,Phi,A, At,Phi_sum, maxiter, step_size, weight, row, col, ColT, X_ori):
y1 = np.zeros((row,col))
begin_time = time.time()
f = At(y,Phi)
for ni in range(maxiter):
fb = A(f,Phi)
y1 = y1+ (y-fb)
f = f + np.multiply(step_size, At( np.divide(y1-fb,Phi_sum),Phi ))
f = denoise_tv_chambolle(f, weight,n_iter_max=30,multichannel=True)
if (ni+1)%5 == 0:
# mse = np.mean(np.sum((y-A(f,Phi))**2,axis=(0,1)))
end_time = time.time()
print("GAP-TV: Iteration %3d, PSNR = %2.2f dB,"
" time = %3.1fs."
% (ni+1, psnr(f, X_ori), end_time-begin_time))
return f
def validate(scale):
dbs = []
for i in range(len(eval_gt_imgs)):
hr = eval_gt_imgs[i]
hr = modular_crop(hr, scale)
h, w, _ = hr.shape
lr = resize_img(hr, h // scale, w // scale)
upcubic = resize_img(lr, h, w)
inputs = np.asarray([lr], np.float32), np.asarray([upcubic],
np.float32), scale
output = net(inputs)
y_pred = clip255(unnormalize(output))
dbs.append(psnr(y_pred, hr))
return np.mean(dbs)
for d in range(Kd):
ki = k[i,:,:,d]
tmp[d,:,:] += conv2(phi_u,ki[::-1,::-1],'full')
u_t[t+1] = np.clip(u_t[t] - crop_zero(tmp, padding, padding) - l*bwd_bayer(fwd_bayer(u_t[t]) - f), 0.0,255.0)
print '.',
#Evaluate
print "\nTest image: %d" % data_config['indices'][example]
#get the result
result = u_t[num_steps]
plt.figure(1)
plt.imshow(swapimdims_3HW_HW3(result).astype('uint8'), interpolation="none")
plt.show()
target = data.target[example]
#compute psnr and ssim on the linear space result image
print "PSNR linear: %.2f dB" % psnr(target, np.round(result), 255.0)
print "SSIM linear: %.3f" % ssim(target, result, 255.0)
#also compute psnr and ssim on the sRGB transformed result image
srgb_params = init_colortransformation_gamma()
result_rgb = apply_colortransformation_gamma(np.expand_dims(result,0), srgb_params)
target_rgb = apply_colortransformation_gamma(np.expand_dims(target,0), srgb_params)
print "PSNR sRGB: %.2f dB" % psnr(target_rgb[0], result_rgb[0], 255.0)
print "SSIM sRGB: %.3f" % ssim(target_rgb[0], result_rgb[0], 255.0)
#save result
plt.imsave("results/" + str(data_config['indices'][example]), swapimdims_3HW_HW3(result).astype('uint8'))