def main():
# Experiment specifications
imagename = 'image_Lena512rgb.png'
# Load noise-free image
y = np.array(Image.open(imagename)) / 255
# Possible noise types to be generated 'gw', 'g1', 'g2', 'g3', 'g4', 'g1w',
# 'g2w', 'g3w', 'g4w'.
noise_type = 'g3'
noise_var = 0.02 # Noise variance
seed = 0 # seed for pseudorandom noise realization
# Generate noise with given PSD
noise, psd, kernel = get_experiment_noise(noise_type, noise_var, seed,
y.shape)
# N.B.: For the sake of simulating a more realistic acquisition scenario,
# the generated noise is *not* circulant. Therefore there is a slight
# discrepancy between PSD and the actual PSD computed from infinitely many
# realizations of this noise with different seeds.
# Generate noisy image corrupted by additive spatially correlated noise
# with noise power spectrum PSD
z = np.atleast_3d(y) + np.atleast_3d(noise)
# Call BM3D With the default settings.
y_est = bm3d_rgb(z, psd)
# To include refiltering:
# y_est = bm3d_rgb(z, psd, 'refilter');
# For other settings, use BM3DProfile.
# profile = BM3DProfile(); # equivalent to profile = BM3DProfile('np');
# profile.gamma = 6; # redefine value of gamma parameter
# y_est = bm3d_rgb(z, psd, profile);
# Note: For white noise, you may instead of the PSD
# also pass a standard deviation
# y_est = bm3d_rgb(z, sqrt(noise_var));
# If the different channels have varying PSDs, you can supply a MxNx3 PSD or a list of 3 STDs:
# y_est = bm3d_rgb(z, np.concatenate((psd1, psd2, psd3), 2))
# y_est = bm3d_rgb(z, [sigma1, sigma2, sigma3])
psnr = get_psnr(y, y_est)
print("PSNR:", psnr)
# PSNR ignoring 16-pixel wide borders (as used in the paper), due to refiltering potentially leaving artifacts
# on the pixels near the boundary of the image when noise is not circulant
psnr_cropped = get_cropped_psnr(y, y_est, [16, 16])
print("PSNR cropped:", psnr_cropped)
# Ignore values outside range for display (or plt gives an error for multichannel input)
y_est = np.minimum(np.maximum(y_est, 0), 1)
z_rang = np.minimum(np.maximum(z, 0), 1)
plt.title("y, z, y_est")
plt.imshow(np.concatenate((y, np.squeeze(z_rang), y_est), axis=1))
plt.show()
def bm3d(img, sigma=None):
if sigma == None:
sigma = est_sigma(img)
if sigma < 4:
sigma **= 2
print("Used sigma = %.2f" % sigma)
try:
tic = time()
dumb = bm3d_rgb(img, sigma)
tock = time()
print("Execute time = %.2f" % (tock - tic), "s")
except:
print("Something went wrong, trying different method.")
tic = time()
dumb = cv.xphoto.bm3dDenoising(img)
tock = time()
print("Execute time = %.2f" % (tock - tic), "s")
dumb = np.clip(dumb, 0.0, 255.0).astype('uint8')
return dumb
def generate_red_img(noisy_img, pred_psd, gt_img):
noisy_img = np.array(noisy_img)
noisy_img = noisy_img / 255.0
gt_img = np.array(gt_img)
gt_img = gt_img / 255.0
profile = BM3DProfile()
profile.bs_ht = random.choice([2, 4, 8])
profile.transform_2d_wiener_name = random.choice(['dct', 'dst'])
profile.bs_wiener = random.choice([4, 5, 6, 7, 8, 9, 10, 11, 12])
cspace = random.choice(['opp', 'YCbCr'])
cff = random.uniform(1, 15)
red_img = bm3d_rgb(noisy_img,
cff * pred_psd[0],
profile,
colorspace=cspace)
red_img = np.minimum(np.maximum(red_img, 0), 1)
psnr = get_psnr(gt_img, red_img)
red_img = Image.fromarray(np.uint8(red_img * 255.0))
return red_img, cff, profile.bs_ht, cspace, profile.transform_2d_wiener_name, profile.bs_wiener, psnr
def generate_red_img(noisy_img, pred_psd, gt_img):
noisy_img = np.array(noisy_img)
noisy_img = noisy_img / 255.0
gt_img = np.array(gt_img)
gt_img = gt_img / 255.0
profile = BM3DProfile()
profile.bs_ht = 4
profile.transform_2d_wiener_name = 'dst'
profile.bs_wiener = choose_neighborhood
cspace = 'YCbCr'
cff = choose_cff
red_img = bm3d_rgb(noisy_img,
cff * pred_psd[0],
profile,
colorspace=cspace)
red_img = np.minimum(np.maximum(red_img, 0), 1)
psnr = get_psnr(gt_img, red_img)
red_img = Image.fromarray(np.uint8(red_img * 255.0))
return red_img, cff, profile.bs_ht, cspace, profile.transform_2d_wiener_name, profile.bs_wiener, psnr
def proxg(x, rho):
return bm3d_rgb(x, bsigma)
"""
Define proximal operator of (implicit, unknown) regularisation term for PPP problem. In this case we use BM3D :cite:`dabov-2008-image` as the denoiser, using the [code](https://pypi.org/project/bm3d) released with :cite:`makinen-2019-exact`.
"""
bsigma = 6.1e-2 # Denoiser parameter
def proxg(x, rho):
return bm3d_rgb(x, bsigma)
"""
Construct a baseline solution and initaliser for the PPP solution by BM3D denoising of a simple bilinear demosaicing solution. The `3 * nsigma` denoising parameter for BM3D is chosen empirically for best performance.
"""
imgb = bm3d_rgb(demosaic(sn), 3 * nsigma)
"""
Set algorithm options for PPP solver, including use of bilinear demosaiced solution as an initial solution.
"""
opt = PPP.Options({
'Verbose': True,
'RelStopTol': 1e-3,
'MaxMainIter': 12,
'rho': 1.8e-1,
'Y0': imgb
})
"""
Create solver object and solve, returning the the demosaiced image ``imgp``.
"""
"""
Define proximal operator of (implicit, unknown) regularisation term for PPP problem. In this case we use BM3D :cite:`dabov-2008-image` as the denoiser, using the [code](https://pypi.org/project/bm3d) released with :cite:`makinen-2019-exact`.
"""
bsigma = 6.1e-2 # Denoiser parameter
def proxg(x, rho):
return bm3d_rgb(x, bsigma)
"""
Construct a baseline solution and initaliser for the PPP solution by BM3D denoising of a simple bilinear demosaicing solution. The `3 * nsigma` denoising parameter for BM3D is chosen empirically for best performance.
"""
imgb = bm3d_rgb(bilinear_demosaic(sn), 3 * nsigma)
"""
Set algorithm options for PPP solver, including use of bilinear demosaiced solution as an initial solution.
"""
opt = PPP.Options({
'Verbose': True,
'RelStopTol': 1e-3,
'MaxMainIter': 12,
'rho': 1.8e-1,
'Y0': imgb
})
"""
Create solver object and solve, returning the the demosaiced image ``imgp``.
"""
def main(t, sigma, args=None):
# Experiment specifications
#imagename = 'image_Lena512rgb.png'
imagepath = args.path + '\\'
imagename = imagepath + 'ref\\' + t + '_r.bmp'
# Load noise-free image
y = np.array(Image.open(imagename)) / 255
# Possible noise types to be generated 'gw', 'g1', 'g2', 'g3', 'g4', 'g1w',
# 'g2w', 'g3w', 'g4w'.
noise_type = 'gw'
noise_var = (sigma / 255)**2 # Noise variance 25 std
seed = 0 # seed for pseudorandom noise realization
# Generate noise with given PSD
noise, psd, kernel = get_experiment_noise(noise_type, noise_var, seed,
y.shape)
# N.B.: For the sake of simulating a more realistic acquisition scenario,
# the generated noise is *not* circulant. Therefore there is a slight
# discrepancy between PSD and the actual PSD computed from infinitely many
# realizations of this noise with different seeds.
# Generate noisy image corrupted by additive spatially correlated noise
# with noise power spectrum PSD
z = np.atleast_3d(y) + np.atleast_3d(noise)
z_rang = np.minimum(np.maximum(z, 0), 1)
noisyimagename = imagepath + 'noisy\\' + t + '_g.bmp'
plt.imsave(noisyimagename, z_rang)
z = np.array(Image.open(noisyimagename)) / 255
# Call BM3D With the default settings.
y_est = bm3d_rgb(z, psd)
# To include refiltering:
# y_est = bm3d_rgb(z, psd, 'refilter');
# For other settings, use BM3DProfile.
# profile = BM3DProfile(); # equivalent to profile = BM3DProfile('np');
# profile.gamma = 6; # redefine value of gamma parameter
# y_est = bm3d_rgb(z, psd, profile);
# Note: For white noise, you may instead of the PSD
# also pass a standard deviation
# y_est = bm3d_rgb(z, sqrt(noise_var));
# If the different channels have varying PSDs, you can supply a MxNx3 PSD or a list of 3 STDs:
# y_est = bm3d_rgb(z, np.concatenate((psd1, psd2, psd3), 2))
# y_est = bm3d_rgb(z, [sigma1, sigma2, sigma3])
psnr = get_psnr(y, y_est)
print("PSNR:", psnr)
# PSNR ignoring 16-pixel wide borders (as used in the paper), due to refiltering potentially leaving artifacts
# on the pixels near the boundary of the image when noise is not circulant
psnr_cropped = get_cropped_psnr(y, y_est, [16, 16])
print("PSNR cropped:", psnr_cropped)
# Ignore values outside range for display (or plt gives an error for multichannel input)
y_est = np.minimum(np.maximum(y_est, 0), 1)
z_rang = np.minimum(np.maximum(z, 0), 1)
opath = 'bm3d_{0}.bmp'.format(t)
plt.imsave(opath.format(t), y_est)
y_est = np.array(Image.open(opath.format(t))) / 255
psnr = get_psnr(y, y_est)
print("PSNR 2:", psnr)
mse = ((y_est - y)**2).mean() * 255
print("MSE:", mse)
#plt.imsave(imagepath+ 'noisy\\' + t + '_g.bmp', z_rang)
# TEST CV2 PSNR
try:
from skimage.metrics import structural_similarity as compare_ssim
except Exception:
from skimage.measure import compare_ssim
import cv2
opath = 'bm3d_{0}.bmp'.format(t)
argref = imagename # path of noisy image
d = cv2.imread(opath)
tref = cv2.imread(argref)
(score, diff) = compare_ssim(tref, d, full=True, multichannel=True)
psnr2 = cv2.PSNR(tref, d)
print('#######################')
print('CV2 PSNR, SSIM: {:.2f}, {:.2f}'.format(psnr2, score))
print('#######################')
print('')
#plt.title("y, z, y_est")
#plt.imshow(np.concatenate((y, np.squeeze(z_rang), y_est), axis=1))
#plt.show()
return psnr, mse
def denoising(noise_im, clean_im, LR=1e-2, sigma=3, rho=1, eta=0.5, total_step=30,
prob1_iter=500, noise_level=None, result_root=None, f=None):
input_depth = 3
latent_dim = 3
en_net = UNet(input_depth, latent_dim).to(device)
de_net = UNet(latent_dim, input_depth).to(device)
parameters = [p for p in en_net.parameters()] + [p for p in de_net.parameters()]
optimizer = torch.optim.Adam(parameters, lr=LR)
l2_loss = torch.nn.MSELoss().cuda()
i0 = np_to_torch(noise_im).to(device)
noise_im_torch = np_to_torch(noise_im).to(device)
i0_til_torch = np_to_torch(noise_im).to(device)
Y = torch.zeros_like(noise_im_torch).to(device)
diff_original_np = noise_im.astype(np.float32) - clean_im.astype(np.float32)
diff_original_name = 'Original_dis.png'
save_hist(diff_original_np, result_root+diff_original_name)
best_psnr = 0
for i in range(total_step):
################################# sub-problem 1 ###############################
for i_1 in range(prob1_iter):
optimizer.zero_grad()
mean = en_net(noise_im_torch)
z = sample_z(mean)
out = de_net(z)
total_loss = 0.5 * l2_loss(out, noise_im_torch)
total_loss += 0.5 * (1/sigma**2)*l2_loss(mean, i0)
total_loss += (rho/2) * l2_loss(i0 + Y, i0_til_torch)
total_loss.backward()
optimizer.step()
with torch.no_grad():
i0 = ((1/sigma**2)*mean.detach() + rho*(i0_til_torch - Y)) / ((1/sigma**2) + rho)
with torch.no_grad():
################################# sub-problem 2 ###############################
i0_np = torch_to_np(i0)
Y_np = torch_to_np(Y)
sig = eval_sigma(i+1, noise_level)
i0_til_np = bm3d.bm3d_rgb(i0_np.transpose(1, 2, 0) + Y_np.transpose(1, 2, 0), sig).transpose(2, 0, 1)
i0_til_torch = np_to_torch(i0_til_np).to(device)
################################# sub-problem 3 ###############################
Y = Y + eta * (i0 - i0_til_torch)
###############################################################################
Y_name = 'Y_{:04d}'.format(i) + '.png'
i0_name = 'i0_num_epoch_{:04d}'.format(i) + '.png'
mean_name = 'Latent_im_num_epoch_{:04d}'.format(i) + '.png'
out_name = 'res_of_dec_num_epoch_{:04d}'.format(i) + '.png'
diff_name = 'Latent_dis_num_epoch_{:04d}'.format(i) + '.png'
Y_np = torch_to_np(Y)
Y_norm_np = np.sqrt((Y_np*Y_np).sum(0))
save_heatmap(Y_norm_np, result_root + Y_name)
save_torch(mean, result_root + mean_name)
save_torch(out, result_root + out_name)
save_torch(i0, result_root + i0_name)
mean_np = torch_to_np(mean)
diff_np = mean_np - clean_im
save_hist(diff_np, result_root + diff_name)
i0_til_np = torch_to_np(i0_til_torch).clip(0, 255)
psnr = compare_psnr(clean_im.transpose(1, 2, 0), i0_til_np.transpose(1, 2, 0), 255)
ssim = compare_ssim(clean_im.transpose(1, 2, 0), i0_til_np.transpose(1, 2, 0), multichannel=True, data_range=255)
i0_til_pil = np_to_pil(i0_til_np)
i0_til_pil.save(os.path.join(result_root, '{}'.format(i) + '.png'))
print('Iteration: {:02d}, VAE Loss: {:f}, PSNR: {:f}, SSIM: {:f}'.format(i, total_loss.item(), psnr, ssim), file=f, flush=True)
if best_psnr < psnr:
best_psnr = psnr
best_ssim = ssim
else:
break
return i0_til_np, best_psnr, best_ssim
def denoising_fixd_noise_level(noise_im,
clean_im,
LR=1e-2,
sigma=5,
rho=1,
eta=0.5,
alpha=0.5,
total_step=19,
prob1_iter=1000,
noise_level=None,
result_root=None,
fo=None):
sig = noise_level
r_bm3d = bm3d.bm3d_rgb(noise_im.transpose(1, 2, 0), sig)
r_bm3d = np.clip(r_bm3d, 0, 255)
psnr_bm3d = compare_psnr(clean_im.transpose(1, 2, 0), r_bm3d, 255)
ssim_bm3d = compare_ssim(r_bm3d,
clean_im.transpose(1, 2, 0),
multichannel=True,
data_range=255)
print('noise level {} '.format(noise_level), file=fo, flush=True)
print('PSNR_BM3D: {}, SSIM_BM3D: {}'.format(psnr_bm3d, ssim_bm3d),
file=fo,
flush=True)
r_bm3d = Image.fromarray(r_bm3d.astype(np.uint8))
r_bm3d.save(result_root + 'bm3d_result.png')
input_depth = 3
latent_dim = 3
en_net = UNet(input_depth, latent_dim, need_sigmoid=True).cuda()
de_net = UNet(latent_dim, input_depth, need_sigmoid=True).cuda()
en_optimizer = torch.optim.Adam(en_net.parameters(), lr=LR)
de_optimizer = torch.optim.Adam(de_net.parameters(), lr=LR)
l2_loss = torch.nn.MSELoss().cuda()
noise_im = noise_im / 255.0
i0 = torch.Tensor(noise_im[None, ...]).cuda()
noise_im_torch = torch.Tensor(noise_im)[None, ...].cuda()
Y = torch.zeros_like(noise_im_torch).cuda()
i0_til_torch = torch.Tensor(noise_im[None, ...]).cuda()
output = None
for i in range(total_step):
############################### sub-problem 1 #################################
for i_1 in range(prob1_iter):
mean_i = en_net(noise_im_torch)
eps = mean_i.clone().normal_()
out = de_net(mean_i + eps)
total_loss = 0.5 * l2_loss(out, noise_im_torch)
total_loss += 1 / (2 * sigma**2) * l2_loss(mean_i, i0)
en_optimizer.zero_grad()
de_optimizer.zero_grad()
total_loss.backward()
en_optimizer.step()
de_optimizer.step()
with torch.no_grad():
i0 = ((1 / sigma**2) * mean_i + rho *
(i0_til_torch - Y) + alpha * noise_im_torch) / (
(1 / sigma**2) + rho + alpha)
with torch.no_grad():
############################### sub-problem 2 #################################
i0_np = i0.cpu().squeeze().detach().numpy()
Y_np = Y.cpu().squeeze().detach().numpy()
tmp = i0_np.transpose(1, 2, 0) + Y_np.transpose(1, 2, 0)
tmp = np.clip(tmp, 0, 1)
sig = noise_level
i0_til_np = bm3d.bm3d_rgb(tmp * 255, sig) / 255
i0_til_torch = torch.Tensor(
i0_til_np.transpose(2, 0, 1)[None, ...]).cuda()
############################### sub-problem 3 #################################
Y = Y + eta * (i0 - i0_til_torch)
###############################################################################
i0_til_np = i0_til_torch.cpu().squeeze().detach().numpy()
i0_til_np = np.clip(i0_til_np, 0, 1)
output = i0_til_np
psnr_gt = compare_psnr(clean_im.transpose(1, 2, 0),
255 * i0_til_np.transpose(1, 2, 0), 255)
ssim_gt = compare_ssim(255 * i0_til_np.transpose(1, 2, 0),
clean_im.transpose(1, 2, 0),
multichannel=True,
data_range=255)
if not i % 5:
denoise_obj_pil = Image.fromarray((tmp * 255).astype(np.uint8))
Y_np = Y.cpu().squeeze().detach().numpy()
i0_np = np.clip(i0_np, 0, 1)
i0_pil = Image.fromarray(
np.uint8(255 * i0_np.transpose(1, 2, 0)))
i0_til_np = i0_til_np.transpose(1, 2, 0)
i0_til_pil = Image.fromarray(
(255 * i0_til_np).astype(np.uint8))
mean_i_np = mean_i.cpu().squeeze().detach().numpy().clip(0, 1)
mean_i_pil = Image.fromarray(
(255 * mean_i_np.transpose(1, 2, 0)).astype(np.uint8))
out_np = out.cpu().squeeze().detach().numpy().clip(0, 1)
out_pil = Image.fromarray(
(255 * out_np.transpose(1, 2, 0)).astype(np.uint8))
denoise_obj_name = 'denoise_obj_{:04d}'.format(i) + '.png'
Y_name = 'Y_{:04d}'.format(i) + '.png'
i0_name = 'i0_num_epoch_{:04d}'.format(i) + '.png'
result_name = 'num_epoch_{:04d}'.format(i) + '.png'
mean_i_name = 'Latent_im_num_epoch_{:04d}'.format(i) + '.png'
out_name = 'res_of_dec_num_epoch_{:04d}'.format(i) + '.png'
denoise_obj_pil.save(result_root + denoise_obj_name)
Y_norm_np = np.sqrt((Y_np * Y_np).sum(0))
save_heatmap(Y_norm_np, result_root + Y_name)
i0_pil.save(result_root + i0_name)
i0_til_pil.save(result_root + result_name)
mean_i_pil.save(result_root + mean_i_name)
out_pil.save(result_root + out_name)
print('Iteration %02d Loss %f PSNR_gt: %f, SSIM_gt: %f' %
(i, total_loss.item(), psnr_gt, ssim_gt),
file=fo,
flush=True)
psnr = psnr_gt
ssim = ssim_gt
output_pil = Image.fromarray(
(255 * output.transpose(1, 2, 0)).astype(np.uint8))
output_pil.save(result_root + 'ours_result.png')
return psnr, ssim, psnr_bm3d, ssim_bm3d
